adapt for ModelHub agent platform

modelhub_submmit_api/runner_common.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+
+DEFAULT_KEY_PATH = Path("KEY.md")
+
+
+def load_key_file(path: Path) -> dict[str, str]:
+    if not path.exists():
+        return {}
+    loaded: dict[str, str] = {}
+    for raw_line in path.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        key, value = line.split("=", 1)
+        loaded[key.strip()] = value.strip()
+    return loaded
+
+
+def ensure_tokens(args: argparse.Namespace) -> None:
+    primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
+    if not primary_key_path.exists():
+        parent_key = primary_key_path.parent.parent / primary_key_path.name
+        if parent_key.exists():
+            primary_key_path = parent_key
+    values = load_key_file(primary_key_path)
+
+    if not getattr(args, "hf_token", None):
+        args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN")
+
+    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")
+
+    args.modelhub_token = modelhub_token
+
+    if args.hf_token:
+        os.environ["HF_TOKEN"] = args.hf_token
+    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")