63 lines
2.4 KiB
Python
63 lines
2.4 KiB
Python
|
"""流水线状态文件 ml/runs/pipeline.json（供训练监控网页显示"现在在哪一步、每步用了多久"）。
|
||
|
|
pipeline.sh 每个阶段开始/结束时自动调用，不需要手动运行：
|
||
|
|
python -m ml.pstate reset <阶段1,阶段2,...> # `all` 开始时：清空并写入计划
|
||
|
|
python -m ml.pstate begin <阶段> [键=值 ...] # 阶段开始（值可为数字），如 target_tasks=50000
|
||
|
|
python -m ml.pstate set <阶段> 键=值 ... # 阶段中途更新（如补量后的目标条数）
|
||
|
|
python -m ml.pstate end <阶段> <返回码>
|
||
|
|
2026-10-18 新增（用户要求：从训练开始到结束都能在监控面板看到进度）。"""
|
||
|
|
import json, os, sys, time
|
||
|
|
|
||
|
|
P = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'ml', 'runs', 'pipeline.json')
|
||
|
|
|
||
|
|
|
||
|
|
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}
|
||
|
|
|
||
|
|
|
||
|
|
def save(d):
|
||
|
|
os.makedirs(os.path.dirname(P), exist_ok=True)
|
||
|
|
tmp = P + '.tmp'
|
||
|
|
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 main():
|
||
|
|
cmd, *rest = sys.argv[1:] or ['?']
|
||
|
|
d = load(); now = time.time()
|
||
|
|
if cmd == 'reset':
|
||
|
|
plan = rest[0].split(',') if rest else []
|
||
|
|
# 从中间某步续跑时保留计划之外（之前已跑完）各步的记录，计划内的清空重来
|
||
|
|
d = {'started': now, 'plan': plan, 'stages': {k: v for k, v in d.get('stages', {}).items() if k not in plan}, 'current': None}
|
||
|
|
elif cmd == 'begin':
|
||
|
|
st = rest[0]
|
||
|
|
d['stages'][st] = {'start': now, 'end': None, 'rc': None, 'info': kv(rest[1:])}
|
||
|
|
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()
|