#!/usr/bin/env python3
"""驱动 v2: 每臂前 docker restart comfyui-secondary 净池(经验:净池1008f可跑,脏池362f OOM,/free无效)"""
import json, time, subprocess, urllib.request

BASE = 'http://127.0.0.1:8190'
CASES = ['case08', 'case07', 'case06']
JSONDIR = '/home/kai/shared/2026-10-17/h3lossless'

def ts(): return time.strftime('%H:%M:%S')

def restart_container():
    print(f'[{ts()}] docker restart comfyui-secondary ...', flush=True)
    subprocess.run(['docker', 'restart', 'comfyui-secondary'], check=True,
                   capture_output=True, timeout=120)
    for i in range(40):
        time.sleep(6)
        try:
            r = urllib.request.urlopen(BASE + '/system_stats', timeout=8)
            d = json.loads(r.read().decode())
            print(f'[{ts()}] healthy after {(i+1)*6}s, vram={d.get("devices",[{}])[0].get("vram_free")}', flush=True)
            return True
        except Exception:
            continue
    print(f'[{ts()}] container NOT healthy after 240s!', flush=True)
    return False

def run_case(case):
    api = json.load(open(f'{JSONDIR}/api_full_{case}_15s.json'))
    body = json.dumps({'prompt': api, 'client_id': 'h3lossless-v2'}).encode()
    req = urllib.request.Request(BASE + '/prompt', data=body,
                                 headers={'Content-Type': 'application/json'})
    resp = json.loads(urllib.request.urlopen(req, timeout=120).read().decode())
    pid = resp.get('prompt_id')
    print(f'[{ts()}] {case} submitted pid={pid}', flush=True)
    t0 = time.time()
    while time.time() - t0 < 90 * 60:
        time.sleep(30)
        try:
            h = json.loads(urllib.request.urlopen(f'{BASE}/history/{pid}', timeout=10).read().decode())
        except Exception:
            continue
        for k, v in h.items():
            st = v.get('status', {}).get('status_str', '')
            if st in ('success', 'error'):
                print(f'[{ts()}] {case} FINISHED status={st} ({(time.time()-t0)/60:.1f}min)', flush=True)
                if st == 'error':
                    for m in v.get('status', {}).get('messages', []):
                        if isinstance(m, list) and m[0] == 'execution_error':
                            print('  ERR:', m[1].get('exception_message', '')[:200], flush=True)
                return st
    print(f'[{ts()}] {case} TIMEOUT', flush=True)
    return 'timeout'

results = {}
for case in CASES:
    if not restart_container():
        results[case] = 'container-unhealthy'
        break
    results[case] = run_case(case)

print('=== DRIVER_V2 DONE ===', flush=True)
for c, r in results.items():
    print(c, '->', r, flush=True)