100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Keep a checkpoint-aware training-loop parent alive across container recreation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
TERMINAL_STATUSES = {"training_complete"}
|
|
|
|
|
|
def read_loop_status(path: Path) -> str | None:
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return "invalid"
|
|
return value.get("status") if isinstance(value, dict) else "invalid"
|
|
|
|
|
|
def process_active(container: str, marker: str) -> bool:
|
|
result = subprocess.run(
|
|
["docker", "top", container, "-eo", "pid,args"],
|
|
capture_output=True, text=True, check=False,
|
|
)
|
|
return result.returncode == 0 and any(marker in line for line in result.stdout.splitlines())
|
|
|
|
|
|
def write_state(path: Path, state: dict) -> None:
|
|
state["updated_at"] = datetime.now(UTC).isoformat()
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(json.dumps(state, indent=2), encoding="utf-8")
|
|
temporary.replace(path)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--container", required=True)
|
|
parser.add_argument("--process-marker", required=True)
|
|
parser.add_argument("--loop-state", type=Path, required=True)
|
|
parser.add_argument("--command-json", type=Path, required=True)
|
|
parser.add_argument("--host-script", type=Path, required=True)
|
|
parser.add_argument("--container-script", required=True)
|
|
parser.add_argument("--state", type=Path, required=True)
|
|
parser.add_argument("--poll-seconds", type=int, default=30)
|
|
parser.add_argument("--max-launches", type=int, default=20)
|
|
parser.add_argument("--once", action="store_true")
|
|
args = parser.parse_args()
|
|
state = {"schema_version": 1, "status": "monitoring", "launch_count": 0}
|
|
if args.state.is_file():
|
|
state.update(json.loads(args.state.read_text(encoding="utf-8")))
|
|
while True:
|
|
loop_status = read_loop_status(args.loop_state)
|
|
state["loop_status"] = loop_status
|
|
if loop_status in TERMINAL_STATUSES:
|
|
state["status"] = "loop_finished"
|
|
write_state(args.state, state)
|
|
return 0
|
|
if loop_status == "invalid":
|
|
state["status"] = "invalid_loop_state"
|
|
write_state(args.state, state)
|
|
return 2
|
|
if not process_active(args.container, args.process_marker):
|
|
if int(state["launch_count"]) >= args.max_launches:
|
|
state["status"] = "launch_budget_exhausted"
|
|
write_state(args.state, state)
|
|
return 3
|
|
command = json.loads(args.command_json.read_text(encoding="utf-8"))
|
|
if not isinstance(command, list) or not command or not all(isinstance(x, str) and x for x in command):
|
|
state["status"] = "invalid_command"
|
|
write_state(args.state, state)
|
|
return 4
|
|
copied = subprocess.run(
|
|
["docker", "cp", str(args.host_script), f"{args.container}:{args.container_script}"],
|
|
check=False,
|
|
)
|
|
launched = subprocess.run(command, check=False) if copied.returncode == 0 else copied
|
|
if launched.returncode != 0:
|
|
state["status"] = "launch_failed"
|
|
write_state(args.state, state)
|
|
return 5
|
|
state["launch_count"] = int(state["launch_count"]) + 1
|
|
state["status"] = "loop_relaunched"
|
|
else:
|
|
state["status"] = "monitoring"
|
|
write_state(args.state, state)
|
|
if args.once:
|
|
return 0
|
|
time.sleep(args.poll_seconds)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|