1 Commits

Author SHA1 Message Date
CoolBoy
c06169d906 Rebuild agent on original pooled runner 2026-07-10 02:14:08 +08:00
12 changed files with 459 additions and 535 deletions

View File

@@ -15,11 +15,11 @@ submission poller in a child process.
## Runtime Environment ## 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 agent platform. Environment variables can override them without rebuilding
the image. 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 - `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 - `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 - `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 "我的适配智能体". Create a tag and submit the repository URL plus tag in "我的适配智能体".
```bash ```bash
git tag agent-v4 git tag agent-v6
git push origin agent-v4 git push origin agent-v6
``` ```

View File

@@ -6,7 +6,7 @@ It currently supports:
- one-shot submission planning via `main.py` - one-shot submission planning via `main.py`
- daily batch execution via `run_daily.sh` - daily batch execution via `run_daily.sh`
- continuous queue refill via `run_poll.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 - automatic task/framework/template selection across the supported GPU catalog
## Layout ## Layout
@@ -15,8 +15,8 @@ It currently supports:
- `daily_runner.py`: daily wave orchestration - `daily_runner.py`: daily wave orchestration
- `poll_runner.py`: long-running queue refiller - `poll_runner.py`: long-running queue refiller
- `runner_common.py`: shared token / key file loading - `runner_common.py`: shared token / key file loading
- `modelscope_discovery.py`: ModelScope model discovery and inspection - `hf_discovery.py`: ModelScope model discovery and inspection (keeps the legacy module name)
- `modelhub_client.py`: ModelHub API client - `modelhub_client.py`: ModelHub API client and token-pool routing
- `history_stats.py`: online history aggregation, ranking, and warnings - `history_stats.py`: online history aggregation, ranking, and warnings
- `template_selector.py`: template lookup and GPU normalization - `template_selector.py`: template lookup and GPU normalization
- `task_registry.py`: task-type and framework selection rules - `task_registry.py`: task-type and framework selection rules
@@ -24,11 +24,12 @@ It currently supports:
## Key Files ## 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 - `templates/public_submit/adapt_task_templates.jsonl`: public submit templates
The agent platform deployment should use environment variables instead of key The runner reads both files automatically. Add more accounts by appending
files. Set `MODELHUB_XC_TOKEN` or `XC_TOKEN` for the single submitting account. `XC_TOKEN3`, `XC_TOKEN4`, and so on to `KEYS.md`.
Template lookup is also relative. The selector searches from the current working Template lookup is also relative. The selector searches from the current working
directory and the module directory. The primary project layout is: 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. - The runner auto-discovers all safe GPU/template combinations from the public submit catalog.
- Each model can be submitted at most once per GPU. - 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. - 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. - The default history threshold is `500` records.
@@ -83,7 +84,7 @@ Common flags:
- `--daily-target`: total target submissions for the day; `0` means unlimited - `--daily-target`: total target submissions for the day; `0` means unlimited
- `--min-downloads`: ModelScope download floor - `--min-downloads`: ModelScope download floor
- `--history-stats-threshold`: local ledger threshold before using online history stats - `--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 - `--scan-multiplier`: multiplier used for auto scan cap derivation from quota/queue capacity
- `--read-concurrency`: concurrent HTTP reads while scanning model candidates (default 4) - `--read-concurrency`: concurrent HTTP reads while scanning model candidates (default 4)
- `--max-submits-per-run`: max tasks to submit per run cycle (0 = unlimited) - `--max-submits-per-run`: max tasks to submit per run cycle (0 = unlimited)
@@ -95,9 +96,9 @@ Common flags:
`run_poll.sh` adds: `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 - `--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 - `--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) - `--max-submits-per-run`: max tasks to submit per poll cycle (0 = unlimited)
- `--skip-outcome-sync`: skip outcome sync before scanning - `--skip-outcome-sync`: skip outcome sync before scanning

View File

