submmit/main.py

from __future__ import annotations

import json
import os
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from modelhub_submmit_api.defaults import EMBEDDED_MODELHUB_XC_TOKEN, EMBEDDED_MODELSCOPE_TOKEN
from modelhub_submmit_api.version import AGENT_VERSION


HOST = "0.0.0.0"
PORT = int(os.getenv("PORT", "8080"))
ROOT = Path(__file__).resolve().parent
WORKER_SCRIPT = ROOT / "modelhub_submmit_api" / "poll_runner.py"
READINESS_PATH = ROOT / ".modelhub_state" / "readiness.json"
WORKER_CRASH_PATH = ROOT / ".modelhub_state" / "worker_crashes.jsonl"
WORKER_STABLE_SECONDS = max(60, int(os.getenv("MODELHUB_AGENT_WORKER_STABLE_SECONDS", "600")))
WORKER_RESTART_MAX_SECONDS = max(30, int(os.getenv("MODELHUB_AGENT_WORKER_RESTART_MAX_SECONDS", "300")))

shutdown_requested = False
worker: subprocess.Popen | None = None
config_error: str | None = None
worker_start_enabled = False
worker_started_at: float | None = None
worker_next_restart_at = 0.0
worker_restart_count = 0
worker_last_exit_code: int | None = None


def _has_modelhub_auth() -> bool:
    return bool(
        os.getenv("MODELHUB_XC_TOKEN")
        or os.getenv("XC_TOKEN")
        or os.getenv("MODELHUB_TOKEN")
        or os.getenv("MODELHUB_JWT_TOKEN")
        or os.getenv("JWT_TOKEN")
        or EMBEDDED_MODELHUB_XC_TOKEN
    )


def _csv_args(env_name: str) -> list[str]:
    value = os.getenv(env_name, "").strip()
    if not value:
        return []
    return [part.strip() for part in value.split() if part.strip()]


def _worker_command() -> list[str]:
    cmd = [
        sys.executable,
        "-u",
        str(WORKER_SCRIPT),
        "--poll-interval-seconds",
        os.getenv("MODELHUB_AGENT_POLL_INTERVAL_SECONDS", "15"),
        "--idle-interval-seconds",
        os.getenv("MODELHUB_AGENT_IDLE_INTERVAL_SECONDS", "60"),
        "--post-cycle-cooldown-seconds",
        os.getenv("MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS", "2"),
        "--skip-history-archive",
        "--state-sync",
    ]

    daily_target = os.getenv("MODELHUB_AGENT_DAILY_TARGET", "").strip()
    if daily_target:
        cmd.extend(["--daily-target", daily_target])

    min_downloads = os.getenv("MODELHUB_AGENT_MIN_DOWNLOADS", "").strip()
    if min_downloads:
        cmd.extend(["--min-downloads", min_downloads])

    gpus = os.getenv("MODELHUB_AGENT_GPUS", "").strip()
    if gpus:
        cmd.extend(["--gpus", gpus])

    cmd.extend(_csv_args("MODELHUB_AGENT_EXTRA_ARGS"))
    # The hosted agent's goal is to fill every currently available slot. Keep
    # this last so an old environment or extra-args value cannot restore "5".
    cmd.extend(["--max-submits-per-run", "0"])
    return cmd


def _config() -> dict[str, object]:
    return {
        "agent_version": AGENT_VERSION,
        "strategy_id_present": bool(os.getenv("STRATEGY_ID")),
        "modelscope_token_present": bool(os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN),
        "modelhub_auth_present": _has_modelhub_auth(),
        "config_error": config_error,
        "worker_running": worker is not None and worker.poll() is None,
        "worker_return_code": worker_last_exit_code if worker is None else worker.poll(),
        "worker_restart_count": worker_restart_count,
        "worker_restart_pending": bool(
            worker_start_enabled
            and worker is None
            and not shutdown_requested
            and config_error is None
        ),
    }


def _readiness() -> dict[str, object]:
    try:
        payload = json.loads(READINESS_PATH.read_text(encoding="utf-8"))
    except (FileNotFoundError, OSError, ValueError, TypeError):
        return {"ready": False, "reason": "readiness_not_reported"}
    return payload if isinstance(payload, dict) else {"ready": False, "reason": "readiness_invalid"}


class Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:
        if self.path == "/health":
            # This endpoint is a liveness probe for the supervisor. Worker
            # readiness is reported separately by /ready.
            status = "ok" if worker is not None and worker.poll() is None else "degraded"
            if config_error:
                status = "config_error"
            self._send_json({"status": status, "config": _config()})
            return

        if self.path == "/ready":
            readiness = _readiness()
            status = 200 if readiness.get("ready") is True else 503
            self._send_json(
                {
                    "status": "ready" if status == 200 else "not_ready",
                    "readiness": readiness,
                    "config": _config(),
                },
                status=status,
            )
            return

        if self.path == "/":
            self._send_json({"name": "modelhub-submmit-agent", "status": "running", "config": _config()})
            return

        self._send_json({"error": "not found"}, status=404)

    def log_message(self, fmt: str, *args: object) -> None:
        print(f"{self.address_string()} - {fmt % args}", flush=True)

    def _send_json(self, body: dict[str, object], status: int = 200) -> None:
        payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)


