Files

53 lines
3.0 KiB
Python

#!/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/medium-split-mtp-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
for attempt in range(60):
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
if attempt==0: print('WAIT production request active; no interruption',flush=True)
time.sleep(3)
else: raise RuntimeError('Production remained busy for 180s; no services stopped')
assert not run('docker','ps','-q','--filter','name=^mike-ai-split-mtp-test$').stdout.strip(),'Existing experiment'
production=json.loads(run('docker','inspect',NAMES[-1]).stdout)[0]
(ROOT/'production.json').write_text(json.dumps({'image':production['Image'],'args':production['Args']},indent=2)+'\n')
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-split-mtp-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)
else: raise RuntimeError('Restored Medium did not become healthy')
run('docker','start',NAMES[1],NAMES[0])
for _ in range(60):
statuses=[run('docker','inspect',name,'--format','{{.State.Health.Status}}').stdout.strip() for name in NAMES[:2]]
if all(status=='healthy' for status in statuses): break
time.sleep(2)
else: raise RuntimeError('Restored router/controller did not become healthy')
print('RESTORED existing medium/controller/router; all healthy',flush=True)