215 lines
11 KiB
Python
215 lines
11 KiB
Python
"""流水线状态与互斥锁（单一写者：所有对 ml/runs/pipeline.json / pipeline.lock 的写都经过本模块，原子替换）。
|
||
pipeline.sh 自动调用，一般不需要手动运行：
|
||
python -m ml.pstate reset <阶段1,...> [--fresh] # `all` 开始：保留 rc=0 的已完成阶段（续跑）；参数变了只作废受影响的阶段；--fresh 全部作废
|
||
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，与子进程竞态）。"""
|
||
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', 'TARGET_SAMPLES', 'CONSENSUS', 'MIN_GAP_MM')},
|
||
'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': {}}
|
||
|
||
|
||
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 mark_stopped(rc=-15):
|
||
"""网页「停止」后调用：当前阶段记为被停止，释放锁。"""
|
||
d = load(); st = d.get('current')
|
||
if st and st in d['stages'] and d['stages'][st].get('end') is None:
|
||
d['stages'][st].update(end=time.time(), 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 {}
|
||
if '--fresh' in rest: bad = set(ORDER)
|
||
else:
|
||
changed = [k for k in PARAM_KEYS if prev and prev.get(k) != pn[k]]
|
||
bad = {st for k in changed for st in AFFECTS[k]}
|
||
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 时把旧参数的东西接上
|
||
arc = archive([st for st in ORDER if st in bad]) if bad else None
|
||
if arc: print(f'旧产物已移到 {arc}/（未删除）', file=sys.stderr)
|
||
d = {'started': now, 'plan': plan, 'stages': keep, 'current': None, 'params': pn}
|
||
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]
|
||
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': {}})
|
||
s['end'], s['rc'] = now, rc
|
||
if d.get('current') == st: d['current'] = None
|
||
else:
|
||
print(__doc__); return
|
||
save(d)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|