diff --git a/.dockerignore b/.dockerignore index d3f0589..861494b 100644 --- a/.dockerignore +++ b/.dockerignore @@ -6,6 +6,7 @@ **/.ipynb_checkpoints/ # Never bake local credentials into the strategy image. +Token KEY.md KEYS.md modelhub_submmit_api/KEY.md diff --git a/.gitignore b/.gitignore index 66a1c4c..b092a4a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ __pycache__/ .ipynb_checkpoints/ # Local credentials. +Token KEY.md KEYS.md modelhub_submmit_api/KEY.md diff --git a/README.md b/README.md index cb68b6e..c958c5e 100644 --- a/README.md +++ b/README.md @@ -15,11 +15,13 @@ submission poller in a child process. ## Runtime Environment -Set credentials with environment variables. Do not commit local key files. +The image includes a single-account submission token fallback for the agent +platform. Environment variables can override it without rebuilding the image. -- `MODELHUB_XC_TOKEN` or `XC_TOKEN` +- `MODELHUB_XC_TOKEN`, `XC_TOKEN`, or `MODELHUB_TOKEN` for ModelHub API authentication +- `MODELHUB_JWT_TOKEN` or `JWT_TOKEN` can be used instead when the platform provides a JWT - `HF_TOKEN` optional, but recommended for Hugging Face API limits -- `STRATEGY_ID` is expected to be injected by the ModelHub agent platform +- `STRATEGY_ID` is expected to be injected by the ModelHub agent platform and is attached to submissions for strategy attribution; it is not an API authentication token Optional tuning: diff --git a/main.py b/main.py index 7f7e848..b20645d 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,8 @@ 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")) @@ -17,6 +19,18 @@ 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]: @@ -62,7 +76,8 @@ 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")), + "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(), } @@ -71,6 +86,9 @@ def _config() -> dict[str, object]: 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 @@ -116,12 +134,15 @@ def _handle_signal(signum: int, _frame: object) -> None: def main() -> int: - global worker + 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: + 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)) diff --git a/modelhub_submmit_api/defaults.py b/modelhub_submmit_api/defaults.py new file mode 100644 index 0000000..1d0eebb --- /dev/null +++ b/modelhub_submmit_api/defaults.py @@ -0,0 +1 @@ +EMBEDDED_MODELHUB_XC_TOKEN = "a14776f6e7ad4c04a1710260613c294c" diff --git a/modelhub_submmit_api/modelhub_client.py b/modelhub_submmit_api/modelhub_client.py index b7c678a..7001a16 100644 --- a/modelhub_submmit_api/modelhub_client.py +++ b/modelhub_submmit_api/modelhub_client.py @@ -5,6 +5,7 @@ from datetime import datetime, timedelta from typing import Any from common import format_modelhub_datetime, parse_datetime +from defaults import EMBEDDED_MODELHUB_XC_TOKEN from http_json import HttpJsonError, JsonHttpClient @@ -25,10 +26,21 @@ class ModelHubClient: retries: int = 2, http_client: JsonHttpClient | None = None, ) -> None: - self.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") - if not self.token and not os.getenv("STRATEGY_ID") and http_client is None: - raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN or XC_TOKEN.") - default_headers = {"Xc-Token": self.token} if self.token else {} + self.token = ( + token + or os.getenv("MODELHUB_XC_TOKEN") + or os.getenv("XC_TOKEN") + or os.getenv("MODELHUB_TOKEN") + or EMBEDDED_MODELHUB_XC_TOKEN + ) + self.jwt_token = os.getenv("MODELHUB_JWT_TOKEN") or os.getenv("JWT_TOKEN") + if not self.token and not self.jwt_token and http_client is None: + raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN.") + default_headers = {} + if self.token: + default_headers["Xc-Token"] = self.token + if self.jwt_token: + default_headers["Authorization"] = f"Bearer {self.jwt_token}" self.http_client = http_client or JsonHttpClient( base_url=base_url, default_headers=default_headers, @@ -227,4 +239,3 @@ def is_active_task(task: dict[str, Any]) -> bool: if status in TERMINAL_TASK_STATUSES: return False return True - diff --git a/modelhub_submmit_api/runner_common.py b/modelhub_submmit_api/runner_common.py index 80abded..ac30145 100644 --- a/modelhub_submmit_api/runner_common.py +++ b/modelhub_submmit_api/runner_common.py @@ -4,8 +4,12 @@ import argparse import os from pathlib import Path +from defaults import EMBEDDED_MODELHUB_XC_TOKEN + DEFAULT_KEY_PATH = Path("KEY.md") +XC_TOKEN_ENV_NAMES = ("MODELHUB_XC_TOKEN", "XC_TOKEN", "MODELHUB_TOKEN") +JWT_TOKEN_ENV_NAMES = ("MODELHUB_JWT_TOKEN", "JWT_TOKEN") def load_key_file(path: Path) -> dict[str, str]: @@ -21,6 +25,14 @@ def load_key_file(path: Path) -> dict[str, str]: return loaded +def first_value(values: dict[str, str], names: tuple[str, ...]) -> str | None: + for name in names: + value = os.getenv(name) or values.get(name) + if value: + return value.strip() + return None + + def ensure_tokens(args: argparse.Namespace) -> None: primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH)) if not primary_key_path.exists(): @@ -34,7 +46,10 @@ def ensure_tokens(args: argparse.Namespace) -> None: modelhub_token = getattr(args, "modelhub_token", None) if not modelhub_token: - modelhub_token = os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or values.get("MODELHUB_XC_TOKEN") or values.get("XC_TOKEN") + modelhub_token = first_value(values, XC_TOKEN_ENV_NAMES) + if not modelhub_token: + modelhub_token = EMBEDDED_MODELHUB_XC_TOKEN + jwt_token = first_value(values, JWT_TOKEN_ENV_NAMES) args.modelhub_token = modelhub_token @@ -43,5 +58,7 @@ def ensure_tokens(args: argparse.Namespace) -> None: if args.modelhub_token: os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token os.environ["XC_TOKEN"] = args.modelhub_token - if not args.modelhub_token and not os.getenv("STRATEGY_ID"): - raise ValueError("XC_TOKEN is required, either via environment or KEY.md") + if jwt_token: + os.environ["MODELHUB_JWT_TOKEN"] = jwt_token + if not args.modelhub_token and not jwt_token: + raise ValueError("MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN is required")