@@ -7,9 +7,9 @@ from pathlib import Path
from typing import Any, Callable from typing import Any, Callable
from common import utc_now, write_json from common import utc_now, write_json
from hf_discovery import HuggingFaceDiscovery
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission
from modelhub_client import ModelHubClient, MultiModelHubClient from modelhub_client import ModelHubClient, ModelHubClientPool
from modelscope_discovery import ModelScopeDiscovery
from outcome_tracker import OutcomeTracker from outcome_tracker import OutcomeTracker
from runner_common import DEFAULT_KEY_PATH, ensure_tokens from runner_common import DEFAULT_KEY_PATH, ensure_tokens
from template_selector import TemplateSelector 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("--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("--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("--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-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("--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("--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_path=base_args.history_archive_path,
history_archive_limit=base_args.history_archive_limit, history_archive_limit=base_args.history_archive_limit,
hf_base_url=base_args.hf_base_url, 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_base_url=base_args.modelhub_base_url,
modelhub_token=base_args.modelhub_token, modelhub_token=base_args.modelhub_token,
) )
@@ -133,18 +131,18 @@ def run_daily_batches(
waves: tuple[WaveSpec, ...] = DEFAULT_WAVES, waves: tuple[WaveSpec, ...] = DEFAULT_WAVES,
now=None, now=None,
run_fn: Callable[..., dict[str, Any]] = run_submission, run_fn: Callable[..., dict[str, Any]] = run_submission,
hf_discovery: ModelScopeDiscovery | None = None, hf_discovery: HuggingFaceDiscovery | None = None,
modelhub_client: ModelHubClient | None = None, modelhub_client: ModelHubClient | None = None,
template_selector: TemplateSelector | None = None, template_selector: TemplateSelector | None = None,
outcome_tracker: OutcomeTracker | None = None, outcome_tracker: OutcomeTracker | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
now = now or utc_now() 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 HuggingFaceDiscovery(base_url=base_args.hf_base_url)
hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
if modelhub_client is None: 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: 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: else:
modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url) modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
template_selector = template_selector or TemplateSelector() template_selector = template_selector or TemplateSelector()
@@ -293,6 +291,7 @@ def main(argv: list[str] | None = None) -> int:
ensure_tokens(args) ensure_tokens(args)
log( log(
f"[daily] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} " 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 [])}" f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
) )
summary = run_daily_batches(base_args=args) summary = run_daily_batches(base_args=args)

View File

