Robot_Arm/ml/pstate.py

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