Add embedded ModelHub auth fallback

modelhub_submmit_api/defaults.py

@@ -0,0 +1 @@
+EMBEDDED_MODELHUB_XC_TOKEN = "a14776f6e7ad4c04a1710260613c294c"

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
-

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