@@ -1,12 +1,14 @@
from __future__ import annotations from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor, as_completed
import os import os
import threading import threading
from pathlib import PurePosixPath from pathlib import PurePosixPath
from typing import Any from typing import Any
from urllib.parse import quote
from common import parse_datetime from common import parse_datetime
from defaults import EMBEDDED_MODELSCOPE_TOKEN
from http_json import HttpJsonError
from http_json import JsonHttpClient from http_json import JsonHttpClient
from models import HFModelSummary, ModelInspection 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") VLLM_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth")
ONNX_WEIGHT_SUFFIXES = (".onnx",) 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: class HuggingFaceDiscovery:
def __init__( def __init__(
self, self,
base_url: str = "https://huggingface.co", base_url: str = "https://modelscope.cn",
http_client: JsonHttpClient | None = None, http_client: JsonHttpClient | None = None,
legacy_http_client: JsonHttpClient | None = None,
timeout: int = 30, timeout: int = 30,
retries: int = 2, retries: int = 2,
) -> None: ) -> 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"} headers = {"User-Agent": "modelhub-submmit-cli/0.1"}
if hf_token: if token:
headers["Authorization"] = f"Bearer {hf_token}" headers["Authorization"] = f"Bearer {token}"
headers["Cookie"] = f"m_session_id={token}"
self.http_client = http_client or JsonHttpClient( 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, base_url=base_url,
default_headers=headers, default_headers=headers,
timeout=timeout, timeout=timeout,
@@ -50,36 +79,15 @@ class HuggingFaceDiscovery:
return [] return []
deduped: dict[str, HFModelSummary] = {} deduped: dict[str, HFModelSummary] = {}
max_workers = min(len(pipeline_tags), max(1, read_concurrency)) del read_concurrency
if max_workers <= 1: for pipeline_tag in pipeline_tags:
tag_results = [ for model in self._query_recent_models(
self._query_recent_models( pipeline_tag=pipeline_tag,
pipeline_tag=pipeline_tag, limit=limit,
limit=limit, min_downloads=min_downloads,
min_downloads=min_downloads, updated_after=updated_after,
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:
current = deduped.get(model.repo_id) current = deduped.get(model.repo_id)
if current is None or (model.last_modified or parse_datetime("1970-01-01")) > ( if current is None or (model.last_modified or parse_datetime("1970-01-01")) > (
current.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, min_downloads: int,
updated_after=None, updated_after=None,
) -> list[HFModelSummary]: ) -> list[HFModelSummary]:
payload = self.http_client.request_json( page_size = min(max(1, limit), 100)
"GET", max_items = min(max(1, limit), 3000)
"/api/models", task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag)
query={
"pipeline_tag": pipeline_tag,
"sort": "lastModified",
"direction": "-1",
"limit": limit,
"full": "true",
},
)
models: list[HFModelSummary] = [] models: list[HFModelSummary] = []
for item in payload or []: for page_number in range(1, (max_items + page_size - 1) // page_size + 1):
model = self._parse_model(item, min_downloads=min_downloads, updated_after=updated_after) try:
if model is not None: payload = self.http_client.request_json(
models.append(model) "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 return models
def list_recent_text_generation_models( def list_recent_text_generation_models(
@@ -140,40 +167,67 @@ class HuggingFaceDiscovery:
if cached is not None: if cached is not None:
return list(cached) return list(cached)
payload = self.http_client.request_json( encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
"GET", try:
f"/api/models/{repo_id}/tree/main", payload = self.legacy_http_client.request_json(
query={"recursive": "1"}, "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]] entries = self._extract_files(payload)
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 = []
with self._repo_tree_lock: with self._repo_tree_lock:
self._repo_tree_cache[repo_id] = list(entries) self._repo_tree_cache[repo_id] = list(entries)
return list(entries) return list(entries)
@staticmethod @staticmethod
def _parse_model(item: dict[str, Any], *, min_downloads: int, updated_after=None) -> HFModelSummary | None: def _extract_models(payload: Any) -> list[dict[str, Any]]:
repo_id = item.get("id") or item.get("modelId") 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: if not repo_id:
return None return None
downloads = int(item.get("downloads") or 0) downloads = int(item.get("downloads") or item.get("Downloads") or 0)
if downloads < min_downloads: if downloads < min_downloads:
return None return None
pipeline_tag = item.get("pipeline_tag") or item.get("pipelineTag") pipeline_tag = fallback_pipeline_tag
last_modified = parse_datetime(item.get("lastModified") or item.get("last_modified")) 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: if updated_after and last_modified and last_modified < updated_after:
return None return None
return HFModelSummary( return HFModelSummary(
@@ -181,7 +235,7 @@ class HuggingFaceDiscovery:
downloads=downloads, downloads=downloads,
last_modified=last_modified, last_modified=last_modified,
pipeline_tag=pipeline_tag, 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")),
) )

View File

@@ -8,6 +8,7 @@ from pathlib import Path
from typing import Any from typing import Any
from common import parse_datetime, utc_now, write_json, write_jsonl from common import parse_datetime, utc_now, write_json, write_jsonl
from hf_discovery import HuggingFaceDiscovery
from history_stats import ( from history_stats import (
append_ledger_entry, append_ledger_entry,
build_empty_pre_submit_report, build_empty_pre_submit_report,
@@ -15,9 +16,8 @@ from history_stats import (
load_ledger, load_ledger,
update_history_archive, update_history_archive,
) )
from modelhub_client import ModelHubAPIError, ModelHubClient, MultiModelHubClient from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool
from models import CandidateModel, HFModelSummary, ModelInspection from models import CandidateModel, HFModelSummary, ModelInspection
from modelscope_discovery import ModelScopeDiscovery
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker 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 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 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-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("--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("--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-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("--modelhub-token", default=os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN"), help=argparse.SUPPRESS)
return parser return parser
@@ -195,7 +194,7 @@ def choose_candidate_for_gpu(
def resolve_submit_concurrency( def resolve_submit_concurrency(
args: argparse.Namespace, args: argparse.Namespace,
*, *,
modelhub_client: ModelHubClient | Any, modelhub_client: ModelHubClient | ModelHubClientPool | Any,
planned_submit_count: int, planned_submit_count: int,
) -> int: ) -> int:
explicit_workers = int(getattr(args, "submit_concurrency", 0) or 0) explicit_workers = int(getattr(args, "submit_concurrency", 0) or 0)
@@ -228,7 +227,7 @@ def resolve_max_submit_count(
def process_model_for_candidates( def process_model_for_candidates(
*, *,
model: HFModelSummary, model: HFModelSummary,
hf_discovery: ModelScopeDiscovery, hf_discovery: HuggingFaceDiscovery,
modelhub_client: ModelHubClient, modelhub_client: ModelHubClient,
template_selector: TemplateSelector, template_selector: TemplateSelector,
target_gpus: list[str], target_gpus: list[str],
@@ -316,6 +315,11 @@ def submit_candidate(
"responseData": response.get("data"), "responseData": response.get("data"),
} }
except ModelHubAPIError as exc: except ModelHubAPIError as exc:
print(
f"[submit] failed repo={candidate['repoId']} gpu={candidate['targetGpu']} "
f"framework={candidate['framework']} reason={exc}",
flush=True,
)
return { return {
"outcome": "failed", "outcome": "failed",
"candidate": candidate, "candidate": candidate,
@@ -340,8 +344,8 @@ def run_submission(
args: argparse.Namespace, args: argparse.Namespace,
*, *,
now=None, now=None,
hf_discovery: ModelScopeDiscovery | None = None, hf_discovery: HuggingFaceDiscovery | None = None,
modelhub_client: ModelHubClient | None = None, modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
template_selector: TemplateSelector | None = None, template_selector: TemplateSelector | None = None,
outcome_tracker: OutcomeTracker | None = None, outcome_tracker: OutcomeTracker | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
@@ -352,12 +356,12 @@ def run_submission(
if not target_gpus: if not target_gpus:
raise RuntimeError("No auto-submittable GPUs are available for the selected task types") 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 HuggingFaceDiscovery(base_url=args.hf_base_url)
hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
if modelhub_client is None: 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: 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: else:
modelhub_client = ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url) modelhub_client = ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url)

View File

@@ -1,9 +1,10 @@
from __future__ import annotations from __future__ import annotations
import os import os
import time
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta from datetime import datetime, timedelta
from itertools import cycle
from threading import Lock
from typing import Any from typing import Any
from common import format_modelhub_datetime, parse_datetime from common import format_modelhub_datetime, parse_datetime
@@ -28,21 +29,10 @@ class ModelHubClient:
retries: int = 2, retries: int = 2,
http_client: JsonHttpClient | None = None, http_client: JsonHttpClient | None = None,
) -> None: ) -> None:
self.token = ( self.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or EMBEDDED_MODELHUB_XC_TOKEN
token if not self.token and http_client is None:
or os.getenv("MODELHUB_XC_TOKEN") raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN or XC_TOKEN.")
or os.getenv("XC_TOKEN") default_headers = {"Xc-Token": self.token} if self.token else {}
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( self.http_client = http_client or JsonHttpClient(
base_url=base_url, base_url=base_url,
default_headers=default_headers, default_headers=default_headers,
@@ -207,73 +197,6 @@ class ModelHubClient:
return payload 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 = { ACTIVE_TASK_STATUSES = {
"waiting", "waiting",
"running", "running",
@@ -308,3 +231,175 @@ def is_active_task(task: dict[str, Any]) -> bool:
if status in TERMINAL_TASK_STATUSES: if status in TERMINAL_TASK_STATUSES:
return False return False
return True 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

View File

@@ -23,13 +23,10 @@ class HFModelSummary:
last_modified: datetime | None last_modified: datetime | None
pipeline_tag: str | None pipeline_tag: str | None
created_at: datetime | None = None created_at: datetime | None = None
source: str = "huggingface"
@property @property
def model_address(self) -> str: def model_address(self) -> str:
if self.source == "modelscope": return f"https://modelscope.cn/models/{self.repo_id}"
return f"https://modelscope.cn/models/{self.repo_id}"
return f"https://huggingface.co/{self.repo_id}"
@dataclass(frozen=True) @dataclass(frozen=True)

View File

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

View File

@@ -7,7 +7,7 @@ from typing import Any
from common import append_jsonl, parse_datetime, read_jsonl, utc_now, write_jsonl from common import append_jsonl, parse_datetime, read_jsonl, utc_now, write_jsonl
from history_stats import classify_failure, is_failure, is_success 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") DEFAULT_OUTCOMES_PATH = Path("outcomes/submissions.jsonl")
@@ -72,11 +72,14 @@ class OutcomeTracker:
self._by_model_gpu[(model_id, target_gpu)].append(record) self._by_model_gpu[(model_id, target_gpu)].append(record)
append_jsonl(self.path, record) append_jsonl(self.path, record)
def sync_from_api(self, client: ModelHubClient) -> int: def sync_from_api(self, client: ModelHubClient | ModelHubClientPool) -> int:
try: try:
begin = self._last_sync_time begin = self._last_sync_time
end = utc_now() 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} 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) tasks = client.list_tasks(**list_kwargs)
except Exception: except Exception:
return 0 return 0

View File

@@ -9,9 +9,9 @@ from typing import Any, Callable
from common import utc_now, write_json from common import utc_now, write_json
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches 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 main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR
from modelhub_client import ModelHubClient, MultiModelHubClient from modelhub_client import ModelHubClient, ModelHubClientPool
from modelscope_discovery import ModelScopeDiscovery
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
from runner_common import DEFAULT_KEY_PATH, ensure_tokens from runner_common import DEFAULT_KEY_PATH, ensure_tokens
from template_selector import TemplateSelector 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("--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("--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("--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-base-url", default="https://modelhub.org.cn", help=argparse.SUPPRESS)
parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS) parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS)
parser.add_argument("--hf-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 return cycle_args
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient: def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool:
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: 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) 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, base_args: argparse.Namespace,
now=None, now=None,
run_fn: Callable[..., dict[str, Any]] = run_daily_batches, run_fn: Callable[..., dict[str, Any]] = run_daily_batches,
hf_discovery: ModelScopeDiscovery | None = None, hf_discovery: HuggingFaceDiscovery | None = None,
modelhub_client: ModelHubClient | None = None, modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
template_selector: TemplateSelector | None = None, template_selector: TemplateSelector | None = None,
outcome_tracker: OutcomeTracker | None = None, outcome_tracker: OutcomeTracker | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
now = now or utc_now() 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 HuggingFaceDiscovery(base_url=base_args.hf_base_url)
hf_discovery = hf_discovery or ModelScopeDiscovery(base_url=discovery_base_url)
modelhub_client = modelhub_client or _build_modelhub_client(base_args) modelhub_client = modelhub_client or _build_modelhub_client(base_args)
template_selector = template_selector or TemplateSelector() template_selector = template_selector or TemplateSelector()
@@ -251,6 +250,7 @@ def main(argv: list[str] | None = None) -> int:
log( log(
f"[poll] modelscope_token={'set' if bool(args.modelscope_token) else 'missing'} " 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 [])}" f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
) )
summary = run_poll_loop(base_args=args) summary = run_poll_loop(base_args=args)

View File

@@ -8,13 +8,20 @@ from defaults import EMBEDDED_HF_TOKEN, EMBEDDED_MODELHUB_XC_TOKENS, EMBEDDED_MO
DEFAULT_KEY_PATH = Path("KEY.md") DEFAULT_KEY_PATH = Path("KEY.md")
DEFAULT_KEYS_PATH = Path("KEYS.md")
MODULE_DIR = Path(__file__).resolve().parent 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") 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]: def load_key_file(path: Path) -> dict[str, str]:
if not path.exists(): if not path.exists():
return {} return {}
@@ -28,109 +35,108 @@ def load_key_file(path: Path) -> dict[str, str]:
return loaded return loaded
def first_value(values: dict[str, str], names: tuple[str, ...]) -> str | None: def load_key_files(*paths: Path) -> dict[str, str]:
for name in names: loaded: dict[str, str] = {}
value = os.getenv(name) or values.get(name) for path in paths:
if value: if path.exists():
return value.strip() loaded.update(load_key_file(path))
return None 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: if not value:
return [] return []
tokens: list[str] = [] tokens: list[str] = []
for normalized in value.replace(",", "\n").replace(";", "\n").splitlines(): for raw in value.replace(",", "\n").replace(";", "\n").splitlines():
token = normalized.strip() token = raw.strip()
if token: if token:
tokens.append(token) tokens.append(token)
return tokens return tokens
def discover_key_paths(primary_path: Path) -> list[Path]: def _add_token(tokens: list[str], token: str | None) -> None:
candidates: list[Path] = [] for item in _split_token_list(token):
search_dirs = [primary_path.parent, primary_path.parent.parent] if item and item not in tokens:
for path in [primary_path, primary_path.parent / "KEYS.md"]: tokens.append(item)
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: def ensure_tokens(args: argparse.Namespace) -> None:
primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH)) 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(): if not primary_key_path.exists():
fallback_paths = [ for candidate in (
MODULE_DIR / primary_key_path.name, MODULE_DIR / primary_key_path.name,
primary_key_path.parent.parent / primary_key_path.name, primary_key_path.parent.parent / primary_key_path.name,
MODULE_DIR.parent / primary_key_path.name, MODULE_DIR.parent / primary_key_path.name,
] ):
for fallback_path in fallback_paths: if candidate.exists():
if fallback_path.exists(): primary_key_path = candidate
primary_key_path = fallback_path supplemental_key_path = candidate.with_name(DEFAULT_KEYS_PATH.name)
break break
values_by_path = [load_key_file(path) for path in discover_key_paths(primary_key_path) if path.exists()] values = load_key_files(primary_key_path, supplemental_key_path)
values: dict[str, str] = {}
for loaded in values_by_path:
values.update(loaded)
if not getattr(args, "hf_token", None): if not getattr(args, "hf_token", None):
args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN") args.hf_token = values.get("HF_TOKEN") or EMBEDDED_HF_TOKEN
if not getattr(args, "hf_token", None):
args.hf_token = EMBEDDED_HF_TOKEN
if not getattr(args, "modelscope_token", None): if not getattr(args, "modelscope_token", None):
args.modelscope_token = first_value(values, MODELSCOPE_TOKEN_ENV_NAMES) args.modelscope_token = (
if not getattr(args, "modelscope_token", None): os.getenv("MODELSCOPE_API_TOKEN")
args.modelscope_token = EMBEDDED_MODELSCOPE_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)) tokens = load_modelhub_tokens(values)
modelhub_token = modelhub_tokens[0] if modelhub_tokens else None env_token_items: list[tuple[str, str]] = []
jwt_token = first_value(values, JWT_TOKEN_ENV_NAMES) 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 modelhub_token = getattr(args, "modelhub_token", None)
args.modelhub_tokens = modelhub_tokens 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: if args.hf_token:
os.environ["HF_TOKEN"] = args.hf_token os.environ["HF_TOKEN"] = args.hf_token
if args.modelscope_token: if args.modelscope_token:
os.environ["MODELSCOPE_API_TOKEN"] = args.modelscope_token os.environ["MODELSCOPE_API_TOKEN"] = args.modelscope_token
os.environ["MODELSCOPE_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: if args.modelhub_token:
os.environ["MODELHUB_XC_TOKEN"] = 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) os.environ["MODELHUB_XC_TOKENS"] = ",".join(args.modelhub_tokens)
if jwt_token: if not args.modelhub_token:
os.environ["MODELHUB_JWT_TOKEN"] = jwt_token raise ValueError("XC_TOKEN is required, either via environment or KEY.md")
if not args.modelhub_token and not jwt_token:
raise ValueError("MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN is required")

View File

@@ -9,7 +9,7 @@ from typing import Any, Callable
from common import parse_datetime, utc_now, write_json, write_jsonl from common import parse_datetime, utc_now, write_json, write_jsonl
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log from daily_runner import DEFAULT_DAILY_RUNS_DIR, log
from history_stats import is_success, update_history_archive 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 models import HFModelSummary
from poll_runner import DEFAULT_POLL_RUNS_DIR, run_poll_loop from poll_runner import DEFAULT_POLL_RUNS_DIR, run_poll_loop
from runner_common import DEFAULT_KEY_PATH, ensure_tokens from runner_common import DEFAULT_KEY_PATH, ensure_tokens
@@ -63,7 +63,7 @@ def build_parser() -> argparse.ArgumentParser:
action="store_true", action="store_true",
help="Accepted for backward compatibility; this runner always restricts to historically successful GPUs.", 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( parser.add_argument(
"--history-stats-threshold", "--history-stats-threshold",
type=int, 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("--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("--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("--runs-dir", default="runs", help=argparse.SUPPRESS)
parser.add_argument("--ledger-path", default="ledger/submissions.jsonl", 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) 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("--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("--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("--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-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("--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("--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("--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("--idle-interval-seconds", type=int, default=30, help="Sleep between cycles when a scan submits nothing")
parser.add_argument( parser.add_argument(
@@ -96,7 +97,10 @@ def build_parser() -> argparse.ArgumentParser:
return parser 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) 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: if value:
return str(value).strip() return str(value).strip()
address = str(task.get("modelAddress") or "").strip().rstrip("/") 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/" marker = "huggingface.co/"
if marker in address: if marker in address:
return address.split(marker, 1)[1].strip("/") return address.split(marker, 1)[1].strip("/")
@@ -243,7 +250,7 @@ def run_success_history_poll(
*, *,
now=None, now=None,
run_fn: Callable[..., dict[str, Any]] = run_poll_loop, 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, template_selector: TemplateSelector | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
now = now or utc_now() now = now or utc_now()
@@ -315,8 +322,9 @@ def main(argv: list[str] | None = None) -> int:
args = parser.parse_args(argv) args = parser.parse_args(argv)
ensure_tokens(args) ensure_tokens(args)
log( log(
f"[success-history] hf_token={'set' if bool(args.hf_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_token={'set' if bool(args.modelhub_token) else 'missing'} "
f"xc_tokens={len(getattr(args, 'modelhub_tokens', []) or [])}"
) )
summary = run_success_history_poll(args) summary = run_success_history_poll(args)
print(f"poll_run_dir={summary['pollRunDir']}") print(f"poll_run_dir={summary['pollRunDir']}")