From c06169d9066c83040632e3bf314fee73a4676ce2 Mon Sep 17 00:00:00 2001 From: CoolBoy Date: Fri, 10 Jul 2026 02:02:08 +0800 Subject: [PATCH] Rebuild agent on original pooled runner --- README.md | 8 +- modelhub_submmit_api/README.md | 21 +- modelhub_submmit_api/daily_runner.py | 17 +- modelhub_submmit_api/hf_discovery.py | 202 +++++++++----- modelhub_submmit_api/main.py | 26 +- modelhub_submmit_api/modelhub_client.py | 263 ++++++++++++------ modelhub_submmit_api/models.py | 5 +- modelhub_submmit_api/modelscope_discovery.py | 243 ---------------- modelhub_submmit_api/outcome_tracker.py | 7 +- modelhub_submmit_api/poll_runner.py | 20 +- modelhub_submmit_api/runner_common.py | 158 ++++++----- .../success_history_runner.py | 24 +- 12 files changed, 459 insertions(+), 535 deletions(-) delete mode 100644 modelhub_submmit_api/modelscope_discovery.py diff --git a/README.md b/README.md index 63df6af..6fef9ac 100644 --- a/README.md +++ b/README.md @@ -15,11 +15,11 @@ submission poller in a child process. ## Runtime Environment -The image includes single-account ModelHub and ModelScope token fallbacks for +The image includes multi-account ModelHub and ModelScope token fallbacks for the agent platform. Environment variables can override them without rebuilding the image. -- `MODELHUB_XC_TOKEN`, `XC_TOKEN`, or `MODELHUB_TOKEN` for ModelHub API authentication +- `MODELHUB_XC_TOKEN`, `XC_TOKEN`, `XC_TOKEN2...`, or `MODELHUB_XC_TOKENS` for ModelHub API authentication - `MODELHUB_JWT_TOKEN` or `JWT_TOKEN` can be used instead when the platform provides a JWT - `MODELSCOPE_API_TOKEN` or `MODELSCOPE_TOKEN` optional override for the embedded ModelScope fallback token - `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 @@ -40,6 +40,6 @@ Optional tuning: Create a tag and submit the repository URL plus tag in "我的适配智能体". ```bash -git tag agent-v4 -git push origin agent-v4 +git tag agent-v6 +git push origin agent-v6 ``` diff --git a/modelhub_submmit_api/README.md b/modelhub_submmit_api/README.md index 28730a8..604598d 100644 --- a/modelhub_submmit_api/README.md +++ b/modelhub_submmit_api/README.md @@ -6,7 +6,7 @@ It currently supports: - one-shot submission planning via `main.py` - daily batch execution via `run_daily.sh` - continuous queue refill via `run_poll.sh` -- a single ModelHub token read from environment variables or `KEY.md` +- multiple ModelHub tokens read from `KEY.md` and `KEYS.md` - automatic task/framework/template selection across the supported GPU catalog ## Layout @@ -15,8 +15,8 @@ It currently supports: - `daily_runner.py`: daily wave orchestration - `poll_runner.py`: long-running queue refiller - `runner_common.py`: shared token / key file loading -- `modelscope_discovery.py`: ModelScope model discovery and inspection -- `modelhub_client.py`: ModelHub API client +- `hf_discovery.py`: ModelScope model discovery and inspection (keeps the legacy module name) +- `modelhub_client.py`: ModelHub API client and token-pool routing - `history_stats.py`: online history aggregation, ranking, and warnings - `template_selector.py`: template lookup and GPU normalization - `task_registry.py`: task-type and framework selection rules @@ -24,11 +24,12 @@ It currently supports: ## Key Files -- `KEY.md`: optional local ModelScope and ModelHub tokens +- `KEY.md`: primary ModelScope and ModelHub tokens +- `KEYS.md`: optional supplemental ModelHub tokens - `templates/public_submit/adapt_task_templates.jsonl`: public submit templates -The agent platform deployment should use environment variables instead of key -files. Set `MODELHUB_XC_TOKEN` or `XC_TOKEN` for the single submitting account. +The runner reads both files automatically. Add more accounts by appending +`XC_TOKEN3`, `XC_TOKEN4`, and so on to `KEYS.md`. Template lookup is also relative. The selector searches from the current working directory and the module directory. The primary project layout is: @@ -72,7 +73,7 @@ bash run_poll.sh --dry-run - The runner auto-discovers all safe GPU/template combinations from the public submit catalog. - Each model can be submitted at most once per GPU. -- One ModelHub account is used per deployed agent instance. +- Multiple ModelHub tokens are pooled and used to route submissions to the account with available async capacity. - History stats are online-only. The local ledger is used for local accounting, but platform history is only used after the local ledger reaches the configured threshold. - The default history threshold is `500` records. @@ -83,7 +84,7 @@ Common flags: - `--daily-target`: total target submissions for the day; `0` means unlimited - `--min-downloads`: ModelScope download floor - `--history-stats-threshold`: local ledger threshold before using online history stats -- `--max-scan-models`: hard cap on scanned ModelScope models for a run (0 = auto) +- `--max-scan-models`: hard cap on scanned HF models for a run (0 = auto) - `--scan-multiplier`: multiplier used for auto scan cap derivation from quota/queue capacity - `--read-concurrency`: concurrent HTTP reads while scanning model candidates (default 4) - `--max-submits-per-run`: max tasks to submit per run cycle (0 = unlimited) @@ -95,9 +96,9 @@ Common flags: `run_poll.sh` adds: -- `--poll-interval-seconds`: sleep when the account has no available slots +- `--poll-interval-seconds`: sleep when all accounts are saturated - `--idle-interval-seconds`: sleep when a cycle submits nothing -- `--max-scan-models`: hard cap on scanned ModelScope models for this cycle (0 = auto) +- `--max-scan-models`: hard cap on scanned HF models for this cycle (0 = auto) - `--scan-multiplier`: multiplier used for auto scan cap derivation from quota/queue capacity - `--max-submits-per-run`: max tasks to submit per poll cycle (0 = unlimited) - `--skip-outcome-sync`: skip outcome sync before scanning diff --git a/modelhub_submmit_api/daily_runner.py b/modelhub_submmit_api/daily_runner.py index 8b75cd7..812d76a 100644 --- a/modelhub_submmit_api/daily_runner.py +++ b/modelhub_submmit_api/daily_runner.py @@ -7,9 +7,9 @@ from pathlib import Path from typing import Any, Callable from common import utc_now, write_json +from hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission -from modelhub_client import ModelHubClient, MultiModelHubClient -from modelscope_discovery import ModelScopeDiscovery +from modelhub_client import ModelHubClient, ModelHubClientPool from outcome_tracker import OutcomeTracker from runner_common import DEFAULT_KEY_PATH, ensure_tokens from template_selector import TemplateSelector @@ -85,7 +85,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS) parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), 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) parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN"), help=argparse.SUPPRESS) @@ -121,7 +120,6 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar history_archive_path=base_args.history_archive_path, history_archive_limit=base_args.history_archive_limit, hf_base_url=base_args.hf_base_url, - modelscope_base_url=getattr(base_args, "modelscope_base_url", base_args.hf_base_url), modelhub_base_url=base_args.modelhub_base_url, modelhub_token=base_args.modelhub_token, ) @@ -133,18 +131,18 @@ def run_daily_batches( waves: tuple[WaveSpec, ...] = DEFAULT_WAVES, now=None, run_fn: Callable[..., dict[str, Any]] = run_submission, - hf_discovery: ModelScopeDiscovery | None = None, + hf_discovery: HuggingFaceDiscovery | None = None, modelhub_client: ModelHubClient | None = None, template_selector: TemplateSelector | None = None, outcome_tracker: OutcomeTracker | None = None, ) -> dict[str, Any]: now = now or utc_now() - discovery_base_url = getattr(base_args, "modelscope_base_url", None) or base_args.hf_base_url - hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url) + hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url) if modelhub_client is None: - modelhub_tokens = list(getattr(base_args, "modelhub_tokens", []) or []) + modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token])) if len(modelhub_tokens) > 1: - modelhub_client = MultiModelHubClient(tokens=modelhub_tokens, base_url=base_args.modelhub_base_url) + clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens] + modelhub_client = ModelHubClientPool(clients) else: modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url) template_selector = template_selector or TemplateSelector() @@ -293,6 +291,7 @@ def main(argv: list[str] | None = None) -> int: ensure_tokens(args) log( f"[daily] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} " + f"xc_token={'set' if bool(args.modelhub_token) else 'missing'} " f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}" ) summary = run_daily_batches(base_args=args) diff --git a/modelhub_submmit_api/hf_discovery.py b/modelhub_submmit_api/hf_discovery.py index e06a628..d7cddb4 100644 --- a/modelhub_submmit_api/hf_discovery.py +++ b/modelhub_submmit_api/hf_discovery.py @@ -1,12 +1,14 @@ from __future__ import annotations -from concurrent.futures import ThreadPoolExecutor, as_completed import os import threading from pathlib import PurePosixPath from typing import Any +from urllib.parse import quote from common import parse_datetime +from defaults import EMBEDDED_MODELSCOPE_TOKEN +from http_json import HttpJsonError from http_json import JsonHttpClient from models import HFModelSummary, ModelInspection @@ -15,20 +17,47 @@ GGUF_PRIORITY = ("q4_0.gguf", "q8_0.gguf", "fp16.gguf") VLLM_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth") ONNX_WEIGHT_SUFFIXES = (".onnx",) +MODELSCOPE_TASK_TAGS = { + "text-generation": "text-generation", + "image-text-to-text": "image-text-to-text", + "visual-question-answering": "visual-question-answering", + "document-question-answering": "document-question-answering", + "video-text-to-text": "video-text-to-text", + "text-to-image": "text-to-image-synthesis", + "image-to-image": "image-to-image", + "automatic-speech-recognition": "auto-speech-recognition", + "question-answering": "question-answering", + "feature-extraction": "feature-extraction", + "sentence-similarity": "sentence-similarity", + "image-classification": "image-classification", + "zero-shot-image-classification": "zero-shot-image-classification", + "text-classification": "text-classification", + "zero-shot-classification": "zero-shot-classification", + "reinforcement-learning": "reinforcement-learning", +} + class HuggingFaceDiscovery: def __init__( self, - base_url: str = "https://huggingface.co", + base_url: str = "https://modelscope.cn", http_client: JsonHttpClient | None = None, + legacy_http_client: JsonHttpClient | None = None, timeout: int = 30, retries: int = 2, ) -> None: - hf_token = os.getenv("HF_TOKEN") + token = os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN headers = {"User-Agent": "modelhub-submmit-cli/0.1"} - if hf_token: - headers["Authorization"] = f"Bearer {hf_token}" + if token: + headers["Authorization"] = f"Bearer {token}" + headers["Cookie"] = f"m_session_id={token}" self.http_client = http_client or JsonHttpClient( + base_url=f"{base_url.rstrip('/')}/openapi/v1", + default_headers=headers, + timeout=timeout, + retries=retries, + ) + self.legacy_http_client = legacy_http_client or JsonHttpClient( base_url=base_url, default_headers=headers, timeout=timeout, @@ -50,36 +79,15 @@ class HuggingFaceDiscovery: return [] deduped: dict[str, HFModelSummary] = {} - max_workers = min(len(pipeline_tags), max(1, read_concurrency)) + del read_concurrency - if max_workers <= 1: - tag_results = [ - self._query_recent_models( - pipeline_tag=pipeline_tag, - limit=limit, - min_downloads=min_downloads, - updated_after=updated_after, - ) - for pipeline_tag in pipeline_tags - ] - else: - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = { - executor.submit( - self._query_recent_models, - pipeline_tag=pipeline_tag, - limit=limit, - min_downloads=min_downloads, - updated_after=updated_after, - ): pipeline_tag - for pipeline_tag in pipeline_tags - } - tag_results = [] - for future in as_completed(futures): - tag_results.append(future.result()) - - for models in tag_results: - for model in models: + for pipeline_tag in pipeline_tags: + for model in self._query_recent_models( + pipeline_tag=pipeline_tag, + limit=limit, + min_downloads=min_downloads, + updated_after=updated_after, + ): current = deduped.get(model.repo_id) if current is None or (model.last_modified or parse_datetime("1970-01-01")) > ( current.last_modified or parse_datetime("1970-01-01") @@ -97,23 +105,42 @@ class HuggingFaceDiscovery: min_downloads: int, updated_after=None, ) -> list[HFModelSummary]: - payload = self.http_client.request_json( - "GET", - "/api/models", - query={ - "pipeline_tag": pipeline_tag, - "sort": "lastModified", - "direction": "-1", - "limit": limit, - "full": "true", - }, - ) - + page_size = min(max(1, limit), 100) + max_items = min(max(1, limit), 3000) + task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag) models: list[HFModelSummary] = [] - for item in payload or []: - model = self._parse_model(item, min_downloads=min_downloads, updated_after=updated_after) - if model is not None: - models.append(model) + for page_number in range(1, (max_items + page_size - 1) // page_size + 1): + try: + payload = self.http_client.request_json( + "GET", + "/models", + query={ + "page_number": page_number, + "page_size": page_size, + "sort": "last_modified", + "filter.task": task_tag, + }, + ) + except HttpJsonError as exc: + print(f"[modelscope] list_models_error task={task_tag} page={page_number} error={exc}", flush=True) + break + + items = self._extract_models(payload) + if not items: + break + for item in items: + model = self._parse_model( + item, + fallback_pipeline_tag=pipeline_tag, + min_downloads=min_downloads, + updated_after=updated_after, + ) + if model is not None: + models.append(model) + if len(models) >= max_items: + return models + if len(items) < page_size: + break return models def list_recent_text_generation_models( @@ -140,40 +167,67 @@ class HuggingFaceDiscovery: if cached is not None: return list(cached) - payload = self.http_client.request_json( - "GET", - f"/api/models/{repo_id}/tree/main", - query={"recursive": "1"}, - ) + encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/")) + try: + payload = self.legacy_http_client.request_json( + "GET", + f"/api/v1/models/{encoded_repo_id}/repo/files", + query={"Revision": "master", "Recursive": "true"}, + ) + except HttpJsonError as exc: + print(f"[modelscope] repo_tree_error repo={repo_id} error={exc}", flush=True) + payload = None - entries: list[dict[str, Any]] - if isinstance(payload, list): - entries = payload - elif isinstance(payload, dict): - for key in ("items", "tree", "siblings"): - value = payload.get(key) - if isinstance(value, list): - entries = value - break - else: - entries = [] - else: - entries = [] + entries = self._extract_files(payload) with self._repo_tree_lock: self._repo_tree_cache[repo_id] = list(entries) return list(entries) @staticmethod - def _parse_model(item: dict[str, Any], *, min_downloads: int, updated_after=None) -> HFModelSummary | None: - repo_id = item.get("id") or item.get("modelId") + def _extract_models(payload: Any) -> list[dict[str, Any]]: + data = payload.get("data") if isinstance(payload, dict) else payload + if isinstance(data, dict): + for key in ("models", "Models", "items", "list", "data", "results"): + value = data.get(key) + if isinstance(value, list): + return [item for item in value if isinstance(item, dict)] + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + return [] + + @staticmethod + def _extract_files(payload: Any) -> list[dict[str, Any]]: + data = payload.get("Data") if isinstance(payload, dict) else payload + if isinstance(data, dict): + for key in ("Files", "files", "items", "tree"): + value = data.get(key) + if isinstance(value, list): + return [item for item in value if isinstance(item, dict)] + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + return [] + + @staticmethod + def _parse_model( + item: dict[str, Any], + *, + fallback_pipeline_tag: str, + min_downloads: int, + updated_after=None, + ) -> HFModelSummary | None: + repo_id = item.get("id") or item.get("model_id") or item.get("modelId") + if not repo_id: + owner = item.get("owner") or item.get("Owner") or item.get("Path") + name = item.get("name") or item.get("Name") + repo_id = f"{owner}/{name}" if owner and name else None if not repo_id: return None - downloads = int(item.get("downloads") or 0) + downloads = int(item.get("downloads") or item.get("Downloads") or 0) if downloads < min_downloads: return None - pipeline_tag = item.get("pipeline_tag") or item.get("pipelineTag") - last_modified = parse_datetime(item.get("lastModified") or item.get("last_modified")) + pipeline_tag = fallback_pipeline_tag + last_modified = parse_datetime(item.get("last_modified") or item.get("UpdatedAt") or item.get("LastUpdatedTime")) if updated_after and last_modified and last_modified < updated_after: return None return HFModelSummary( @@ -181,7 +235,7 @@ class HuggingFaceDiscovery: downloads=downloads, last_modified=last_modified, pipeline_tag=pipeline_tag, - created_at=parse_datetime(item.get("createdAt") or item.get("created_at")), + created_at=parse_datetime(item.get("created_at") or item.get("CreatedAt")), ) diff --git a/modelhub_submmit_api/main.py b/modelhub_submmit_api/main.py index 3f7463a..0d6a8c0 100644 --- a/modelhub_submmit_api/main.py +++ b/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) diff --git a/modelhub_submmit_api/modelhub_client.py b/modelhub_submmit_api/modelhub_client.py index dd3aa39..e322395 100644 --- a/modelhub_submmit_api/modelhub_client.py +++ b/modelhub_submmit_api/modelhub_client.py @@ -1,9 +1,10 @@ from __future__ import annotations import os +import time +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta -from itertools import cycle -from threading import Lock from typing import Any from common import format_modelhub_datetime, parse_datetime @@ -28,21 +29,10 @@ 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") - 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.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or EMBEDDED_MODELHUB_XC_TOKEN + if not self.token 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.http_client = http_client or JsonHttpClient( base_url=base_url, default_headers=default_headers, @@ -207,73 +197,6 @@ class ModelHubClient: return payload -class MultiModelHubClient: - def __init__( - self, - *, - tokens: list[str], - base_url: str = "https://modelhub.org.cn", - timeout: int = 30, - retries: int = 2, - ) -> None: - deduped: list[str] = [] - for token in tokens: - clean = token.strip() - if clean and clean not in deduped: - deduped.append(clean) - if not deduped: - raise ValueError("At least one ModelHub token is required") - - self.clients = [ - ModelHubClient(token=token, base_url=base_url, timeout=timeout, retries=retries) - for token in deduped - ] - self._cycle = cycle(self.clients) - self._lock = Lock() - - def _next_client(self) -> ModelHubClient: - with self._lock: - return next(self._cycle) - - def add_task(self, payload: dict[str, Any]) -> dict[str, Any]: - return self._next_client().add_task(payload) - - def search_by_model_id(self, model_id: str) -> dict[str, Any]: - return self.clients[0].search_by_model_id(model_id) - - def is_model_processed_for_gpu(self, model_id: str, target_gpu: str) -> bool: - return self.clients[0].is_model_processed_for_gpu(model_id, target_gpu) - - def get_verify_result_map(self, model_id: str) -> dict[str, Any]: - return self.clients[0].get_verify_result_map(model_id) - - def processed_gpus_for_model(self, model_id: str) -> set[str]: - return self.clients[0].processed_gpus_for_model(model_id) - - def list_tasks_page(self, **kwargs: Any) -> dict[str, Any]: - return self.clients[0].list_tasks_page(**kwargs) - - def list_tasks(self, **kwargs: Any) -> list[dict[str, Any]]: - return self.clients[0].list_tasks(**kwargs) - - def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None: - return self.clients[0].find_recent_task_id(model_id, gpu_type, submitted_after) - - def count_active_tasks(self, **kwargs: Any) -> int: - return self.clients[0].count_active_tasks(**kwargs) - - def active_task_counts(self) -> list[int]: - counts: list[int] = [] - now = datetime.utcnow() - begin_time = now - timedelta(days=1) - for client in self.clients: - counts.append(client.count_active_tasks(begin_time=begin_time, end_time=now, max_count=100)) - return counts - - def available_submit_slots(self, max_active_per_account: int = 5) -> int: - return sum(max(0, max_active_per_account - count) for count in self.active_task_counts()) - - ACTIVE_TASK_STATUSES = { "waiting", "running", @@ -308,3 +231,175 @@ def is_active_task(task: dict[str, Any]) -> bool: if status in TERMINAL_TASK_STATUSES: return False return True + + +class ModelHubClientPool: + def __init__(self, clients: list[ModelHubClient], *, active_task_cap: int = 100) -> None: + if not clients: + raise ValueError("At least one ModelHub client is required") + self.clients = clients + self.active_task_cap = active_task_cap + self._active_counts: list[int] = [0 for _ in clients] + self._active_refresh_at: float = 0.0 + self._active_counts_ttl: float = 30.0 + self._state_lock = threading.Lock() + # Per-cycle cache for search_by_model_id results (model_id -> merged verify result map) + self._verify_cache: dict[str, dict[str, Any]] = {} + # Single reader client to avoid fanout on read operations + self._reader = clients[0] + + def _safe_count_active_tasks(self, client: ModelHubClient) -> int: + try: + return client.count_active_tasks(max_count=self.active_task_cap, page_size=200) + except Exception: + return self.active_task_cap + + def _safe_search_by_model_id(self, client: ModelHubClient, model_id: str) -> dict[str, Any]: + try: + payload = client.search_by_model_id(model_id) + return payload if isinstance(payload, dict) else {} + except Exception: + return {} + + def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + try: + return client.list_tasks(**kwargs) + except Exception: + return [] + + def _refresh_active_counts(self, *, force: bool = False) -> None: + now = time.time() + if (not force) and self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl: + return + + def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]: + return index, self._safe_count_active_tasks(client) + + results: list[tuple[int, int]] = [] + with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor: + futures = {executor.submit(_to_indexed_result, index, client): index for index, client in enumerate(self.clients)} + for future in as_completed(futures): + index = futures[future] + try: + results.append(future.result()) + except Exception: + results.append((index, self.active_task_cap)) + + self._active_counts = [0 for _ in self.clients] + for index, count in sorted(results, key=lambda item: item[0]): + self._active_counts[index] = count + # Use current time after refresh completes, not the stale 'now' from function start + self._active_refresh_at = time.time() + + def active_task_counts(self) -> list[int]: + with self._state_lock: + now = time.time() + if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl: + pass # use cached counts + else: + self._refresh_active_counts() + return list(self._active_counts) + + def available_submit_slots(self) -> int: + with self._state_lock: + # Only refresh if cache is stale (respects TTL) or counts are empty + now = time.time() + if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl: + pass # use cached counts + else: + self._refresh_active_counts() + return sum(max(0, self.active_task_cap - count) for count in self._active_counts) + + def list_tasks(self, **kwargs): # noqa: ANN003, ANN001 + # Default: use only the reader client to avoid fanout amplification. + # The fanout_all parameter allows explicit cross-account merging when needed. + fanout_all = kwargs.pop("_fanout_all", False) + if not fanout_all: + return self._safe_list_tasks(self._reader, kwargs) + + merged: list[dict[str, Any]] = [] + seen: set[str] = set() + with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor: + futures = {executor.submit(self._safe_list_tasks, client, kwargs): index for index, client in enumerate(self.clients)} + for future in as_completed(futures): + try: + tasks = future.result() + except Exception: + tasks = [] + for task in tasks: + task_id = str(task.get("taskId")) if task.get("taskId") is not None else None + if task_id and task_id in seen: + continue + if task_id: + seen.add(task_id) + merged.append(task) + return merged + + def _read_from_cache(self, model_id: str) -> dict[str, Any] | None: + cached = self._verify_cache.get(model_id) + if cached is None: + return None + return dict(cached) + + def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None: + self._verify_cache[model_id] = payload + + def search_by_model_id(self, model_id: str) -> dict[str, Any]: + # Check cache first (per-cycle cache to avoid repeated API calls for the same model) + with self._state_lock: + cached = self._read_from_cache(model_id) + if cached is not None: + return cached + + # Only query ONE client (the reader) instead of fanning out to all clients. + # verifyResult is model-specific platform data, not account-specific. + response = self._safe_search_by_model_id(self._reader, model_id) + if not isinstance(response, dict): + response = {"code": 0, "data": {"verifyResult": {}}} + + with self._state_lock: + self._write_to_cache(model_id, response) + return response + + def get_verify_result_map(self, model_id: str) -> dict[str, Any]: + payload = self.search_by_model_id(model_id) + return ((payload.get("data") or {}).get("verifyResult") or {}) + + def processed_gpus_for_model(self, model_id: str) -> set[str]: + return set(self.get_verify_result_map(model_id).keys()) + + def add_task(self, payload: dict[str, Any]) -> dict[str, Any]: + with self._state_lock: + self._refresh_active_counts(force=True) + selected_index = None + best_remaining = -1 + for index, active_count in enumerate(self._active_counts): + remaining = self.active_task_cap - active_count + if remaining > best_remaining: + best_remaining = remaining + selected_index = index + if selected_index is None or best_remaining <= 0: + raise ModelHubAPIError( + f"当前等待中或运行中的异步模型验证任务数量已达上限（{self.active_task_cap}）" + ) + self._active_counts[selected_index] += 1 + + selected_client = self.clients[selected_index] + try: + response = selected_client.add_task(payload) + except Exception: + with self._state_lock: + self._active_counts[selected_index] -= 1 + raise + return response + + def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001 + """Single-page task listing via the reader client (no fanout).""" + return self._reader.list_tasks_page(**kwargs) + + def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None: + for client in self.clients: + task_id = client.find_recent_task_id(model_id, gpu_type, submitted_after) + if task_id is not None: + return task_id + return None diff --git a/modelhub_submmit_api/models.py b/modelhub_submmit_api/models.py index 8e73592..d2dc088 100644 --- a/modelhub_submmit_api/models.py +++ b/modelhub_submmit_api/models.py @@ -23,13 +23,10 @@ class HFModelSummary: last_modified: datetime | None pipeline_tag: str | None created_at: datetime | None = None - source: str = "huggingface" @property def model_address(self) -> str: - if self.source == "modelscope": - return f"https://modelscope.cn/models/{self.repo_id}" - return f"https://huggingface.co/{self.repo_id}" + return f"https://modelscope.cn/models/{self.repo_id}" @dataclass(frozen=True) diff --git a/modelhub_submmit_api/modelscope_discovery.py b/modelhub_submmit_api/modelscope_discovery.py deleted file mode 100644 index 481998b..0000000 --- a/modelhub_submmit_api/modelscope_discovery.py +++ /dev/null @@ -1,243 +0,0 @@ -from __future__ import annotations - -import os -import threading -from typing import Any -from urllib.parse import quote - -from common import parse_datetime -from defaults import EMBEDDED_MODELSCOPE_TOKEN -from hf_discovery import inspect_repo_tree -from http_json import HttpJsonError, JsonHttpClient -from models import HFModelSummary, ModelInspection - - -MODELSCOPE_TASK_TAGS = { - "text-generation": "text-generation", - "image-text-to-text": "image-text-to-text", - "visual-question-answering": "visual-question-answering", - "document-question-answering": "document-question-answering", - "video-text-to-text": "video-text-to-text", - "text-to-image": "text-to-image-synthesis", - "image-to-image": "image-to-image", - "automatic-speech-recognition": "auto-speech-recognition", - "question-answering": "question-answering", - "feature-extraction": "feature-extraction", - "sentence-similarity": "sentence-similarity", - "image-classification": "image-classification", - "zero-shot-image-classification": "zero-shot-image-classification", - "text-classification": "text-classification", - "zero-shot-classification": "zero-shot-classification", - "reinforcement-learning": "reinforcement-learning", -} - - -class ModelScopeDiscovery: - def __init__( - self, - base_url: str = "https://modelscope.cn", - http_client: JsonHttpClient | None = None, - legacy_http_client: JsonHttpClient | None = None, - timeout: int = 30, - retries: int = 2, - ) -> None: - token = os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN - headers = {"User-Agent": "modelhub-submmit-cli/0.1"} - if token: - headers["Authorization"] = f"Bearer {token}" - headers["Cookie"] = f"m_session_id={token}" - self.http_client = http_client or JsonHttpClient( - base_url=f"{base_url.rstrip('/')}/openapi/v1", - default_headers=headers, - timeout=timeout, - retries=retries, - ) - self.legacy_http_client = legacy_http_client or JsonHttpClient( - base_url=base_url, - default_headers=headers, - timeout=timeout, - retries=retries, - ) - self._repo_tree_cache: dict[str, list[dict[str, Any]]] = {} - self._repo_tree_lock = threading.Lock() - - def list_recent_models( - self, - *, - pipeline_tags: list[str], - limit: int, - min_downloads: int, - updated_after=None, - read_concurrency: int = 1, - ) -> list[HFModelSummary]: - del read_concurrency - if not pipeline_tags: - return [] - - deduped: dict[str, HFModelSummary] = {} - for pipeline_tag in pipeline_tags: - for model in self._query_recent_models( - pipeline_tag=pipeline_tag, - limit=limit, - min_downloads=min_downloads, - updated_after=updated_after, - ): - current = deduped.get(model.repo_id) - if current is None or (model.last_modified or parse_datetime("1970-01-01")) > ( - current.last_modified or parse_datetime("1970-01-01") - ): - deduped[model.repo_id] = model - - models = list(deduped.values()) - models.sort(key=lambda item: item.last_modified or parse_datetime("1970-01-01"), reverse=True) - return models - - def _query_recent_models( - self, - *, - pipeline_tag: str, - limit: int, - min_downloads: int, - updated_after=None, - ) -> list[HFModelSummary]: - page_size = min(max(1, limit), 100) - max_items = min(max(1, limit), 3000) - task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag) - models: list[HFModelSummary] = [] - - for page_number in range(1, (max_items + page_size - 1) // page_size + 1): - try: - payload = self.http_client.request_json( - "GET", - "/models", - query={ - "page_number": page_number, - "page_size": page_size, - "sort": "last_modified", - "filter.task": task_tag, - }, - ) - except HttpJsonError as exc: - print(f"[modelscope] list_models_error task={task_tag} page={page_number} error={exc}", flush=True) - break - - items = self._extract_models(payload) - if not items: - break - for item in items: - model = self._parse_model( - item, - fallback_pipeline_tag=pipeline_tag, - min_downloads=min_downloads, - updated_after=updated_after, - ) - if model is not None: - models.append(model) - if len(models) >= max_items: - return models - if len(items) < page_size: - break - - return models - - def list_recent_text_generation_models( - self, - *, - limit: int, - min_downloads: int, - updated_after=None, - ) -> list[HFModelSummary]: - return self.list_recent_models( - pipeline_tags=["text-generation"], - limit=limit, - min_downloads=min_downloads, - updated_after=updated_after, - ) - - def inspect_model(self, model: HFModelSummary) -> ModelInspection: - entries = self.list_repo_tree(model.repo_id) - return inspect_repo_tree(model.repo_id, entries) - - def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]: - with self._repo_tree_lock: - cached = self._repo_tree_cache.get(repo_id) - if cached is not None: - return list(cached) - - encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/")) - try: - payload = self.legacy_http_client.request_json( - "GET", - f"/api/v1/models/{encoded_repo_id}/repo/files", - query={"Revision": "master", "Recursive": "true"}, - ) - except HttpJsonError as exc: - print(f"[modelscope] repo_tree_error repo={repo_id} error={exc}", flush=True) - payload = None - - entries = self._extract_files(payload) - with self._repo_tree_lock: - self._repo_tree_cache[repo_id] = list(entries) - return list(entries) - - @staticmethod - def _extract_models(payload: Any) -> list[dict[str, Any]]: - data = payload.get("data") if isinstance(payload, dict) else payload - if isinstance(data, dict): - for key in ("models", "Models", "items", "list", "data", "results"): - value = data.get(key) - if isinstance(value, list): - return [item for item in value if isinstance(item, dict)] - if isinstance(data, list): - return [item for item in data if isinstance(item, dict)] - return [] - - @staticmethod - def _extract_files(payload: Any) -> list[dict[str, Any]]: - data = payload.get("Data") if isinstance(payload, dict) else payload - if isinstance(data, dict): - for key in ("Files", "files", "items", "tree"): - value = data.get(key) - if isinstance(value, list): - return [item for item in value if isinstance(item, dict)] - if isinstance(data, list): - return [item for item in data if isinstance(item, dict)] - return [] - - @staticmethod - def _parse_model( - item: dict[str, Any], - *, - fallback_pipeline_tag: str, - min_downloads: int, - updated_after=None, - ) -> HFModelSummary | None: - repo_id = item.get("id") or item.get("model_id") or item.get("modelId") - if not repo_id: - owner = item.get("owner") or item.get("Owner") or item.get("Path") - name = item.get("name") or item.get("Name") - repo_id = f"{owner}/{name}" if owner and name else None - if not repo_id: - return None - - downloads = int(item.get("downloads") or item.get("Downloads") or 0) - if downloads < min_downloads: - return None - - tasks = item.get("tasks") or item.get("Tasks") or [] - if isinstance(tasks, list) and tasks: - pipeline_tag = str(tasks[0]) - else: - pipeline_tag = fallback_pipeline_tag - last_modified = parse_datetime(item.get("last_modified") or item.get("UpdatedAt") or item.get("LastUpdatedTime")) - if updated_after and last_modified and last_modified < updated_after: - return None - - return HFModelSummary( - repo_id=repo_id, - downloads=downloads, - last_modified=last_modified, - pipeline_tag=pipeline_tag, - created_at=parse_datetime(item.get("created_at") or item.get("CreatedAt")), - source="modelscope", - ) diff --git a/modelhub_submmit_api/outcome_tracker.py b/modelhub_submmit_api/outcome_tracker.py index 75cc70a..584ed79 100644 --- a/modelhub_submmit_api/outcome_tracker.py +++ b/modelhub_submmit_api/outcome_tracker.py @@ -7,7 +7,7 @@ from typing import Any from common import append_jsonl, parse_datetime, read_jsonl, utc_now, write_jsonl from history_stats import classify_failure, is_failure, is_success -from modelhub_client import ModelHubClient +from modelhub_client import ModelHubClient, ModelHubClientPool DEFAULT_OUTCOMES_PATH = Path("outcomes/submissions.jsonl") @@ -72,11 +72,14 @@ class OutcomeTracker: self._by_model_gpu[(model_id, target_gpu)].append(record) append_jsonl(self.path, record) - def sync_from_api(self, client: ModelHubClient) -> int: + def sync_from_api(self, client: ModelHubClient | ModelHubClientPool) -> int: try: begin = self._last_sync_time end = utc_now() + # Use fanout for pools to sync outcomes across all accounts list_kwargs: dict[str, Any] = {"begin_time": begin, "end_time": end, "page_size": 100, "only_mine": True} + if isinstance(client, ModelHubClientPool): + list_kwargs["_fanout_all"] = True tasks = client.list_tasks(**list_kwargs) except Exception: return 0 diff --git a/modelhub_submmit_api/poll_runner.py b/modelhub_submmit_api/poll_runner.py index 1e46b88..9f1a201 100644 --- a/modelhub_submmit_api/poll_runner.py +++ b/modelhub_submmit_api/poll_runner.py @@ -9,9 +9,9 @@ from typing import Any, Callable from common import utc_now, write_json from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches +from hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR -from modelhub_client import ModelHubClient, MultiModelHubClient -from modelscope_discovery import ModelScopeDiscovery +from modelhub_client import ModelHubClient, ModelHubClientPool from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker from runner_common import DEFAULT_KEY_PATH, ensure_tokens from template_selector import TemplateSelector @@ -64,7 +64,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS) parser.add_argument("--hf-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS) - parser.add_argument("--modelscope-base-url", default="https://modelscope.cn", help=argparse.SUPPRESS) parser.add_argument("--modelhub-base-url", default="https://modelhub.org.cn", help=argparse.SUPPRESS) parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS) parser.add_argument("--hf-token", default=None, help=argparse.SUPPRESS) @@ -86,10 +85,11 @@ def _make_cycle_args(base_args: argparse.Namespace) -> argparse.Namespace: return cycle_args -def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient: - modelhub_tokens = list(getattr(base_args, "modelhub_tokens", []) or []) +def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool: + modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token])) if len(modelhub_tokens) > 1: - return MultiModelHubClient(tokens=modelhub_tokens, base_url=base_args.modelhub_base_url) + clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens] + return ModelHubClientPool(clients) return ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url) @@ -98,14 +98,13 @@ def run_poll_loop( base_args: argparse.Namespace, now=None, run_fn: Callable[..., dict[str, Any]] = run_daily_batches, - 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]: now = now or utc_now() - discovery_base_url = getattr(base_args, "modelscope_base_url", None) or base_args.hf_base_url - hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url) + hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url) modelhub_client = modelhub_client or _build_modelhub_client(base_args) template_selector = template_selector or TemplateSelector() @@ -251,6 +250,7 @@ def main(argv: list[str] | None = None) -> int: log( f"[poll] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} " + f"xc_token={'set' if bool(args.modelhub_token) else 'missing'} " f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}" ) summary = run_poll_loop(base_args=args) diff --git a/modelhub_submmit_api/runner_common.py b/modelhub_submmit_api/runner_common.py index ee3ab19..e63bf0b 100644 --- a/modelhub_submmit_api/runner_common.py +++ b/modelhub_submmit_api/runner_common.py @@ -8,13 +8,20 @@ from defaults import EMBEDDED_HF_TOKEN, EMBEDDED_MODELHUB_XC_TOKENS, EMBEDDED_MO DEFAULT_KEY_PATH = Path("KEY.md") +DEFAULT_KEYS_PATH = Path("KEYS.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 _token_sort_key(key: str) -> tuple[int, str]: + suffix = key.removeprefix("XC_TOKEN") + if not suffix: + return (0, key) + if suffix.isdigit(): + return (int(suffix), key) + return (10_000, key) + + def load_key_file(path: Path) -> dict[str, str]: if not path.exists(): return {} @@ -28,109 +35,108 @@ 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 load_key_files(*paths: Path) -> dict[str, str]: + loaded: dict[str, str] = {} + for path in paths: + if path.exists(): + loaded.update(load_key_file(path)) + return loaded -def split_token_list(value: str | None) -> list[str]: +def load_modelhub_tokens(values: dict[str, str]) -> list[str]: + tokens: list[tuple[str, str]] = [] + seen: set[str] = set() + for key, value in values.items(): + if key == "MODELHUB_XC_TOKEN" or key.startswith("XC_TOKEN"): + token = value.strip() + if token and token not in seen: + seen.add(token) + tokens.append((key, token)) + tokens.sort(key=lambda item: _token_sort_key(item[0])) + return [token for _, token in tokens] + + +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() + for raw in value.replace(",", "\n").replace(";", "\n").splitlines(): + token = raw.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 _add_token(tokens: list[str], token: str | None) -> None: + for item in _split_token_list(token): + if item and item not in tokens: + tokens.append(item) def ensure_tokens(args: argparse.Namespace) -> None: primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH)) + supplemental_key_path = primary_key_path.with_name(DEFAULT_KEYS_PATH.name) if not primary_key_path.exists(): - fallback_paths = [ + for candidate in ( 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 + ): + if candidate.exists(): + primary_key_path = candidate + supplemental_key_path = candidate.with_name(DEFAULT_KEYS_PATH.name) 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) + values = load_key_files(primary_key_path, supplemental_key_path) 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 + args.hf_token = values.get("HF_TOKEN") or 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 + args.modelscope_token = ( + os.getenv("MODELSCOPE_API_TOKEN") + or os.getenv("MODELSCOPE_TOKEN") + or values.get("MODELSCOPE_API_TOKEN") + or values.get("MODELSCOPE_TOKEN") + or EMBEDDED_MODELSCOPE_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) + tokens = load_modelhub_tokens(values) + env_token_items: list[tuple[str, str]] = [] + for key, value in os.environ.items(): + if key == "MODELHUB_XC_TOKEN" or key.startswith("XC_TOKEN"): + token = value.strip() + if token: + env_token_items.append((key, token)) + env_token_items.sort(key=lambda item: _token_sort_key(item[0])) + env_tokens = [token for _, token in env_token_items] - args.modelhub_token = modelhub_token - args.modelhub_tokens = modelhub_tokens + modelhub_token = getattr(args, "modelhub_token", None) + if modelhub_token and modelhub_token not in env_tokens: + env_tokens.insert(0, modelhub_token) + if not env_tokens: + env_tokens = tokens + else: + for token in tokens: + if token not in env_tokens: + env_tokens.append(token) + for env_name in ("MODELHUB_XC_TOKENS", "XC_TOKENS", "MODELHUB_TOKENS"): + for token in _split_token_list(os.getenv(env_name)): + _add_token(env_tokens, token) + for token in EMBEDDED_MODELHUB_XC_TOKENS: + _add_token(env_tokens, token) + + args.modelhub_tokens = env_tokens + args.modelhub_token = env_tokens[0] if env_tokens else None 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 + for index, token in enumerate(args.modelhub_tokens, start=1): + env_name = "XC_TOKEN" if index == 1 else f"XC_TOKEN{index}" + os.environ[env_name] = 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: - raise ValueError("MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN is required") + if not args.modelhub_token: + raise ValueError("XC_TOKEN is required, either via environment or KEY.md") diff --git a/modelhub_submmit_api/success_history_runner.py b/modelhub_submmit_api/success_history_runner.py index 7f41ffd..6dcd15d 100644 --- a/modelhub_submmit_api/success_history_runner.py +++ b/modelhub_submmit_api/success_history_runner.py @@ -9,7 +9,7 @@ from typing import Any, Callable from common import parse_datetime, utc_now, write_json, write_jsonl from daily_runner import DEFAULT_DAILY_RUNS_DIR, log from history_stats import is_success, update_history_archive -from modelhub_client import ModelHubClient +from modelhub_client import ModelHubClient, ModelHubClientPool from models import HFModelSummary from poll_runner import DEFAULT_POLL_RUNS_DIR, run_poll_loop from runner_common import DEFAULT_KEY_PATH, ensure_tokens @@ -63,7 +63,7 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Accepted for backward compatibility; this runner always restricts to historically successful GPUs.", ) - parser.add_argument("--min-downloads", type=int, default=50, help="Minimum Hugging Face download threshold") + parser.add_argument("--min-downloads", type=int, default=50, help="Minimum ModelScope download threshold") parser.add_argument( "--history-stats-threshold", type=int, @@ -72,7 +72,7 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--read-concurrency", type=int, default=4, help="Concurrency for read-only remote calls") parser.add_argument("--dry-run", action="store_true", help="Plan the queue without creating tasks") - parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing HF_TOKEN/XC_TOKEN") + parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN") parser.add_argument("--runs-dir", default="runs", help=argparse.SUPPRESS) parser.add_argument("--ledger-path", default="ledger/submissions.jsonl", help=argparse.SUPPRESS) parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS) @@ -80,10 +80,11 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--poll-runs-dir", default=str(DEFAULT_POLL_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--outcomes-path", default="outcomes/submissions.jsonl", help=argparse.SUPPRESS) - parser.add_argument("--hf-base-url", default=os.getenv("HF_BASE_URL", "https://huggingface.co"), help=argparse.SUPPRESS) + parser.add_argument("--hf-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) parser.add_argument("--hf-token", default=os.getenv("HF_TOKEN"), help=argparse.SUPPRESS) + parser.add_argument("--modelscope-token", default=os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN"), help=argparse.SUPPRESS) parser.add_argument("--poll-interval-seconds", type=int, default=60, help="Sleep between polling cycles when no slots are available") parser.add_argument("--idle-interval-seconds", type=int, default=30, help="Sleep between cycles when a scan submits nothing") parser.add_argument( @@ -96,7 +97,10 @@ def build_parser() -> argparse.ArgumentParser: return parser -def make_modelhub_client(args: argparse.Namespace) -> ModelHubClient: +def make_modelhub_client(args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool: + tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token])) + if len(tokens) > 1: + return ModelHubClientPool([ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in tokens]) return ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url) @@ -115,6 +119,9 @@ def model_id_from_task(task: dict[str, Any]) -> str | None: if value: return str(value).strip() address = str(task.get("modelAddress") or "").strip().rstrip("/") + marker = "modelscope.cn/models/" + if marker in address: + return address.split(marker, 1)[1].strip("/") marker = "huggingface.co/" if marker in address: return address.split(marker, 1)[1].strip("/") @@ -243,7 +250,7 @@ def run_success_history_poll( *, now=None, run_fn: Callable[..., dict[str, Any]] = run_poll_loop, - modelhub_client: ModelHubClient | None = None, + modelhub_client: ModelHubClient | ModelHubClientPool | None = None, template_selector: TemplateSelector | None = None, ) -> dict[str, Any]: now = now or utc_now() @@ -315,8 +322,9 @@ def main(argv: list[str] | None = None) -> int: args = parser.parse_args(argv) ensure_tokens(args) log( - f"[success-history] hf_token={'set' if bool(args.hf_token) else 'missing'} " - f"xc_token={'set' if bool(args.modelhub_token) else 'missing'}" + f"[success-history] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} " + f"xc_token={'set' if bool(args.modelhub_token) else 'missing'} " + f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}" ) summary = run_success_history_poll(args) print(f"poll_run_dir={summary['pollRunDir']}")