cbb
c2732fb3cd
- train.py：evaluate 改 bs=64，结束后 empty_cache()。WDDM 下 15.9 GiB 的评估峰值用完不退， 整卡顶到 96% 后驱动把训练张量换出（实测 Local 18.8 GB / Shared 1.9 GB），吞吐从 5.5k 掉到 444 samples/s、8000 步 eval 卡死 68 分钟；改后峰值 5.32 GiB、eval 106 s、吞吐 6.4–7.0k。 - pstate/pipeline：新增 --fresh-from=<阶段>：只把该阶段及之后的产物移进 ml/archive/（不删）， 与参数失效取并集；配合 FROM 实现"只重训某一阶段，之后全部重做"（沙盒实测选 train 只归档 ml/runs/base）。 - monitor：运行方式加「从指定阶段重训」+ 起始阶段下拉；TARGET_SAMPLES 默认改 200 万； 「最好 X%」改标最好那次 eval 的步数（原来错标当前步，18,000 步显示成 0.2% 的出处）。 连同工作区已有的停止补扫 sweep_nt（按命令行补杀 git-bash 下逃逸的孙进程）一并提交。 - eval/diag_tail.py：闭环末位误差 / 是否进过 5 mm 窗口的逐题诊断（手动运行，不进流水线）。
243 lines
13 KiB
Python
243 lines
13 KiB
Python
"""流水线状态与互斥锁（单一写者：所有对 ml/runs/pipeline.json / pipeline.lock 的写都经过本模块，原子替换）。
|
||
pipeline.sh 自动调用，一般不需要手动运行：
|
||
python -m ml.pstate reset <阶段1,...> [--fresh] # `all` 开始：保留 rc=0 的已完成阶段（续跑）；参数变了只作废受影响的阶段；--fresh 全部作废；--fresh-from=<阶段> 只作废该阶段及之后
|
||
python -m ml.pstate todo # 打印计划中尚未完成的阶段（空格分隔）
|
||
python -m ml.pstate begin <阶段> [键=值 ...] # 阶段开始
|
||
python -m ml.pstate set <阶段> 键=值 ... # 阶段中途更新
|
||
python -m ml.pstate get <阶段> <键> # 读值；不存在则非零退出（不静默回空）
|
||
python -m ml.pstate end <阶段> <返回码>
|
||
python -m ml.pstate timeout <阶段> # 该阶段超时秒数（估时 ×3，TIMEOUT_SCALE 可再乘）
|
||
python -m ml.pstate lock acquire <所有者> [pid]|release|beat|show
|
||
2026-10-18 第五轮：
|
||
- 加锁（审查 N-7）：网页 / 终端 / `all` 都走同一把锁 ml/runs/pipeline.lock（pid + 心跳），已有一条在跑时立即拒绝；
|
||
心跳超过 LOCK_STALE_S 未刷新视为陈旧（进程被强杀），可接管。
|
||
- 续跑（缺口 1）：reset 默认保留已完成阶段；参数变了只作废受它影响的阶段（AFFECTS），其产物移到 ml/archive/，避免新旧参数混用。
|
||
- 阶段表 STAGES 是估时 / 超时 / 网页进度权重的唯一来源（审查 N-14：网页权重与耗时表同源）。
|
||
- 停止也走本模块（审查 N-12：此前 monitor 直写 pipeline.json，与子进程竞态）。
|
||
- 累计用时（2026-10-18）：每个阶段跑过的时间并进顶部 time{阶段: 秒}，跨「停止 → 继续训练」累加，
|
||
网页的「总用时 / 本步用时」据此显示；产物被归档（要重做）的阶段清零，其余保留。"""
|
||
import glob, json, os, shutil, socket, sys, time
|
||
|
||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
P = os.path.join(ROOT, 'ml', 'runs', 'pipeline.json')
|
||
LOCK = os.path.join(ROOT, 'ml', 'runs', 'pipeline.lock')
|
||
LOCK_STALE_S = 300 # 心跳每 60 s 一次；5 分钟没刷新 = 持锁进程已死
|
||
|
||
# 阶段顺序、名称、估时（小时，按 README §耗时表的共识分支上界；网页进度权重即此值）
|
||
STAGES = [
|
||
('setup', '安装依赖', 0.2), ('kin', '对拍', 0.2), ('r4', '推理延迟', 0.2), ('gate1', '一致性门禁', 0.5),
|
||
('gen', '生成老师轨迹', 5.0), ('prep', '打包样本', 0.5), ('train', '预训练', 3.0), ('m5', '头门禁', 0.5),
|
||
('apg', 'APG 微调', 4.0), ('mine', '挖错循环', 9.0), ('export', '导出', 0.1), ('bench', '部署评测', 1.0),
|
||
('accept', '验收汇总', 0.5)]
|
||
ORDER = [s[0] for s in STAGES]
|
||
EST_H = {s[0]: s[2] for s in STAGES}
|
||
# m5 不过时流程内自动训 WTA 头（训练 3 h + 门禁），超时按分支后的总量给
|
||
TIMEOUT_H = {k: max(1.0, 3 * v) for k, v in EST_H.items()}
|
||
TIMEOUT_H['m5'] = 3 * (EST_H['train'] + 2 * EST_H['m5'])
|
||
# 参数 → 受它影响的阶段（这些阶段即使上次完成也要重跑，其产物先挪进 ml/archive/）。只改挖错轮数不会让 5 小时的生成作废。
|
||
_from = lambda st: ORDER[ORDER.index(st):]
|
||
AFFECTS = {**{k: _from('gen') for k in ('N_TASKS', 'NPROC', 'CONSENSUS', 'MIN_GAP_MM')},
|
||
# TARGET_SAMPLES 只决定"补量到多少"：不作废 gen 的已有分片（重跑 gen 会跳过写完的 shard、按差额补量），
|
||
# 但样本集会变，所以 prep 及之后作废。2026-10-18 实测：把它当"从 gen 作废"会让 5 小时的 5 万条分片白跑。
|
||
'TARGET_SAMPLES': ['gen'] + _from('prep'),
|
||
'CFG': ['r4'] + _from('train'), **{k: _from('mine') for k in ('ROUNDS', 'APG_STEPS', 'MINE_N')}}
|
||
PARAM_KEYS = list(AFFECTS)
|
||
# 各阶段产物（相对仓库根，可含通配）；作废 / 从头重跑时整体移到 ml/archive/<时间>/，不删除
|
||
PRODUCTS = {
|
||
'r4': ['ml/runs/r4', 'ml/runs/r4.json'], 'gate1': ['ml/data/consistency.json'],
|
||
'gen': ['ml/data/shard_*', 'ml/data/log_*', 'ml/data/stats.json'], 'prep': ['ml/data/samples', 'ml/data/r2.json'],
|
||
'train': ['ml/runs/base', 'ml/runs/wta', 'ml/runs/small', 'ml/runs/bench_*'],
|
||
'm5': ['ml/runs/m5_*', 'ml/runs/selected_ckpt'], 'apg': ['ml/runs/apg', 'ml/runs/current.pt'],
|
||
'mine': ['ml/runs/mine', 'ml/runs/apg_r*', 'ml/data/e2', 'ml/data/samples_e2'],
|
||
'export': ['webapp/models/traj_policy.onnx', 'webapp/models/traj_policy.meta.json'],
|
||
'bench': ['ml/runs/policy_bench.json'], 'accept': ['ml/runs/acceptance.json']}
|
||
|
||
|
||
def archive(stages):
|
||
"""把这些阶段的产物移到 ml/archive/<时间>/（保持相对路径），返回目标目录或 None。"""
|
||
# 2026-10-18：同一秒内多次归档会撞目录、覆盖旧文件，加序号后缀保证唯一
|
||
base = os.path.join(ROOT, 'ml', 'archive', time.strftime('%Y%m%d_%H%M%S')); dst, n = base, 1
|
||
while os.path.exists(dst): dst = '%s_%d' % (base, n); n += 1
|
||
moved = 0
|
||
for st in stages:
|
||
for pat in PRODUCTS.get(st, []):
|
||
for f in glob.glob(os.path.join(ROOT, pat)):
|
||
rel = os.path.relpath(f, ROOT); os.makedirs(os.path.dirname(os.path.join(dst, rel)), exist_ok=True)
|
||
shutil.move(f, os.path.join(dst, rel)); moved += 1
|
||
return os.path.relpath(dst, ROOT) if moved else None
|
||
|
||
|
||
def load():
|
||
try:
|
||
with open(P, encoding='utf-8') as f: return json.load(f)
|
||
except Exception:
|
||
return {'started': time.time(), 'plan': None, 'stages': {}, 'current': None, 'params': {}, 'time': {}}
|
||
|
||
|
||
def save(d):
|
||
os.makedirs(os.path.dirname(P), exist_ok=True)
|
||
tmp = P + '.%d.tmp' % os.getpid()
|
||
with open(tmp, 'w', encoding='utf-8') as f: json.dump(d, f, ensure_ascii=False, indent=1)
|
||
os.replace(tmp, P)
|
||
|
||
|
||
def kv(args):
|
||
out = {}
|
||
for a in args:
|
||
k, _, v = a.partition('=')
|
||
try: v = float(v) if '.' in v else int(v)
|
||
except ValueError: pass
|
||
out[k] = v
|
||
return out
|
||
|
||
|
||
def params_now():
|
||
return {k: os.environ.get(k, '') for k in PARAM_KEYS}
|
||
|
||
|
||
# ---------------------------------------------------------------- 锁
|
||
def lock_read():
|
||
try:
|
||
with open(LOCK, encoding='utf-8') as f: return json.load(f)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def lock_alive(lk=None):
|
||
lk = lk if lk is not None else lock_read()
|
||
return bool(lk) and time.time() - lk.get('beat', 0) < LOCK_STALE_S
|
||
|
||
|
||
def lock_acquire(owner, pid=None):
|
||
"""原子创建锁文件；已有活锁则返回 (False, 锁内容)。陈旧锁直接接管。"""
|
||
os.makedirs(os.path.dirname(LOCK), exist_ok=True)
|
||
now = time.time()
|
||
rec = {'pid': pid or os.getppid(), 'owner': owner, 'start': now, 'beat': now, 'host': socket.gethostname()}
|
||
for _ in range(2):
|
||
try:
|
||
fd = os.open(LOCK, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
||
with os.fdopen(fd, 'w', encoding='utf-8') as f: json.dump(rec, f)
|
||
try: os.remove(os.path.join(ROOT, 'ml', 'runs', 'STOP')) # 上次的停止请求不能让这次一开训就退出
|
||
except FileNotFoundError: pass
|
||
return True, rec
|
||
except FileExistsError:
|
||
lk = lock_read()
|
||
if lk is not None and lock_alive(lk): return False, lk
|
||
try: os.remove(LOCK) # 陈旧锁（或刚好被写了一半）：接管
|
||
except FileNotFoundError: pass
|
||
return False, lock_read()
|
||
|
||
|
||
def lock_beat():
|
||
lk = lock_read()
|
||
if not lk: return
|
||
lk['beat'] = time.time()
|
||
tmp = LOCK + '.%d.tmp' % os.getpid()
|
||
with open(tmp, 'w', encoding='utf-8') as f: json.dump(lk, f)
|
||
os.replace(tmp, LOCK)
|
||
|
||
|
||
def lock_release():
|
||
try: os.remove(LOCK)
|
||
except FileNotFoundError: pass
|
||
|
||
|
||
# ---------------------------------------------------------------- 状态
|
||
def todo(d=None):
|
||
d = d or load()
|
||
return [s for s in (d.get('plan') or []) if (d['stages'].get(s) or {}).get('rc') != 0]
|
||
|
||
|
||
def _acc_time(d, st, now):
|
||
"""把"这一段"的用时并进累计 time{阶段: 秒}。幂等：stage.end 已写就不再加，停止（mark_stopped）
|
||
与流水线的 end 都会调，重复调用不会双记。"""
|
||
s = d['stages'].get(st)
|
||
if s and s.get('start') and s.get('end') is None:
|
||
t = d.setdefault('time', {})
|
||
t[st] = round(t.get(st, 0) + max(0.0, now - s['start']), 1)
|
||
|
||
|
||
def mark_stopped(rc=-15):
|
||
"""网页「停止」后调用：当前阶段记为被停止、这段用时并进累计，释放锁。"""
|
||
d = load(); st = d.get('current'); now = time.time()
|
||
if st and st in d['stages'] and d['stages'][st].get('end') is None:
|
||
_acc_time(d, st, now)
|
||
d['stages'][st].update(end=now, rc=rc)
|
||
d['current'] = None
|
||
save(d); lock_release()
|
||
|
||
|
||
def main():
|
||
cmd, *rest = sys.argv[1:] or ['?']
|
||
if cmd == 'lock':
|
||
sub = rest[0] if rest else 'show'
|
||
if sub == 'acquire':
|
||
ok, lk = lock_acquire(rest[1] if len(rest) > 1 else 'terminal', int(rest[2]) if len(rest) > 2 else None)
|
||
if not ok:
|
||
age = time.time() - lk.get('beat', 0)
|
||
print(f"已有一条流水线在运行（{lk.get('owner')}，pid {lk.get('pid')}，{age:.0f} s 前心跳）。"
|
||
f"锁文件 {os.path.relpath(LOCK, ROOT)}；确认它已经不在了，等 {LOCK_STALE_S // 60} 分钟自动失效或手动删除该文件。", file=sys.stderr)
|
||
sys.exit(9)
|
||
elif sub == 'beat': lock_beat()
|
||
elif sub == 'release': lock_release()
|
||
else: print(json.dumps(lock_read()))
|
||
return
|
||
if cmd == 'timeout':
|
||
print(int(TIMEOUT_H.get(rest[0], 3.0) * 3600 * float(os.environ.get('TIMEOUT_SCALE') or 1))); return
|
||
if cmd == 'get':
|
||
v = (load()['stages'].get(rest[0]) or {}).get('info', {}).get(rest[1])
|
||
if v is None: print(f'pipeline.json 里没有 {rest[0]}.{rest[1]}', file=sys.stderr); sys.exit(1)
|
||
print(v); return
|
||
if cmd == 'todo':
|
||
print(' '.join(todo())); return
|
||
d = load(); now = time.time()
|
||
if cmd == 'reset':
|
||
plan = rest[0].split(',') if rest and rest[0] else []
|
||
pn = params_now(); prev = d.get('params') or {}
|
||
changed = [k for k in PARAM_KEYS if prev and prev.get(k) != pn[k]]
|
||
ff = next((r.split('=', 1)[1] for r in rest if r.startswith('--fresh-from=')), '')
|
||
if '--fresh' in rest: bad = set(ORDER)
|
||
else:
|
||
bad = {st for k in changed for st in AFFECTS[k]}
|
||
if ff in ORDER: bad |= set(ORDER[ORDER.index(ff):])
|
||
if changed: print('参数与上次不同：' + json.dumps({k: (prev.get(k), pn[k]) for k in changed}, ensure_ascii=False)
|
||
+ '；受影响的阶段重跑：' + ' '.join(st for st in ORDER if st in bad), file=sys.stderr)
|
||
keep = {k: v for k, v in d.get('stages', {}).items() if v.get('rc') == 0 and k not in bad}
|
||
# 受影响 / 未完成的阶段：旧产物挪走，免得续训、跳过 shard 时把旧参数的东西接上。
|
||
# 例外：只改了目标样本量时，gen 要重跑补量，但已有分片必须留着（run_parallel 跳过写完的 shard、
|
||
# 按差额补量），否则 5 小时的生成白跑。
|
||
no_archive = {'gen'} if 'TARGET_SAMPLES' in changed else set()
|
||
arch = [st for st in ORDER if st in bad and st not in no_archive]
|
||
arc = archive(arch) if arch else None
|
||
if arc: print(f'旧产物已移到 {arc}/（未删除）', file=sys.stderr)
|
||
# 累计用时跨续跑保留；只有产物被挪走的阶段（要从头重做）才清零
|
||
tm = {k: v for k, v in (d.get('time') or {}).items() if k not in arch}
|
||
d = {'started': now, 'plan': plan, 'stages': keep, 'current': None, 'params': pn, 'time': tm}
|
||
done = [s for s in plan if s in keep]
|
||
if done: print('跳过已完成的 %d 个阶段：%s' % (len(done), ' '.join(done)), file=sys.stderr)
|
||
save(d); return
|
||
elif cmd == 'begin':
|
||
st = rest[0]
|
||
prev = d['stages'].get(st) or {}
|
||
# 老记录（本次改动之前完成的）没有累计字段：用上一次的整段补上，免得续跑时用时从 0 起
|
||
if prev.get('start') and prev.get('end') and st not in (d.get('time') or {}):
|
||
d.setdefault('time', {})[st] = round(prev['end'] - prev['start'], 1)
|
||
d['stages'][st] = {'start': now, 'end': None, 'rc': None, 'info': kv(rest[1:]), 'timeout_h': TIMEOUT_H.get(st)}
|
||
d['current'] = st
|
||
elif cmd == 'set':
|
||
s = d['stages'].setdefault(rest[0], {'start': now, 'end': None, 'rc': None, 'info': {}})
|
||
s['info'].update(kv(rest[1:]))
|
||
elif cmd == 'end':
|
||
st, rc = rest[0], int(rest[1]) if len(rest) > 1 else 0
|
||
s = d['stages'].setdefault(st, {'start': now, 'info': {}})
|
||
_acc_time(d, st, now)
|
||
s['end'], s['rc'] = now, rc
|
||
if d.get('current') == st: d['current'] = None
|
||
else:
|
||
print(__doc__); return
|
||
save(d)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|