adapt for ModelHub agent platform

2026-07-10 00:22:50 +08:00
commit 8c4bbff413
28 changed files with 3991 additions and 0 deletions
--- a/main.py
+++ b/main.py
@@ -0,0 +1,150 @@
+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
+
+
+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
+
+
+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_token_present": bool(os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN")),
+        "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 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 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:
+        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())