fix: 缓存改用文本CSV(逐行写)替代pickle，避免容器cgroup OOM静默杀进程
pickle.dump 150万记录的memo瞬间撑爆容器内存上限被杀；改为流式逐行写 保留的历史行到 cache_filtered_history.csv，读回用 load_sample_files。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
OwnerSunshine530
+64
-111
@@ -3,15 +3,19 @@
|
||||
不进提交包。**以子进程方式运行**（AI Studio 内核禁止 import torch）：
|
||||
|
||||
%cd /home/aistudio/code
|
||||
!python bench.py --diag # 诊断：序列长度分布 + sign-id 超界比例
|
||||
!python bench.py --smoke 50 # 冒烟：只跑前 50 batch
|
||||
!python bench.py # 默认基线
|
||||
!python bench.py --fp32 # FP32 天花板（Task 3）
|
||||
!python bench.py --fp32 # FP32 天花板
|
||||
!python bench.py --rebuild # 强制重建过滤缓存
|
||||
|
||||
关键设计——只保留“测试用户”的数据：
|
||||
不同用户被因果 mask 完全隔离，非测试用户的前向输出不参与打分；过滤掉它们
|
||||
对测试样本的 AUC/PCOC 没有任何影响，却能把数据量从 924 万条降到一小部分，
|
||||
避免 CTRTestSeqDataset 构造时 OOM。过滤后的数据缓存到磁盘，后续秒级复用。
|
||||
只保留“测试用户”的数据：不同用户被因果 mask 完全隔离，非测试用户的前向输出
|
||||
不参与打分；过滤掉它们对测试样本的 AUC/PCOC 没有任何影响，却能把数据量从
|
||||
924 万条降到一小部分。
|
||||
|
||||
缓存用**文本 CSV**而非 pickle：容器 cgroup 内存有限，pickle.dump 大对象的 memo
|
||||
会瞬间撑爆内存被静默 OOM-kill；逐行写 CSV 内存几乎不涨，再用 load_sample_files
|
||||
读回，稳。
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
@@ -46,22 +50,28 @@ def _test_user_ids(test_csv):
|
||||
return users
|
||||
|
||||
|
||||
def _load_filtered(history_dir, test_csv, test_users):
|
||||
"""流式读取所有文件，只保留 userid ∈ test_users 的记录（不持有完整字典，防 OOM）。
|
||||
|
||||
解析逻辑与 infer.load_sample_files 完全一致，只是多了一道用户过滤。
|
||||
def _stream_build(ref, cache_csv_path=None):
|
||||
"""流式过滤：构建 item_dict/user_seq；若给 cache_csv_path，同时把保留的历史行
|
||||
原样逐行写入（低内存文本缓存，test.csv 直接复用、不进缓存）。
|
||||
"""
|
||||
files = (sorted(history_dir.glob("*.csv")) if history_dir.exists() else []) + [test_csv]
|
||||
test_csv = ref / "test.csv"
|
||||
history = ref / "history"
|
||||
test_users = _test_user_ids(test_csv)
|
||||
files = (sorted(history.glob("*.csv")) if history.exists() else []) + [test_csv]
|
||||
print(f"[BENCH] 流式过滤加载 {len(files)} 个文件（仅保留 {len(test_users)} 个测试用户）...")
|
||||
|
||||
item_dict = {}
|
||||
user_logs = defaultdict(list)
|
||||
cf = open(cache_csv_path, "w") if cache_csv_path else None
|
||||
try:
|
||||
for fp in files:
|
||||
has_clk = infer._detect_has_clk(fp)
|
||||
min_parts = 5 if has_clk else 4
|
||||
is_test = (Path(fp).name == test_csv.name)
|
||||
kept = 0
|
||||
with open(fp) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
for raw in f:
|
||||
line = raw.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split(",")
|
||||
@@ -70,6 +80,8 @@ def _load_filtered(history_dir, test_csv, test_users):
|
||||
userid = int(parts[1])
|
||||
if userid not in test_users:
|
||||
continue
|
||||
if cf is not None and not is_test: # 只缓存历史行
|
||||
cf.write(raw if raw.endswith("\n") else raw + "\n")
|
||||
logid = int(parts[0])
|
||||
adid = int(parts[2])
|
||||
if has_clk:
|
||||
@@ -94,68 +106,58 @@ def _load_filtered(history_dir, test_csv, test_users):
|
||||
}
|
||||
user_logs[userid].append((timestamp, logid))
|
||||
kept += 1
|
||||
print(f" {fp.name}: has_clk={has_clk}, kept={kept}")
|
||||
print(f" {Path(fp).name}: has_clk={has_clk}, kept={kept}")
|
||||
finally:
|
||||
if cf is not None:
|
||||
cf.flush()
|
||||
os.fsync(cf.fileno())
|
||||
cf.close()
|
||||
|
||||
user_seq = {}
|
||||
for u, logs in user_logs.items():
|
||||
logs.sort(key=lambda x: x[0])
|
||||
user_seq[u] = [lid for _, lid in logs]
|
||||
print(f"[BENCH] 过滤后：{len(item_dict)} 条记录，{len(user_seq)} 个用户")
|
||||
if cache_csv_path:
|
||||
print(f"[BENCH] 已缓存历史行 -> {cache_csv_path}（下次快速读取）")
|
||||
return item_dict, user_seq
|
||||
|
||||
|
||||
def _cache_path(cur):
|
||||
return cur / "bench_filtered_cache.pkl"
|
||||
|
||||
|
||||
def _build_filtered(ref):
|
||||
test_csv = ref / "test.csv"
|
||||
history = ref / "history"
|
||||
test_users = _test_user_ids(test_csv)
|
||||
return _load_filtered(history, test_csv, test_users)
|
||||
|
||||
|
||||
def _load_cache(cache):
|
||||
import pickle
|
||||
with open(cache, "rb") as f:
|
||||
d = pickle.load(f)
|
||||
return d["item_dict"], d["user_seq"]
|
||||
|
||||
|
||||
def _save_cache(cache, item_dict, user_seq):
|
||||
"""原子写 + fsync + 写后校验；任何异常都不留毒文件。
|
||||
|
||||
用 pickle 而非 torch.save：AI Studio overlay 文件系统对 torch 的 zip/mmap
|
||||
读取会间歇性报 [Errno 38]。pickle.dump 大对象较慢但顺序写更稳。
|
||||
"""
|
||||
import pickle
|
||||
try:
|
||||
with open(cache, "wb") as f:
|
||||
pickle.dump({"item_dict": item_dict, "user_seq": user_seq}, f,
|
||||
protocol=pickle.HIGHEST_PROTOCOL)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
print(f"[BENCH] 已缓存 -> {cache}（下次秒级读取；读不出会自动重建）")
|
||||
except Exception as e:
|
||||
print(f"[BENCH][WARN] 缓存写入失败（{e}），本次不缓存（不影响结果）")
|
||||
try:
|
||||
os.remove(cache)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _get_data(cur, ref, rebuild=False):
|
||||
"""取过滤后的 (item_dict, user_seq)，优先读磁盘缓存。"""
|
||||
cache = _cache_path(cur)
|
||||
if cache.exists() and not rebuild:
|
||||
print(f"[BENCH] 读取过滤缓存：{cache}")
|
||||
"""取过滤后的 (item_dict, user_seq)，优先读 CSV 缓存。"""
|
||||
cache_csv = cur / "cache_filtered_history.csv"
|
||||
test_csv = ref / "test.csv"
|
||||
if cache_csv.exists() and not rebuild:
|
||||
print(f"[BENCH] 读取过滤缓存(CSV)：{cache_csv}")
|
||||
try:
|
||||
return _load_cache(cache)
|
||||
return infer.load_sample_files([str(cache_csv), str(test_csv)])
|
||||
except Exception as e:
|
||||
print(f"[BENCH][WARN] 缓存读取失败（{e}），重新构建")
|
||||
item_dict, user_seq = _build_filtered(ref)
|
||||
_save_cache(cache, item_dict, user_seq)
|
||||
return item_dict, user_seq
|
||||
return _stream_build(ref, cache_csv_path=str(cache_csv))
|
||||
|
||||
|
||||
def run_diag(rebuild=False):
|
||||
"""诊断：测试用户序列长度分布 + sign-id 是否超界（判断上下文与 modulo 的价值）。"""
|
||||
cur = Path(__file__).parent
|
||||
ref = cur / "dataset"
|
||||
item_dict, user_seq = _get_data(cur, ref, rebuild=rebuild)
|
||||
lens = np.array([len(v) for v in user_seq.values()]) if user_seq else np.array([0])
|
||||
print(f"[DIAG] 测试用户数={len(user_seq)} 总记录数={len(item_dict)}")
|
||||
print(f"[DIAG] 每用户序列长度 min/median/mean/max = "
|
||||
f"{int(lens.min())}/{int(np.median(lens))}/{lens.mean():.1f}/{int(lens.max())}")
|
||||
print(f"[DIAG] 序列长度>1 的用户占比 = {(lens > 1).mean():.1%}")
|
||||
VOCAB = 5_000_000
|
||||
mx, over, tot = 0, 0, 0
|
||||
for rec in item_dict.values():
|
||||
s = rec["signs"]
|
||||
if s.size:
|
||||
m = int(s.max())
|
||||
if m > mx:
|
||||
mx = m
|
||||
over += int((s >= VOCAB).sum())
|
||||
tot += int(s.size)
|
||||
print(f"[DIAG] max_sign_id={mx} vocab={VOCAB} "
|
||||
f"超界sign占比={over}/{tot}={(over / max(tot, 1)):.2%}")
|
||||
|
||||
|
||||
def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
@@ -174,7 +176,6 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
test_csv = ref / "test.csv"
|
||||
label_file = ref / "label_data.txt"
|
||||
|
||||
# ----- 取数据（过滤+缓存）-----
|
||||
item_dict, user_seq = _get_data(cur, ref, rebuild=rebuild)
|
||||
test_logids = infer.load_logids_from_file(test_csv)
|
||||
ds = infer.CTRTestSeqDataset(
|
||||
@@ -191,15 +192,12 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
if max_batches is not None and len(batches) >= max_batches:
|
||||
break
|
||||
|
||||
# 释放构造期内存，降低推理峰值
|
||||
del item_dict, user_seq, ds, loader
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
# ----- 加载模型 -----
|
||||
model, dev = infer.load_model(ckpt_path=None)
|
||||
|
||||
# ----- 推理 + 同步计时 -----
|
||||
logid2p = {}
|
||||
t_sum = 0.0
|
||||
cuda = (dev.type == "cuda")
|
||||
@@ -218,7 +216,6 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
for lid, p in zip(b["logid"][pm].cpu().tolist(), probs[pm].cpu().tolist()):
|
||||
logid2p[lid] = p
|
||||
|
||||
# ----- 按 test.csv 顺序写 predict.txt 并打分 -----
|
||||
order = [int(l.split(",")[0]) for l in open(test_csv) if l.strip()]
|
||||
missing = [lid for lid in order if lid not in logid2p]
|
||||
if missing:
|
||||
@@ -238,54 +235,10 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
return res
|
||||
|
||||
|
||||
def run_diag(rebuild=False):
|
||||
"""诊断：测试用户序列长度分布 + sign-id 是否超界（判断上下文与 modulo 的价值）。
|
||||
|
||||
先打印诊断，再写缓存——避免缓存写入卡住时看不到诊断结果。
|
||||
"""
|
||||
cur = Path(__file__).parent
|
||||
ref = cur / "dataset"
|
||||
cache = _cache_path(cur)
|
||||
loaded = False
|
||||
item_dict = user_seq = None
|
||||
if cache.exists() and not rebuild:
|
||||
print(f"[BENCH] 读取过滤缓存：{cache}")
|
||||
try:
|
||||
item_dict, user_seq = _load_cache(cache)
|
||||
loaded = True
|
||||
except Exception as e:
|
||||
print(f"[BENCH][WARN] 缓存读取失败（{e}），重新构建")
|
||||
if not loaded:
|
||||
item_dict, user_seq = _build_filtered(ref)
|
||||
|
||||
lens = np.array([len(v) for v in user_seq.values()]) if user_seq else np.array([0])
|
||||
print(f"[DIAG] 测试用户数={len(user_seq)} 总记录数={len(item_dict)}")
|
||||
print(f"[DIAG] 每用户序列长度 min/median/mean/max = "
|
||||
f"{int(lens.min())}/{int(np.median(lens))}/{lens.mean():.1f}/{int(lens.max())}")
|
||||
print(f"[DIAG] 序列长度>1 的用户占比 = {(lens > 1).mean():.1%} "
|
||||
f"（占比低=大量测试样本没有历史上下文 → 生成式模型发挥不出来）")
|
||||
VOCAB = 5_000_000
|
||||
mx, over, tot = 0, 0, 0
|
||||
for rec in item_dict.values():
|
||||
s = rec["signs"]
|
||||
if s.size:
|
||||
m = int(s.max())
|
||||
if m > mx:
|
||||
mx = m
|
||||
over += int((s >= VOCAB).sum())
|
||||
tot += int(s.size)
|
||||
print(f"[DIAG] max_sign_id={mx} vocab={VOCAB} "
|
||||
f"超界sign占比={over}/{tot}={(over / max(tot, 1)):.2%} "
|
||||
f"（占比高=clamp 在污染 embedding → modulo 可能找回 AUC）")
|
||||
|
||||
if not loaded:
|
||||
_save_cache(_cache_path(cur), item_dict, user_seq)
|
||||
|
||||
|
||||
def _parse_args():
|
||||
import argparse
|
||||
ap = argparse.ArgumentParser(description="CTI 推理测量闭环（子进程跑：!python bench.py ...）")
|
||||
ap.add_argument("--diag", action="store_true", help="只跑诊断（序列长度分布 + sign-id 超界比例），不推理")
|
||||
ap.add_argument("--diag", action="store_true", help="只跑诊断，不推理")
|
||||
ap.add_argument("--smoke", type=int, default=None, help="只跑前 N 个 batch（冒烟）")
|
||||
ap.add_argument("--bs", type=int, default=50, help="batch_size（本地参考）")
|
||||
ap.add_argument("--fp32", action="store_true", help="FP32 天花板 = 关 fp16 + 关 expert 合并")
|
||||
|
||||
Reference in New Issue
Block a user