def _stop_worker() -> None:
    global worker
    if worker is None or worker.poll() is not None:
        return
    print("stopping submission worker", flush=True)
    worker.terminate()
    try:
        worker.wait(timeout=25)
    except subprocess.TimeoutExpired:
        worker.kill()
        worker.wait(timeout=5)


def _restart_delay(restart_count: int) -> int:
    return min(WORKER_RESTART_MAX_SECONDS, 5 * (2 ** min(max(0, restart_count - 1), 6)))


def _write_supervisor_readiness(reason: str, *, exit_code: int | None = None) -> None:
    READINESS_PATH.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "ready": False,
        "reason": reason,
        "updatedAt": datetime.now(timezone.utc).isoformat(),
        "supervisor": {
            "workerRestartCount": worker_restart_count,
            "lastExitCode": exit_code,
        },
    }
    temporary = READINESS_PATH.with_name(f".{READINESS_PATH.name}.supervisor-{os.getpid()}")
    temporary.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
    os.replace(temporary, READINESS_PATH)


def _record_worker_crash(*, exit_code: int | None, uptime_seconds: float, reason: str) -> None:
    WORKER_CRASH_PATH.parent.mkdir(parents=True, exist_ok=True)
    event = {
        "at": datetime.now(timezone.utc).isoformat(),
        "exitCode": exit_code,
        "uptimeSeconds": round(max(0.0, uptime_seconds), 3),
        "restartCount": worker_restart_count,
        "reason": reason[:500],
    }
    with WORKER_CRASH_PATH.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(event, ensure_ascii=False, separators=(",", ":")) + "\n")


def _launch_worker(*, now: float | None = None) -> bool:
    global worker, worker_started_at, worker_last_exit_code
    now = time.monotonic() if now is None else now
    cmd = _worker_command()
    print("starting submission worker: " + " ".join(cmd), flush=True)
    try:
        worker = subprocess.Popen(cmd, cwd=str(ROOT))
    except OSError as exc:
        worker = None
        worker_started_at = None
        worker_last_exit_code = None
        print(f"submission worker spawn failed: {type(exc).__name__}: {exc}", flush=True)
        return False
    worker_started_at = now
    worker_last_exit_code = None
    return True


def _supervise_worker(*, now: float | None = None) -> None:
    global worker, worker_started_at, worker_next_restart_at
    global worker_restart_count, worker_last_exit_code
    if not worker_start_enabled or shutdown_requested or config_error:
        return
    now = time.monotonic() if now is None else now
    if worker is not None:
        return_code = worker.poll()
        if return_code is None:
            if worker_started_at is not None and now - worker_started_at >= WORKER_STABLE_SECONDS:
                worker_restart_count = 0
            return
        uptime = max(0.0, now - worker_started_at) if worker_started_at is not None else 0.0
        if uptime >= WORKER_STABLE_SECONDS:
            worker_restart_count = 0
        worker_restart_count += 1
        worker_last_exit_code = int(return_code)
        delay = _restart_delay(worker_restart_count)
        worker_next_restart_at = now + delay
        print(
            f"submission worker exited code={return_code} uptime={uptime:.1f}s "
            f"restart_in={delay}s restart_count={worker_restart_count}",
            flush=True,
        )
        _record_worker_crash(exit_code=return_code, uptime_seconds=uptime, reason="worker_exited")
        _write_supervisor_readiness("worker_restarting", exit_code=return_code)
        worker = None
        worker_started_at = None
        return
    if now < worker_next_restart_at:
        return
    if not _launch_worker(now=now):
        worker_restart_count += 1
        delay = _restart_delay(worker_restart_count)
        worker_next_restart_at = now + delay
        _record_worker_crash(exit_code=None, uptime_seconds=0.0, reason="worker_spawn_failed")
        _write_supervisor_readiness("worker_spawn_failed")


def _handle_signal(signum: int, _frame: object) -> None:
    global shutdown_requested
    shutdown_requested = True
    print(f"received signal {signum}, shutting down", flush=True)
    _stop_worker()


def main() -> int:
    global config_error, worker_start_enabled
    print(f"modelhub-submmit-agent version={AGENT_VERSION}", flush=True)
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    worker_start_enabled = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
    if worker_start_enabled and not _has_modelhub_auth():
        config_error = "missing ModelHub auth: set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN"
        print(config_error, flush=True)
    elif worker_start_enabled:
        _launch_worker()
    else:
        print("submission worker disabled by MODELHUB_AGENT_START_WORKER", flush=True)

    server = ThreadingHTTPServer((HOST, PORT), Handler)
    server.timeout = 1
    print(f"modelhub-submmit-agent listening on {HOST}:{PORT}", flush=True)

    try:
        while not shutdown_requested:
            server.handle_request()
            _supervise_worker()
            time.sleep(0.1)
    finally:
        server.server_close()
        _stop_worker()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())