168 lines
6.4 KiB
Python
168 lines
6.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Resume one checkpointed YOLO run after container recreation, fail closed."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
if str(SCRIPT_DIR) not in sys.path:
|
|
sys.path.insert(0, str(SCRIPT_DIR))
|
|
|
|
from training_release_manifest import ( # noqa: E402
|
|
TrainingReleaseError,
|
|
assert_training_release_eligible,
|
|
)
|
|
|
|
|
|
def container_running(container: str) -> bool:
|
|
result = subprocess.run(
|
|
["docker", "inspect", "--format", "{{.State.Running}}", container],
|
|
capture_output=True, text=True, check=False,
|
|
)
|
|
return result.returncode == 0 and result.stdout.strip() == "true"
|
|
|
|
|
|
def training_active(container: str, run_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(
|
|
"train" in line and run_marker in line
|
|
for line in result.stdout.splitlines()
|
|
)
|
|
|
|
|
|
def write_state(path: Path, payload: dict) -> None:
|
|
payload["updated_at"] = datetime.now(UTC).isoformat()
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
|
temporary.replace(path)
|
|
|
|
|
|
def load_completion_command(path: Path) -> list[str]:
|
|
command = json.loads(path.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):
|
|
raise ValueError("completion command must be a non-empty JSON list of non-empty strings")
|
|
return command
|
|
|
|
|
|
def verify_training_release(
|
|
*,
|
|
train_yaml: Path,
|
|
corpus_manifest: Path,
|
|
fixture_mode: bool,
|
|
) -> None:
|
|
"""Re-check the exact release before each detached resume command."""
|
|
|
|
assert_training_release_eligible(
|
|
train_yaml=train_yaml,
|
|
corpus_manifest=corpus_manifest,
|
|
fixture_mode=fixture_mode,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--container", required=True)
|
|
parser.add_argument("--host-run-dir", type=Path, required=True)
|
|
parser.add_argument("--container-checkpoint", required=True)
|
|
parser.add_argument("--train-yaml", type=Path, required=True)
|
|
parser.add_argument("--corpus-manifest", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--fixture-mode",
|
|
action="store_true",
|
|
help="Only valid for an explicitly fixture-only frozen corpus release.",
|
|
)
|
|
parser.add_argument("--run-marker", required=True)
|
|
parser.add_argument("--yolo", default="/opt/geointel/venv/bin/yolo")
|
|
parser.add_argument("--poll-seconds", type=int, default=30)
|
|
parser.add_argument("--max-resumes", type=int, default=20)
|
|
parser.add_argument("--completion-command-json", type=Path)
|
|
parser.add_argument("--once", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.poll_seconds < 1 or args.max_resumes < 1:
|
|
raise SystemExit("poll-seconds and max-resumes must be positive")
|
|
|
|
state_path = args.host_run_dir / "supervisor-state.json"
|
|
state = {"schema_version": 1, "status": "monitoring", "resume_count": 0}
|
|
if state_path.is_file():
|
|
state.update(json.loads(state_path.read_text(encoding="utf-8")))
|
|
|
|
while True:
|
|
if (args.host_run_dir / "results.png").is_file():
|
|
if state.get("completion_handoff_started"):
|
|
state["status"] = "training_finished_handoff_already_started"
|
|
write_state(state_path, state)
|
|
return 0
|
|
if args.completion_command_json:
|
|
try:
|
|
command = load_completion_command(args.completion_command_json)
|
|
except (OSError, json.JSONDecodeError, ValueError) as exc:
|
|
state["status"] = "invalid_completion_command"
|
|
state["completion_handoff_error"] = str(exc)
|
|
write_state(state_path, state)
|
|
return 4
|
|
result = subprocess.run(command, check=False)
|
|
if result.returncode != 0:
|
|
state["status"] = "completion_handoff_failed"
|
|
state["completion_handoff_returncode"] = result.returncode
|
|
write_state(state_path, state)
|
|
return 5
|
|
state["completion_handoff_started"] = True
|
|
state["completion_command"] = command
|
|
state["status"] = "training_finished_handoff_started"
|
|
else:
|
|
state["status"] = "training_finished"
|
|
write_state(state_path, state)
|
|
return 0
|
|
checkpoint = args.host_run_dir / "weights" / "last.pt"
|
|
if not checkpoint.is_file() or checkpoint.stat().st_size < 1024 * 1024:
|
|
state["status"] = "checkpoint_missing_or_incomplete"
|
|
write_state(state_path, state)
|
|
return 2
|
|
if container_running(args.container) and not training_active(args.container, args.run_marker):
|
|
if int(state["resume_count"]) >= args.max_resumes:
|
|
state["status"] = "resume_budget_exhausted"
|
|
write_state(state_path, state)
|
|
return 3
|
|
try:
|
|
verify_training_release(
|
|
train_yaml=args.train_yaml,
|
|
corpus_manifest=args.corpus_manifest,
|
|
fixture_mode=args.fixture_mode,
|
|
)
|
|
except TrainingReleaseError as exc:
|
|
state["status"] = "training_release_verification_failed"
|
|
state["training_release_error"] = str(exc)
|
|
write_state(state_path, state)
|
|
return 6
|
|
result = subprocess.run(
|
|
["docker", "exec", "-d", args.container, args.yolo, "train",
|
|
f"resume={args.container_checkpoint}", "device=0"],
|
|
check=False,
|
|
)
|
|
if result.returncode == 0:
|
|
state["resume_count"] = int(state["resume_count"]) + 1
|
|
state["status"] = "resumed"
|
|
else:
|
|
state["status"] = "resume_command_failed"
|
|
write_state(state_path, state)
|
|
else:
|
|
state["status"] = "monitoring"
|
|
write_state(state_path, state)
|
|
if args.once:
|
|
return 0
|
|
time.sleep(args.poll_seconds)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|