submmit/main.py

from __future__ import annotations

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

from modelhub_submmit_api.defaults import EMBEDDED_MODELHUB_XC_TOKEN


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"

shutdown_requested = False
worker: subprocess.Popen | None = None
config_error: str | 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", "300"),
        "--idle-interval-seconds",
        os.getenv("MODELHUB_AGENT_IDLE_INTERVAL_SECONDS", "600"),
        "--post-cycle-cooldown-seconds",
        os.getenv("MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS", "30"),
        "--max-submits-per-run",
        os.getenv("MODELHUB_AGENT_MAX_SUBMITS_PER_RUN", "5"),
        "--skip-history-archive",
    ]

    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"))
    return cmd


def _config() -> dict[str, object]:
    return {
        "strategy_id_present": bool(os.getenv("STRATEGY_ID")),
        "hf_token_present": bool(os.getenv("HF_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": None if worker is None else worker.poll(),
    }


class Handler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:
        if self.path == "/health":
            if config_error:
                self._send_json({"status": "config_error", "config": _config()}, status=500)
                return
            if worker is not None and worker.poll() is not None:
                self._send_json({"status": "worker_exited", "config": _config()}, status=500)
                return
            self._send_json({"status": "ok"})
            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 _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
    signal.signal(signal.SIGTERM, _handle_signal)
    signal.signal(signal.SIGINT, _handle_signal)

    start_worker = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
    if start_worker 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 start_worker:
        cmd = _worker_command()
        print("starting submission worker: " + " ".join(cmd), flush=True)
        worker = subprocess.Popen(cmd, cwd=str(ROOT))
    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()
            if worker is not None and worker.poll() is not None:
                print(f"submission worker exited with code {worker.returncode}", flush=True)
                return worker.returncode or 1
            time.sleep(0.1)
    finally:
        server.server_close()
        _stop_worker()

    return 0


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