Switch agent discovery to ModelScope

modelhub_submmit_api/runner_common.py

@@ -4,12 +4,15 @@ import argparse
 import os
 from pathlib import Path
 
-from defaults import EMBEDDED_HF_TOKEN, EMBEDDED_MODELHUB_XC_TOKEN
+from defaults import EMBEDDED_HF_TOKEN, EMBEDDED_MODELHUB_XC_TOKENS, EMBEDDED_MODELSCOPE_TOKEN
 
 
 DEFAULT_KEY_PATH = Path("KEY.md")
+MODULE_DIR = Path(__file__).resolve().parent
 XC_TOKEN_ENV_NAMES = ("MODELHUB_XC_TOKEN", "XC_TOKEN", "MODELHUB_TOKEN")
+XC_TOKEN_LIST_ENV_NAMES = ("MODELHUB_XC_TOKENS", "XC_TOKENS", "MODELHUB_TOKENS")
 JWT_TOKEN_ENV_NAMES = ("MODELHUB_JWT_TOKEN", "JWT_TOKEN")
+MODELSCOPE_TOKEN_ENV_NAMES = ("MODELSCOPE_API_TOKEN", "MODELSCOPE_TOKEN")
 
 
 def load_key_file(path: Path) -> dict[str, str]:
@@ -33,33 +36,100 @@ def first_value(values: dict[str, str], names: tuple[str, ...]) -> str | None:
     return None
 
 
+def split_token_list(value: str | None) -> list[str]:
+    if not value:
+        return []
+    tokens: list[str] = []
+    for normalized in value.replace(",", "\n").replace(";", "\n").splitlines():
+        token = normalized.strip()
+        if token:
+            tokens.append(token)
+    return tokens
+
+
+def discover_key_paths(primary_path: Path) -> list[Path]:
+    candidates: list[Path] = []
+    search_dirs = [primary_path.parent, primary_path.parent.parent]
+    for path in [primary_path, primary_path.parent / "KEYS.md"]:
+        if path not in candidates:
+            candidates.append(path)
+    for directory in search_dirs:
+        if not directory.exists():
+            continue
+        for path in sorted(directory.glob("KEYS*")):
+            if path.is_file() and path not in candidates:
+                candidates.append(path)
+    return candidates
+
+
+def collect_modelhub_tokens(values_by_path: list[dict[str, str]], explicit_token: str | None) -> list[str]:
+    tokens: list[str] = []
+
+    def add(value: str | None) -> None:
+        for token in split_token_list(value):
+            if token and token not in tokens:
+                tokens.append(token)
+
+    add(explicit_token)
+    for env_name in XC_TOKEN_LIST_ENV_NAMES:
+        add(os.getenv(env_name))
+    for env_name in XC_TOKEN_ENV_NAMES:
+        add(os.getenv(env_name))
+
+    for values in values_by_path:
+        for key, value in values.items():
+            normalized = key.strip().upper()
+            if normalized in XC_TOKEN_ENV_NAMES or normalized.startswith("XC_TOKEN") or normalized.startswith("MODELHUB_XC_TOKEN"):
+                add(value)
+
+    for token in EMBEDDED_MODELHUB_XC_TOKENS:
+        add(token)
+    return tokens
+
+
 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)
+        fallback_paths = [
+            MODULE_DIR / primary_key_path.name,
+            primary_key_path.parent.parent / primary_key_path.name,
+            MODULE_DIR.parent / primary_key_path.name,
+        ]
+        for fallback_path in fallback_paths:
+            if fallback_path.exists():
+                primary_key_path = fallback_path
+                break
+    values_by_path = [load_key_file(path) for path in discover_key_paths(primary_key_path) if path.exists()]
+    values: dict[str, str] = {}
+    for loaded in values_by_path:
+        values.update(loaded)
 
     if not getattr(args, "hf_token", None):
         args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN")
     if not getattr(args, "hf_token", None):
         args.hf_token = EMBEDDED_HF_TOKEN
+    if not getattr(args, "modelscope_token", None):
+        args.modelscope_token = first_value(values, MODELSCOPE_TOKEN_ENV_NAMES)
+    if not getattr(args, "modelscope_token", None):
+        args.modelscope_token = EMBEDDED_MODELSCOPE_TOKEN
 
-    modelhub_token = getattr(args, "modelhub_token", None)
-    if not modelhub_token:
-        modelhub_token = first_value(values, XC_TOKEN_ENV_NAMES)
-    if not modelhub_token:
-        modelhub_token = EMBEDDED_MODELHUB_XC_TOKEN
+    modelhub_tokens = collect_modelhub_tokens(values_by_path, getattr(args, "modelhub_token", None))
+    modelhub_token = modelhub_tokens[0] if modelhub_tokens else None
     jwt_token = first_value(values, JWT_TOKEN_ENV_NAMES)
 
     args.modelhub_token = modelhub_token
+    args.modelhub_tokens = modelhub_tokens
 
     if args.hf_token:
         os.environ["HF_TOKEN"] = args.hf_token
+    if args.modelscope_token:
+        os.environ["MODELSCOPE_API_TOKEN"] = args.modelscope_token
+        os.environ["MODELSCOPE_TOKEN"] = args.modelscope_token
     if args.modelhub_token:
         os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
         os.environ["XC_TOKEN"] = args.modelhub_token
+    if args.modelhub_tokens:
+        os.environ["MODELHUB_XC_TOKENS"] = ",".join(args.modelhub_tokens)
     if jwt_token:
         os.environ["MODELHUB_JWT_TOKEN"] = jwt_token
     if not args.modelhub_token and not jwt_token: