"""流水线状态与互斥锁（单一写者：所有对 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()