"""Run via ssh_task.py after uploading comfy_progress.py. No GPU workers restart."""
import datetime, pathlib, subprocess, json
payload = pathlib.Path('/home/ps/huobao-progress-update/comfy_progress.py').read_text()
marker = '# BEGIN HUOBAO READ-ONLY PROGRESS\n'
compile(payload, 'comfy_progress.py', 'exec')
names = ['canvas-h3-prompt-normalizer-healthy-20261026', 'huobao-vdn-normalizer']
backup = pathlib.Path('/home/ps/huobao-progress-update') / ('backup-' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
backup.mkdir(mode=0o700)
changed = []
for name in names:
    path = pathlib.Path('/home/ps/huobao-gpu-selection') / (name + '.candidate.py')
    original = path.read_text(); source = original.split(marker)[0]
    for old,new in [
      ('return {**result, "worker_state": "running"}', 'return {**result, "worker_state": "running", "progress": await huobao_progress(provider, prompt_id, queue, provider_headers(provider))}'),
      ('return {**result, "worker_state": "queued", "queue_position": pending.index(prompt_id) + 1}', 'return {**result, "worker_state": "queued", "queue_position": pending.index(prompt_id) + 1, "progress": await huobao_progress(provider, prompt_id, queue, provider_headers(provider))}')]:
        if old in source:
            assert source.count(old) == 1
            source = source.replace(old,new)
        else: assert new in source
    source = source.rstrip() + '\n\n' + marker + payload
    compile(source, str(path), 'exec')
    (backup / (name + '.py')).write_text(original)
    # Preserve the inode of the individually bind-mounted source file.
    path.write_text(source); changed.append(name)
for name,path in [('huobao-model-bridge','/home/ps/huobao-model-bridge/bridge.py'),('huobao-vdn-bridge','/home/ps/huobao-vdn/bridge/bridge.py')]:
    path = pathlib.Path(path); original = path.read_text(); source = original
    marker_line = '        task["progress"] = result.get("progress")'
    if marker_line not in source:
        lines = source.splitlines()
        indexes = [i for i,line in enumerate(lines) if line.startswith('        task = {"id": jid, "status": state, "duration": duration,')]
        assert len(indexes) == 1
        lines.insert(indexes[0]+1,marker_line); source = '\n'.join(lines)+'\n'
    compile(source,str(path),'exec'); (backup / (name+'.py')).write_text(original)
    path.write_text(source); changed.append(name)
for name in changed:
    subprocess.run(['docker','restart','--time','15',name],check=True,capture_output=True)
    d=json.loads(subprocess.check_output(['docker','inspect',name]))[0]
    assert d['State']['Running'], name
    print(name,'running')
print('backup',str(backup))