#!/usr/bin/env python3
"""Stop only existing router/controller/model; always restore the same containers."""
import json, pathlib, subprocess, sys, time, signal
ROOT=pathlib.Path('/data/benchmarks/dflash2-medium-20260920')
NAMES=['mike-ai-router','mike-ai-profile-controller','mike-ai-llama-medium']
def run(*args,check=True,timeout=90):
    return subprocess.run(args,capture_output=True,text=True,check=check,timeout=timeout)
def stop_signal(*_): raise RuntimeError('Supervisor interrupted')
signal.signal(signal.SIGTERM,stop_signal); signal.signal(signal.SIGINT,stop_signal)
# Refuse if the known production state has changed, or if requests are active.
for name in NAMES:
    assert run('docker','inspect',name,'--format','{{.State.Running}}').stdout.strip()=='true',name
slots=json.loads(run('docker','exec',NAMES[-1],'curl','-fsS','http://127.0.0.1:8080/slots').stdout)
assert not any(s['is_processing'] for s in slots),'Production request active'
assert not run('docker','ps','-q','--filter','name=^mike-ai-dflash2-test$').stdout.strip(),'Existing experiment'
child=None
try:
    run('docker','stop','-t','30',*NAMES[:2])
    # Drain requests already handed to the model, before unloading it.
    for _ in range(120):
        slots=json.loads(run('docker','exec',NAMES[-1],'curl','-fsS','http://127.0.0.1:8080/slots').stdout)
        if not any(s['is_processing'] for s in slots): break
        time.sleep(2)
    else: raise RuntimeError('Model did not drain')
    run('docker','stop','-t','30',NAMES[-1])
    child=subprocess.Popen(['python3',str(ROOT/'run.py'),sys.argv[1]])
    code=child.wait(timeout=600)
    if code: raise RuntimeError('Benchmark failed: '+str(code))
finally:
    if child is not None and child.poll() is None:
        child.terminate()
        try: child.wait(timeout=20)
        except subprocess.TimeoutExpired: child.kill(); child.wait(timeout=10)
    run('docker','rm','-f','mike-ai-dflash2-test',check=False)
    run('docker','start',NAMES[-1])
    for _ in range(150):
        if run('docker','inspect',NAMES[-1],'--format','{{.State.Health.Status}}').stdout.strip()=='healthy': break
        time.sleep(2)
    run('docker','start',NAMES[1],NAMES[0])
    print('RESTORED existing medium/controller/router',flush=True)