Rebuild agent on original pooled runner

modelhub_submmit_api/main.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from typing import Any
 
 from common import parse_datetime, utc_now, write_json, write_jsonl
+from hf_discovery import HuggingFaceDiscovery
 from history_stats import (
     append_ledger_entry,
     build_empty_pre_submit_report,
@@ -15,9 +16,8 @@ from history_stats import (
     load_ledger,
     update_history_archive,
 )
-from modelhub_client import ModelHubAPIError, ModelHubClient, MultiModelHubClient
+from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool
 from models import CandidateModel, HFModelSummary, ModelInspection
-from modelscope_discovery import ModelScopeDiscovery
 from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
 from task_registry import TASK_SPEC_BY_TYPE, all_task_types, choose_framework_for_task, choose_text_generation_framework, pipeline_tags_for_task_types, task_specs_for_model
 from template_selector import TemplateSelector
@@ -74,7 +74,6 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
     parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
     parser.add_argument("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
-    parser.add_argument("--modelscope-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
     parser.add_argument("--modelhub-base-url", default=os.getenv("MODELHUB_BASE_URL", "https://modelhub.org.cn"), help=argparse.SUPPRESS)
     parser.add_argument("--modelhub-token", default=os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN"), help=argparse.SUPPRESS)
     return parser
@@ -195,7 +194,7 @@ def choose_candidate_for_gpu(
 def resolve_submit_concurrency(
     args: argparse.Namespace,
     *,
-    modelhub_client: ModelHubClient | Any,
+    modelhub_client: ModelHubClient | ModelHubClientPool | Any,
     planned_submit_count: int,
 ) -> int:
     explicit_workers = int(getattr(args, "submit_concurrency", 0) or 0)
@@ -228,7 +227,7 @@ def resolve_max_submit_count(
 def process_model_for_candidates(
     *,
     model: HFModelSummary,
-    hf_discovery: ModelScopeDiscovery,
+    hf_discovery: HuggingFaceDiscovery,
     modelhub_client: ModelHubClient,
     template_selector: TemplateSelector,
     target_gpus: list[str],
@@ -316,6 +315,11 @@ def submit_candidate(
             "responseData": response.get("data"),
         }
     except ModelHubAPIError as exc:
+        print(
+            f"[submit] failed repo={candidate['repoId']} gpu={candidate['targetGpu']} "
+            f"framework={candidate['framework']} reason={exc}",
+            flush=True,
+        )
         return {
             "outcome": "failed",
             "candidate": candidate,
@@ -340,8 +344,8 @@ def run_submission(
     args: argparse.Namespace,
     *,
     now=None,
-    hf_discovery: ModelScopeDiscovery | None = None,
-    modelhub_client: ModelHubClient | None = None,
+    hf_discovery: HuggingFaceDiscovery | None = None,
+    modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
     template_selector: TemplateSelector | None = None,
     outcome_tracker: OutcomeTracker | None = None,
 ) -> dict[str, Any]:
@@ -352,12 +356,12 @@ def run_submission(
     if not target_gpus:
         raise RuntimeError("No auto-submittable GPUs are available for the selected task types")
 
-    discovery_base_url = getattr(args, "modelscope_base_url", None) or args.hf_base_url
-    hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
+    hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=args.hf_base_url)
     if modelhub_client is None:
-        modelhub_tokens = list(getattr(args, "modelhub_tokens", []) or [])
+        modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
         if len(modelhub_tokens) > 1:
-            modelhub_client = MultiModelHubClient(tokens=modelhub_tokens, base_url=args.modelhub_base_url)
+            clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in modelhub_tokens]
+            modelhub_client = ModelHubClientPool(clients)
         else:
             modelhub_client = ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url)