diff --git a/README.md b/README.md index 049337a4..cdab283b 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ Optional tuning: - `MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS` default `2` - `MODELHUB_AGENT_MAX_SUBMITS_PER_RUN` default `0` (fill all currently available slots) - `MODELHUB_AGENT_ACTIVE_TASK_CAP` default `100` per account +- `MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES` default `3` +- `MODELHUB_CAPACITY_STATE_PATH` default `.modelhub_state/account_capacity.json` - `MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS` default `15` - `MODELHUB_AGENT_RESERVATION_TTL_SECONDS` default `120` - `MODELHUB_AGENT_INSTANCE_ID` optional stable worker identity used to spread concurrent agents across accounts and candidates @@ -39,6 +41,27 @@ Optional tuning: - `MODELHUB_AGENT_MIN_DOWNLOADS` - `MODELHUB_AGENT_GPUS` - `MODELHUB_AGENT_EXTRA_ARGS` +- `MODELHUB_GPU_STRATEGY_STATE_PATH` default `.modelhub_state/gpu_strategy.json` +- `MODELSCOPE_PAGE_INTERVAL_SECONDS` default `0.25` +- `MODELSCOPE_PAGE_CACHE_TTL_SECONDS` default `900` + +## Adaptive GPU Strategy + +When no explicit GPU override is supplied, the worker uses a local 50/30/20 +strategy generation: + +- 50%: the three long-term GPUs with the best Wilson lower confidence score and at least 100 terminal samples +- 30%: round-robin exploration across every currently supported GPU +- 20%: the best GPU among the latest 1,000 terminal tasks + +Only platform-accepted tasks count. After exactly 200 accepted tasks, the next +poll cycle reloads all account history, generates a new immutable strategy snapshot, +and resets the generation counters to 100/60/40 targets. The active snapshot and +progress are stored in `.modelhub_state/gpu_strategy.json`. + +ModelScope HTTP 429 responses use exponential backoff and `Retry-After`. Successful +pages remain cached, so a later cycle retries the failed page instead of restarting +the whole pagination scan. ## Concurrent Agents @@ -56,7 +79,7 @@ authority for account capacity and model/GPU uniqueness. If the platform reports that a model/GPU is already being validated, the claim is retained and the runner immediately draws replacement candidates from the same scan instead of retrying the duplicate every cycle. Startup logs and the -health response expose `agent_version`; version `2026.08.02.2` or newer includes +health response expose `agent_version`; version `2026.08.02.3` or newer includes this behavior. ## Deploy diff --git a/modelhub_submmit_api/README.md b/modelhub_submmit_api/README.md index cc6bc3a7..37126c16 100644 --- a/modelhub_submmit_api/README.md +++ b/modelhub_submmit_api/README.md @@ -8,6 +8,7 @@ It currently supports: - continuous queue refill via `run_poll.sh` - multiple ModelHub tokens read from `KEY.md` and `KEYS.md` - automatic task/framework/template selection across the supported GPU catalog +- adaptive long-term/exploration/recent GPU scheduling with a persistent local snapshot ## Layout @@ -72,13 +73,21 @@ bash run_poll.sh --dry-run ## Behavior - The runner auto-discovers all safe GPU/template combinations from the public submit catalog. +- Automatic GPU selection uses smooth 50/30/20 scheduling: long-term Wilson-ranked top 3 GPUs, + all supported GPUs, and the best GPU from the latest 1,000 terminal tasks. +- A strategy generation lasts exactly 200 platform-accepted submissions. Rejected API calls and + duplicates do not advance it. The next cycle refreshes platform history before submitting again. +- Strategy state is stored in `.modelhub_state/gpu_strategy.json`; a generation never recalculates + during candidate submission. - Each model can be submitted at most once per GPU. - Multiple ModelHub tokens are pooled and used to route submissions to the account with available async capacity. - Concurrent submissions reserve account slots locally, and an account-capacity race automatically falls through to another account. - Concurrent local processes claim model/GPU pairs in `.modelhub_state/submission_claims.jsonl`; shared ledger, history, and outcome files use process locks and atomic replacement. -- Isolated agent containers diversify candidate order by instance identity to reduce cross-container duplicate submissions. -- 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. +- ModelScope list pages are paced and cached for 15 minutes. HTTP 429 responses use exponential + backoff and `Retry-After`; pages already downloaded remain usable and the failed page is retried + on the next cycle. +- Every third poll cycle, a full account gets one controlled capacity probe. A successful probe + raises that account's persisted known limit; a capacity rejection enters cooldown. ## Important Flags @@ -95,6 +104,8 @@ Common flags: - `--skip-outcome-sync`: skip outcome sync before scanning - `--skip-history-archive`: skip history archive download for this run - `--dry-run`: plan only, do not submit +- `--gpu-strategy-refresh-submissions`: accepted tasks per strategy generation (default `200`) +- `--disable-gpu-strategy`: restore legacy ordering; explicit `--gpu/--gpus` also bypasses adaptive selection - `run_daily.sh` injects `--daily-target 3` when no daily-target flag is provided. Set `SUBMIT_DAILY_TARGET` or pass `--daily-target` explicitly for a different target. `run_poll.sh` adds: @@ -127,6 +138,11 @@ Each run typically includes: - `skipped.jsonl` - `failed.jsonl` +Persistent local scheduler state is written under `.modelhub_state/`: + +- `gpu_strategy.json`: GPU ranks, generation progress, and 50/30/20 accepted counters +- `account_capacity.json`: learned per-account active-task limits + ## Verification Run the full test suite: diff --git a/modelhub_submmit_api/daily_runner.py b/modelhub_submmit_api/daily_runner.py index 278b10f9..0f09564a 100644 --- a/modelhub_submmit_api/daily_runner.py +++ b/modelhub_submmit_api/daily_runner.py @@ -7,9 +7,10 @@ from pathlib import Path from typing import Any, Callable from common import utc_now, write_json +from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH from hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir, run_submission -from modelhub_client import ModelHubClient, ModelHubClientPool +from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool from outcome_tracker import OutcomeTracker from runner_common import DEFAULT_KEY_PATH, ensure_tokens from submission_claims import DEFAULT_CLAIMS_PATH @@ -79,6 +80,13 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--skip-outcome-sync", action="store_true", help="Skip outcome sync from ModelHub before scanning") parser.add_argument("--skip-history-archive", action="store_true", help="Skip historical task archive download for this run") parser.add_argument("--dry-run", action="store_true", help="Plan the day without creating tasks") + parser.add_argument("--disable-gpu-strategy", action="store_true", help="Disable adaptive 50/30/20 GPU scheduling") + parser.add_argument( + "--gpu-strategy-refresh-submissions", + type=int, + default=200, + help="Recalculate adaptive GPU choices after this many accepted submissions", + ) 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=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS) @@ -89,6 +97,24 @@ 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( + "--gpu-strategy-state-path", + default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument("--gpu-strategy-recent-window", type=int, default=1000, help=argparse.SUPPRESS) + parser.add_argument("--gpu-strategy-min-long-samples", type=int, default=100, help=argparse.SUPPRESS) + parser.add_argument( + "--capacity-state-path", + default=os.getenv("MODELHUB_CAPACITY_STATE_PATH", str(DEFAULT_CAPACITY_STATE_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--capacity-probe-interval-cycles", + type=int, + default=int(os.getenv("MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES", "3")), + 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("--modelhub-base-url", default=os.getenv("MODELHUB_BASE_URL", "https://modelhub.org.cn"), help=argparse.SUPPRESS) @@ -118,6 +144,13 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar scan_multiplier=getattr(base_args, "scan_multiplier", 4), skip_outcome_sync=getattr(base_args, "skip_outcome_sync", False), skip_history_archive=getattr(base_args, "skip_history_archive", False), + disable_gpu_strategy=getattr(base_args, "disable_gpu_strategy", False), + gpu_strategy_refresh_submissions=getattr(base_args, "gpu_strategy_refresh_submissions", 200), + gpu_strategy_state_path=getattr(base_args, "gpu_strategy_state_path", str(DEFAULT_GPU_STRATEGY_PATH)), + gpu_strategy_recent_window=getattr(base_args, "gpu_strategy_recent_window", 1000), + gpu_strategy_min_long_samples=getattr(base_args, "gpu_strategy_min_long_samples", 100), + capacity_state_path=getattr(base_args, "capacity_state_path", str(DEFAULT_CAPACITY_STATE_PATH)), + capacity_probe_interval_cycles=getattr(base_args, "capacity_probe_interval_cycles", 3), submit_concurrency=getattr(base_args, "submit_concurrency", 1), max_submits_per_run=getattr(base_args, "max_submits_per_run", 0), runs_dir=base_args.runs_dir, @@ -149,7 +182,11 @@ def run_daily_batches( modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token])) token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token] clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values] - modelhub_client = ModelHubClientPool(clients) + modelhub_client = ModelHubClientPool( + clients, + capacity_probe_interval_cycles=max(0, int(getattr(base_args, "capacity_probe_interval_cycles", 3) or 0)), + capacity_state_path=Path(getattr(base_args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)), + ) template_selector = template_selector or TemplateSelector() daily_run_dir = make_run_dir(Path(base_args.daily_runs_dir), now) @@ -172,6 +209,11 @@ def run_daily_batches( log(f"[daily] round={round_index} start") for wave in waves: attempted_waves += 1 + if ( + hasattr(modelhub_client, "configure_capacity_probe") + and not bool(getattr(base_args, "capacity_probe_already_configured", False)) + ): + modelhub_client.configure_capacity_probe(attempted_waves) wave_args = make_wave_namespace(base_args, wave) log( f"[daily] round={round_index} wave={wave.name} " diff --git a/modelhub_submmit_api/gpu_strategy.py b/modelhub_submmit_api/gpu_strategy.py new file mode 100644 index 00000000..6bd8463e --- /dev/null +++ b/modelhub_submmit_api/gpu_strategy.py @@ -0,0 +1,480 @@ +from __future__ import annotations + +import math +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime, timedelta +from pathlib import Path +from typing import Any, Callable + +from common import parse_datetime, read_json, utc_now, write_json +from history_stats import is_failure, is_success +from modelhub_client import ModelHubClientPool +from submission_claims import candidate_key + + +DEFAULT_GPU_STRATEGY_PATH = Path(".modelhub_state/gpu_strategy.json") +DEFAULT_REFRESH_SUBMISSIONS = 200 +DEFAULT_RECENT_TERMINAL_WINDOW = 1000 +DEFAULT_LONG_TERM_MIN_SAMPLES = 100 +STRATEGY_STATE_VERSION = 1 + +LONG_TERM = "long_term" +ALL_SUPPORTED = "all_supported" +RECENT = "recent" +CATEGORIES = (LONG_TERM, ALL_SUPPORTED, RECENT) +CATEGORY_WEIGHTS = {LONG_TERM: 5, ALL_SUPPORTED: 3, RECENT: 2} + +# The long-term half is split 50/30/20 across its three ranked GPUs. +LONG_TERM_GPU_PATTERN = (0, 1, 0, 2, 0, 1, 0, 1, 0, 2) + + +def _empty_category_counts() -> dict[str, int]: + return {category: 0 for category in CATEGORIES} + + +def _gpu_name(record: dict[str, Any]) -> str: + return str(record.get("gpuType") or record.get("targetGpu") or "").strip() + + +def _terminal_outcome(record: dict[str, Any]) -> str | None: + try: + if is_success(record): + return "success" + if is_failure(record): + return "failure" + except TypeError: + # Be tolerant of APIs that serialize verifyResult as a string. + copied = dict(record) + try: + copied["verifyResult"] = float(record.get("verifyResult")) + except (TypeError, ValueError): + copied["verifyResult"] = None + if is_success(copied): + return "success" + if is_failure(copied): + return "failure" + return None + + +def _task_sort_key(record: dict[str, Any]) -> tuple[float, str]: + timestamp = ( + parse_datetime(record.get("updateTime")) + or parse_datetime(record.get("createTime")) + or parse_datetime(record.get("submitTime")) + ) + return (timestamp.timestamp() if timestamp else 0.0, str(record.get("taskId") or "")) + + +def wilson_lower_bound(successes: int, total: int, *, z: float = 1.96) -> float: + if total <= 0: + return 0.0 + probability = successes / total + z_squared = z * z + denominator = 1.0 + z_squared / total + centre = probability + z_squared / (2.0 * total) + margin = z * math.sqrt( + (probability * (1.0 - probability) / total) + z_squared / (4.0 * total * total) + ) + return max(0.0, (centre - margin) / denominator) + + +def _summarize_gpu_records( + records: list[dict[str, Any]], + supported_gpus: list[str], +) -> dict[str, dict[str, Any]]: + supported = set(supported_gpus) + counts: dict[str, dict[str, int]] = { + gpu: {"success": 0, "failure": 0} + for gpu in supported_gpus + } + for record in records: + gpu = _gpu_name(record) + if gpu not in supported: + continue + outcome = _terminal_outcome(record) + if outcome is not None: + counts[gpu][outcome] += 1 + + summaries: dict[str, dict[str, Any]] = {} + for gpu in supported_gpus: + success = counts[gpu]["success"] + failure = counts[gpu]["failure"] + terminal = success + failure + summaries[gpu] = { + "gpu": gpu, + "success": success, + "failure": failure, + "terminal": terminal, + "successRate": success / terminal if terminal else 0.0, + "wilsonLowerBound": wilson_lower_bound(success, terminal), + } + return summaries + + +def _rank_gpu_summaries(summaries: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: + return sorted( + summaries.values(), + key=lambda item: ( + -float(item["wilsonLowerBound"]), + -float(item["successRate"]), + -int(item["terminal"]), + str(item["gpu"]), + ), + ) + + +def build_strategy_snapshot( + tasks: list[dict[str, Any]], + *, + supported_gpus: list[str], + generated_at: datetime | None = None, + generation: int = 0, + recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW, + long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES, + refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS, +) -> dict[str, Any]: + generated_at = generated_at or utc_now() + supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu)) + if not supported_gpus: + raise ValueError("At least one supported GPU is required") + + all_time = _summarize_gpu_records(tasks, supported_gpus) + all_time_ranked = _rank_gpu_summaries(all_time) + qualified_long = [ + item for item in all_time_ranked + if int(item["terminal"]) >= max(1, int(long_term_min_samples)) + ] + long_term_gpus = [str(item["gpu"]) for item in qualified_long[:3]] + if len(long_term_gpus) < min(3, len(supported_gpus)): + for item in all_time_ranked: + gpu = str(item["gpu"]) + if gpu not in long_term_gpus: + long_term_gpus.append(gpu) + if len(long_term_gpus) >= min(3, len(supported_gpus)): + break + + terminal_tasks = [record for record in tasks if _terminal_outcome(record) is not None] + terminal_tasks.sort(key=_task_sort_key, reverse=True) + recent_tasks = terminal_tasks[: max(1, int(recent_terminal_window))] + recent = _summarize_gpu_records(recent_tasks, supported_gpus) + recent_ranked = _rank_gpu_summaries(recent) + recent_qualified = [item for item in recent_ranked if int(item["terminal"]) >= 20] + if recent_qualified: + recent_gpu = str(recent_qualified[0]["gpu"]) + elif recent_tasks and recent_ranked: + recent_gpu = str(recent_ranked[0]["gpu"]) + else: + recent_gpu = long_term_gpus[0] + + return { + "version": STRATEGY_STATE_VERSION, + "generation": max(0, int(generation)), + "generatedAt": generated_at.isoformat(), + "historyReady": True, + "historyTaskCount": len(tasks), + "supportedGpus": supported_gpus, + "longTermGpus": long_term_gpus, + "recentGpu": recent_gpu, + "recentTerminalCount": len(recent_tasks), + "refreshSubmissions": max(1, int(refresh_submissions)), + "recentTerminalWindow": max(1, int(recent_terminal_window)), + "longTermMinSamples": max(1, int(long_term_min_samples)), + "acceptedSinceRefresh": 0, + "acceptedTotal": 0, + "acceptedByCategory": _empty_category_counts(), + "longTermStats": all_time_ranked, + "recentStats": recent_ranked, + "lastRefreshError": None, + "refreshRetryAfter": None, + } + + +def choose_next_category(counts: dict[str, int]) -> str: + normalized = {category: max(0, int(counts.get(category, 0))) for category in CATEGORIES} + next_total = sum(normalized.values()) + 1 + + def deficit(category: str) -> int: + return CATEGORY_WEIGHTS[category] * next_total - normalized[category] * sum(CATEGORY_WEIGHTS.values()) + + return max(CATEGORIES, key=lambda category: (deficit(category), -CATEGORIES.index(category))) + + +class GPUStrategyManager: + def __init__( + self, + path: Path | str = DEFAULT_GPU_STRATEGY_PATH, + *, + refresh_submissions: int = DEFAULT_REFRESH_SUBMISSIONS, + recent_terminal_window: int = DEFAULT_RECENT_TERMINAL_WINDOW, + long_term_min_samples: int = DEFAULT_LONG_TERM_MIN_SAMPLES, + log_fn: Callable[[str], None] | None = None, + ) -> None: + self.path = Path(path) + self.refresh_submissions = max(1, int(refresh_submissions)) + self.recent_terminal_window = max(1, int(recent_terminal_window)) + self.long_term_min_samples = max(1, int(long_term_min_samples)) + self.log = log_fn or (lambda message: print(message, flush=True)) + self.state: dict[str, Any] | None = None + + def _load(self) -> dict[str, Any] | None: + try: + value = read_json(self.path) + except (FileNotFoundError, ValueError): + return None + return value if isinstance(value, dict) else None + + def _is_compatible(self, state: dict[str, Any] | None, supported_gpus: list[str]) -> bool: + if not state or int(state.get("version") or 0) != STRATEGY_STATE_VERSION: + return False + return list(state.get("supportedGpus") or []) == supported_gpus + + def _refresh_due(self, state: dict[str, Any] | None, supported_gpus: list[str], now: datetime) -> tuple[bool, str]: + if not self._is_compatible(state, supported_gpus): + return True, "initial_or_gpu_catalog_changed" + retry_after = parse_datetime(state.get("refreshRetryAfter")) if state is not None else None + if state is not None and state.get("lastRefreshError") and retry_after is not None and retry_after > now: + return False, "refresh_error_backoff" + if not bool(state.get("historyReady", False)): + return (retry_after is None or retry_after <= now), "history_not_ready" + accepted = int(state.get("acceptedSinceRefresh") or 0) + refresh_every = max(1, int(state.get("refreshSubmissions") or self.refresh_submissions)) + return accepted >= refresh_every, "accepted_submission_threshold" + + @staticmethod + def _load_platform_history(client: Any) -> list[dict[str, Any]]: + if isinstance(client, ModelHubClientPool): + by_account: dict[int, list[dict[str, Any]]] = {} + errors: list[int] = [] + with ThreadPoolExecutor(max_workers=min(len(client.clients), 12)) as executor: + futures = { + executor.submit(account_client.list_tasks, page_size=100, only_mine=True): index + for index, account_client in enumerate(client.clients, start=1) + } + for future in as_completed(futures): + index = futures[future] + try: + by_account[index] = future.result() + except Exception: + errors.append(index) + if errors: + joined = ",".join(str(index) for index in sorted(errors)) + raise RuntimeError(f"ModelHub history fetch failed for account indexes: {joined}") + tasks = [task for index in sorted(by_account) for task in by_account[index]] + else: + tasks = client.list_tasks(page_size=100, only_mine=True) + + deduped: list[dict[str, Any]] = [] + seen_task_ids: set[str] = set() + for task in tasks: + if not isinstance(task, dict): + continue + task_id = str(task.get("taskId")) if task.get("taskId") is not None else None + if task_id and task_id in seen_task_ids: + continue + if task_id: + seen_task_ids.add(task_id) + deduped.append(task) + return deduped + + def prepare( + self, + client: Any, + *, + supported_gpus: list[str], + now: datetime | None = None, + ) -> dict[str, Any]: + now = now or utc_now() + supported_gpus = list(dict.fromkeys(gpu for gpu in supported_gpus if gpu)) + state = self._load() + refresh_due, reason = self._refresh_due(state, supported_gpus, now) + if not refresh_due and state is not None: + self.state = state + self._log_state("loaded") + return state + + previous_generation = int(state.get("generation", -1)) if state is not None else -1 + self.log(f"[strategy] refresh_start reason={reason} supported_gpus={len(supported_gpus)}") + try: + tasks = self._load_platform_history(client) + refreshed = build_strategy_snapshot( + tasks, + supported_gpus=supported_gpus, + generated_at=now, + generation=previous_generation + 1, + recent_terminal_window=self.recent_terminal_window, + long_term_min_samples=self.long_term_min_samples, + refresh_submissions=self.refresh_submissions, + ) + refreshed["acceptedTotal"] = int((state or {}).get("acceptedTotal") or 0) + self.state = refreshed + write_json(self.path, refreshed) + self._log_state("refreshed") + return refreshed + except Exception as exc: + self.log(f"[strategy] refresh_error reason={type(exc).__name__}: {exc}") + if self._is_compatible(state, supported_gpus): + assert state is not None + state["lastRefreshError"] = str(exc) + state["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat() + self.state = state + write_json(self.path, state) + return state + + fallback = build_strategy_snapshot( + [], + supported_gpus=supported_gpus, + generated_at=now, + generation=0, + recent_terminal_window=self.recent_terminal_window, + long_term_min_samples=self.long_term_min_samples, + refresh_submissions=self.refresh_submissions, + ) + fallback["historyReady"] = False + fallback["lastRefreshError"] = str(exc) + fallback["refreshRetryAfter"] = (now + timedelta(minutes=5)).isoformat() + self.state = fallback + write_json(self.path, fallback) + self._log_state("fallback") + return fallback + + def _log_state(self, action: str) -> None: + if self.state is None: + return + counts = self.state.get("acceptedByCategory") or {} + self.log( + f"[strategy] {action} generation={self.state.get('generation', 0)} " + f"accepted={self.state.get('acceptedSinceRefresh', 0)}/{self.state.get('refreshSubmissions', self.refresh_submissions)} " + f"categories={counts.get(LONG_TERM, 0)},{counts.get(ALL_SUPPORTED, 0)},{counts.get(RECENT, 0)} " + f"long={','.join(self.state.get('longTermGpus') or [])} recent={self.state.get('recentGpu') or 'n/a'}" + ) + + @property + def submissions_until_refresh(self) -> int: + if self.state is None: + return self.refresh_submissions + refresh_every = max(1, int(self.state.get("refreshSubmissions") or self.refresh_submissions)) + accepted = max(0, int(self.state.get("acceptedSinceRefresh") or 0)) + return max(0, refresh_every - accepted) + + def _category_gpu_order(self, category: str, occurrence: int) -> list[str]: + assert self.state is not None + supported = list(self.state.get("supportedGpus") or []) + long_term = [gpu for gpu in self.state.get("longTermGpus") or [] if gpu in supported] + recent_gpu = str(self.state.get("recentGpu") or "") + + if category == LONG_TERM and long_term: + desired_index = LONG_TERM_GPU_PATTERN[occurrence % len(LONG_TERM_GPU_PATTERN)] + desired = long_term[min(desired_index, len(long_term) - 1)] + return [desired, *(gpu for gpu in long_term if gpu != desired)] + if category == RECENT and recent_gpu in supported: + return [recent_gpu] + if supported: + offset = occurrence % len(supported) + return [*supported[offset:], *supported[:offset]] + return [] + + def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: + if self.state is None or not candidates: + return list(candidates) + + by_gpu: dict[str, list[dict[str, Any]]] = defaultdict(list) + for candidate in candidates: + by_gpu[str(candidate.get("targetGpu") or "")].append(candidate) + gpu_indexes: dict[str, int] = defaultdict(int) + used: set[str] = set() + ordered: list[dict[str, Any]] = [] + virtual_counts = { + category: max(0, int((self.state.get("acceptedByCategory") or {}).get(category, 0))) + for category in CATEGORIES + } + + def take_from_gpu(gpu: str) -> dict[str, Any] | None: + pool = by_gpu.get(gpu) or [] + index = gpu_indexes[gpu] + while index < len(pool): + candidate = pool[index] + index += 1 + gpu_indexes[gpu] = index + if candidate_key(candidate) not in used: + return candidate + gpu_indexes[gpu] = index + return None + + while len(ordered) < len(candidates): + planned_category = choose_next_category(virtual_counts) + selected: dict[str, Any] | None = None + actual_category = planned_category + categories_to_try = [planned_category, *(category for category in CATEGORIES if category != planned_category)] + for category in categories_to_try: + occurrence = virtual_counts[category] + for gpu in self._category_gpu_order(category, occurrence): + selected = take_from_gpu(gpu) + if selected is not None: + actual_category = category + break + if selected is not None: + break + + if selected is None: + # Unknown/custom GPUs can only appear when callers bypass the normal resolver. + selected = next((candidate for candidate in candidates if candidate_key(candidate) not in used), None) + actual_category = ALL_SUPPORTED + if selected is None: + break + + key = candidate_key(selected) + used.add(key) + annotated = dict(selected) + annotated["strategyCategory"] = actual_category + annotated["strategyPlannedCategory"] = planned_category + annotated["strategyGeneration"] = int(self.state.get("generation") or 0) + ordered.append(annotated) + virtual_counts[actual_category] += 1 + + return ordered + + def record_accepted(self, candidates: list[dict[str, Any]]) -> dict[str, Any] | None: + if not candidates: + return self.state + state = self._load() or self.state + if state is None: + return None + + category_counts = { + category: max(0, int((state.get("acceptedByCategory") or {}).get(category, 0))) + for category in CATEGORIES + } + for candidate in candidates: + category = str(candidate.get("strategyCategory") or ALL_SUPPORTED) + if category not in category_counts: + category = ALL_SUPPORTED + category_counts[category] += 1 + + accepted_count = len(candidates) + state["acceptedSinceRefresh"] = int(state.get("acceptedSinceRefresh") or 0) + accepted_count + state["acceptedTotal"] = int(state.get("acceptedTotal") or 0) + accepted_count + state["acceptedByCategory"] = category_counts + self.state = state + write_json(self.path, state) + self._log_state("progress") + return state + + def summary(self) -> dict[str, Any]: + if self.state is None: + return {"enabled": False} + return { + "enabled": True, + "statePath": str(self.path), + "generation": int(self.state.get("generation") or 0), + "generatedAt": self.state.get("generatedAt"), + "historyReady": bool(self.state.get("historyReady", False)), + "historyTaskCount": int(self.state.get("historyTaskCount") or 0), + "acceptedSinceRefresh": int(self.state.get("acceptedSinceRefresh") or 0), + "refreshSubmissions": int(self.state.get("refreshSubmissions") or self.refresh_submissions), + "acceptedByCategory": dict(self.state.get("acceptedByCategory") or {}), + "longTermGpus": list(self.state.get("longTermGpus") or []), + "recentGpu": self.state.get("recentGpu"), + "recentTerminalCount": int(self.state.get("recentTerminalCount") or 0), + "refreshDue": self.submissions_until_refresh <= 0, + } diff --git a/modelhub_submmit_api/hf_discovery.py b/modelhub_submmit_api/hf_discovery.py index 1ed73c5c..af38a7c7 100644 --- a/modelhub_submmit_api/hf_discovery.py +++ b/modelhub_submmit_api/hf_discovery.py @@ -2,6 +2,7 @@ from __future__ import annotations import os import threading +import time from pathlib import PurePosixPath from typing import Any from urllib.parse import quote @@ -44,7 +45,9 @@ class HuggingFaceDiscovery: http_client: JsonHttpClient | None = None, legacy_http_client: JsonHttpClient | None = None, timeout: int = 30, - retries: int = 2, + retries: int = 5, + page_interval_seconds: float | None = None, + page_cache_ttl_seconds: float | None = None, ) -> None: token = os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN headers = {"User-Agent": "modelhub-submmit-cli/0.1"} @@ -56,6 +59,7 @@ class HuggingFaceDiscovery: default_headers=headers, timeout=timeout, retries=retries, + backoff_seconds=2.0, ) self.legacy_http_client = legacy_http_client or JsonHttpClient( base_url=base_url, @@ -65,6 +69,24 @@ class HuggingFaceDiscovery: ) self._repo_tree_cache: dict[str, list[dict[str, Any]]] = {} self._repo_tree_lock = threading.Lock() + self._model_page_cache: dict[tuple[str, int, int], tuple[float, list[dict[str, Any]]]] = {} + self._model_page_cache_ttl = max( + 0.0, + float( + page_cache_ttl_seconds + if page_cache_ttl_seconds is not None + else os.getenv("MODELSCOPE_PAGE_CACHE_TTL_SECONDS", "900") + ), + ) + self._page_interval_seconds = max( + 0.0, + float( + page_interval_seconds + if page_interval_seconds is not None + else os.getenv("MODELSCOPE_PAGE_INTERVAL_SECONDS", "0.25") + ), + ) + self._last_model_page_request_at = 0.0 def list_recent_models( self, @@ -111,22 +133,37 @@ class HuggingFaceDiscovery: 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 + cache_key = (task_tag, page_number, page_size) + cached = self._model_page_cache.get(cache_key) + if cached is not None and time.monotonic() - cached[0] < self._model_page_cache_ttl: + items = list(cached[1]) + else: + elapsed = time.monotonic() - self._last_model_page_request_at + if self._last_model_page_request_at > 0 and elapsed < self._page_interval_seconds: + time.sleep(self._page_interval_seconds - elapsed) + 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, + }, + ) + self._last_model_page_request_at = time.monotonic() + except HttpJsonError as exc: + self._last_model_page_request_at = time.monotonic() + print( + f"[modelscope] list_models_error task={task_tag} page={page_number} " + f"partial_models={len(models)} retry_next_cycle=true error={exc}", + flush=True, + ) + break - items = self._extract_models(payload) + items = self._extract_models(payload) + self._model_page_cache[cache_key] = (time.monotonic(), list(items)) if not items: break for item in items: diff --git a/modelhub_submmit_api/history_stats.py b/modelhub_submmit_api/history_stats.py index 085f1f00..a3d1993e 100644 --- a/modelhub_submmit_api/history_stats.py +++ b/modelhub_submmit_api/history_stats.py @@ -231,23 +231,32 @@ def summarize_records(records: list[dict[str, Any]]) -> dict[str, Any]: def is_success(record: dict[str, Any]) -> bool: - verify_result = record.get("verifyResult") + verify_result = _numeric_verify_result(record.get("verifyResult")) status = str(record.get("status") or "").lower() return verify_result is not None and verify_result > 0 and status == "success" def is_failure(record: dict[str, Any]) -> bool: - verify_result = record.get("verifyResult") + verify_result = _numeric_verify_result(record.get("verifyResult")) status = str(record.get("status") or "").lower() if verify_result is not None and verify_result < 0: return True - return status in {"failed", "error", "cancelled"} + return status in {"failed", "error", "cancelled", "canceled", "rejected", "timeout"} + + +def _numeric_verify_result(value: Any) -> float | None: + if value in (None, ""): + return None + try: + return float(value) + except (TypeError, ValueError): + return None def classify_failure(record: dict[str, Any]) -> str: status = str(record.get("status") or "").lower() - verify_result = record.get("verifyResult") - if status in {"failed", "error", "cancelled"}: + verify_result = _numeric_verify_result(record.get("verifyResult")) + if status in {"failed", "error", "cancelled", "canceled", "rejected", "timeout"}: return "参数/模板问题" if status in WAITING_STATUSES or status in RUNNING_STATUSES or verify_result is None: return "排队中/未知" diff --git a/modelhub_submmit_api/http_json.py b/modelhub_submmit_api/http_json.py index 21ef902f..7d7083dc 100644 --- a/modelhub_submmit_api/http_json.py +++ b/modelhub_submmit_api/http_json.py @@ -4,6 +4,8 @@ import json import socket import time from dataclasses import dataclass +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime from http.client import IncompleteRead, RemoteDisconnected from typing import Any from urllib.error import HTTPError, URLError @@ -76,17 +78,34 @@ class JsonHttpClient: except HTTPError as exc: payload = self._decode_error_payload(exc) if exc.code in retryable_statuses and attempt < self.retries: - time.sleep(self.backoff_seconds * (attempt + 1)) + retry_after = exc.headers.get("Retry-After") if exc.headers is not None else None + time.sleep(self._retry_delay(attempt, retry_after=retry_after)) continue raise HttpJsonError(f"HTTP request failed for {url}", status_code=exc.code, payload=payload) from exc except (URLError, TimeoutError, socket.timeout, RemoteDisconnected, IncompleteRead) as exc: if attempt < self.retries: - time.sleep(self.backoff_seconds * (attempt + 1)) + time.sleep(self._retry_delay(attempt)) continue raise HttpJsonError(f"Network request failed for {url}") from exc raise HttpJsonError(f"Request exhausted retries for {url}") + def _retry_delay(self, attempt: int, *, retry_after: str | None = None) -> float: + exponential_delay = self.backoff_seconds * (2**attempt) + retry_after_delay = 0.0 + if retry_after: + try: + retry_after_delay = max(0.0, float(retry_after)) + except ValueError: + try: + parsed = parsedate_to_datetime(retry_after) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + retry_after_delay = max(0.0, (parsed - datetime.now(timezone.utc)).total_seconds()) + except (TypeError, ValueError, OverflowError): + retry_after_delay = 0.0 + return min(60.0, max(exponential_delay, retry_after_delay)) + def _build_url(self, path: str, query: dict[str, Any] | None) -> str: if path.startswith("http://") or path.startswith("https://"): url = path diff --git a/modelhub_submmit_api/main.py b/modelhub_submmit_api/main.py index c34059fe..e92663a3 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, runtime_instance_id, utc_now, write_json, write_jsonl +from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH, GPUStrategyManager from hf_discovery import HuggingFaceDiscovery from history_stats import ( append_ledger_entry, @@ -16,7 +17,7 @@ from history_stats import ( load_ledger, update_history_archive, ) -from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool, is_duplicate_submission_error +from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubAPIError, ModelHubClient, ModelHubClientPool, is_duplicate_submission_error from models import CandidateModel, HFModelSummary, ModelInspection from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker from submission_claims import DEFAULT_CLAIMS_PATH, SubmissionClaimStore, candidate_key, diversify_candidates @@ -44,6 +45,17 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--min-downloads", type=int, default=50, help="Minimum downloads threshold") parser.add_argument("--daily-target", type=int, default=0, help="Daily submission target across all auto runs; 0 means unlimited") parser.add_argument("--dry-run", action="store_true", help="Only write artifacts without creating tasks") + parser.add_argument( + "--disable-gpu-strategy", + action="store_true", + help="Disable adaptive 50/30/20 GPU scheduling and keep the legacy candidate order", + ) + parser.add_argument( + "--gpu-strategy-refresh-submissions", + type=int, + default=200, + help="Recalculate adaptive GPU choices after this many accepted submissions", + ) parser.add_argument("--skip-outcome-sync", action="store_true", help="Skip outcome sync from ModelHub before scanning") parser.add_argument("--skip-history-archive", action="store_true", help="Skip historical task archive download for this run") group = parser.add_mutually_exclusive_group() @@ -79,6 +91,24 @@ 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( + "--gpu-strategy-state-path", + default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument("--gpu-strategy-recent-window", type=int, default=1000, help=argparse.SUPPRESS) + parser.add_argument("--gpu-strategy-min-long-samples", type=int, default=100, help=argparse.SUPPRESS) + parser.add_argument( + "--capacity-state-path", + default=os.getenv("MODELHUB_CAPACITY_STATE_PATH", str(DEFAULT_CAPACITY_STATE_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--capacity-probe-interval-cycles", + type=int, + default=int(os.getenv("MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES", "3")), + 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) @@ -376,7 +406,11 @@ def run_submission( modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token])) token_values: list[str | None] = modelhub_tokens or [args.modelhub_token] clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in token_values] - modelhub_client = ModelHubClientPool(clients) + modelhub_client = ModelHubClientPool( + clients, + capacity_probe_interval_cycles=max(0, int(getattr(args, "capacity_probe_interval_cycles", 3) or 0)), + capacity_state_path=Path(getattr(args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)), + ) if hasattr(modelhub_client, "begin_cycle"): modelhub_client.begin_cycle() @@ -389,6 +423,13 @@ def run_submission( history_archive_path.parent.mkdir(parents=True, exist_ok=True) outcome_tracker = outcome_tracker or OutcomeTracker(Path(args.outcomes_path)) + strategy_manager: GPUStrategyManager | None = None + strategy_summary: dict[str, Any] = {"enabled": False} + strategy_enabled = ( + not bool(getattr(args, "disable_gpu_strategy", False)) + and not bool(getattr(args, "gpu", None) or getattr(args, "gpus", None)) + and len(target_gpus) > 1 + ) synced_count = 0 if not getattr(args, "skip_outcome_sync", False): try: @@ -466,6 +507,7 @@ def run_submission( "historyStatsThreshold": getattr(args, "history_stats_threshold", 500), "historyArchivePath": str(history_archive_path), "historyArchiveRecordCount": len(archived_history), + "gpuStrategy": strategy_summary, "scanLimit": 0, "scannedModels": 0, "candidateCount": 0, @@ -481,6 +523,16 @@ def run_submission( "outcomeSyncCount": synced_count, } + if strategy_enabled: + strategy_manager = GPUStrategyManager( + Path(getattr(args, "gpu_strategy_state_path", DEFAULT_GPU_STRATEGY_PATH)), + refresh_submissions=max(1, int(getattr(args, "gpu_strategy_refresh_submissions", 200) or 200)), + recent_terminal_window=max(1, int(getattr(args, "gpu_strategy_recent_window", 1000) or 1000)), + long_term_min_samples=max(1, int(getattr(args, "gpu_strategy_min_long_samples", 100) or 100)), + ) + strategy_manager.prepare(modelhub_client, supported_gpus=target_gpus, now=now) + strategy_summary = strategy_manager.summary() + if getattr(args, "skip_history_archive", False): archived_history = [] else: @@ -561,8 +613,6 @@ def run_submission( skipped.extend(model_skipped) failed.extend(model_failed) - write_jsonl(run_dir / "candidates.jsonl", candidates) - submitted: list[dict[str, Any]] = [] duplicate_candidates: list[dict[str, Any]] = [] target_submit_count = resolve_max_submit_count( @@ -570,8 +620,13 @@ def run_submission( planned_count=len(candidates), remaining_daily_quota=remaining_daily_quota, ) + if strategy_manager is not None: + target_submit_count = min(target_submit_count, strategy_manager.submissions_until_refresh) instance_id = runtime_instance_id() diversified_candidates = diversify_candidates(candidates, instance_id=instance_id) + if strategy_manager is not None: + diversified_candidates = strategy_manager.order_candidates(diversified_candidates) + write_jsonl(run_dir / "candidates.jsonl", diversified_candidates) claim_store: SubmissionClaimStore | None = None attempted_candidates: list[dict[str, Any]] = [] submit_workers = 1 @@ -698,6 +753,8 @@ def run_submission( claim_store.mark_submitted([*batch_submitted_candidates, *batch_duplicate_candidates]) claim_store.release(batch_failed_candidates) + if strategy_manager is not None: + strategy_manager.record_accepted(batch_submitted_candidates) if hasattr(modelhub_client, "available_submit_slots") and modelhub_client.available_submit_slots() <= 0: break @@ -706,6 +763,9 @@ def run_submission( write_jsonl(run_dir / "skipped.jsonl", skipped) write_jsonl(run_dir / "failed.jsonl", failed) + if strategy_manager is not None: + strategy_summary = strategy_manager.summary() + outcome_tracker.save() summary = { @@ -723,6 +783,7 @@ def run_submission( "historyStatsThreshold": getattr(args, "history_stats_threshold", 500), "historyArchivePath": str(history_archive_path), "historyArchiveRecordCount": len(archived_history), + "gpuStrategy": strategy_summary, "scanLimit": scan_limit, "scannedModels": len(models), "candidateCount": len(candidates), diff --git a/modelhub_submmit_api/modelhub_client.py b/modelhub_submmit_api/modelhub_client.py index 5bb287b4..ddbab7ed 100644 --- a/modelhub_submmit_api/modelhub_client.py +++ b/modelhub_submmit_api/modelhub_client.py @@ -6,9 +6,10 @@ import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta +from pathlib import Path from typing import Any -from common import format_modelhub_datetime, parse_datetime, runtime_instance_id +from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, write_json from defaults import EMBEDDED_MODELHUB_XC_TOKEN from http_json import HttpJsonError, JsonHttpClient @@ -255,6 +256,8 @@ DUPLICATE_SUBMISSION_MARKERS = ( "already in progress", ) +DEFAULT_CAPACITY_STATE_PATH = Path(".modelhub_state/account_capacity.json") + def is_capacity_error(error: ModelHubAPIError) -> bool: if error.code in {409, 429}: @@ -281,12 +284,24 @@ class ModelHubClientPool: active_counts_ttl: float | None = None, reservation_ttl: float | None = None, instance_id: str | None = None, + capacity_probe_interval_cycles: int = 3, + capacity_probe_cooldown_cycles: int = 3, + capacity_state_path: Path | str | None = None, ) -> None: if not clients: raise ValueError("At least one ModelHub client is required") self.clients = clients configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100") self.active_task_cap = max(1, int(configured_cap)) + self._capacity_state_path = Path(capacity_state_path) if capacity_state_path else None + self._account_keys = [self._account_key(client, index) for index, client in enumerate(clients)] + self._account_caps = self._load_account_caps(self.active_task_cap) + self._capacity_probe_interval_cycles = max(0, int(capacity_probe_interval_cycles)) + self._capacity_probe_cooldown_cycles = max(1, int(capacity_probe_cooldown_cycles)) + self._capacity_probe_cycle = 0 + self._capacity_probe_enabled = False + self._capacity_probe_attempted: set[int] = set() + self._capacity_probe_cooldown_until: list[int] = [0 for _ in clients] configured_ttl = active_counts_ttl if active_counts_ttl is not None else os.getenv("MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS", "15") configured_reservation_ttl = ( reservation_ttl @@ -310,11 +325,65 @@ class ModelHubClientPool: # Single reader client to avoid fanout on read operations self._reader = clients[0] - def _safe_count_active_tasks(self, client: ModelHubClient) -> int: + @staticmethod + def _account_key(client: ModelHubClient, index: int) -> str: + token = str(getattr(client, "token", "") or "") + identity = token if token else f"account-index:{index}" + return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20] + + def _load_account_caps(self, default_cap: int) -> list[int]: + stored: dict[str, Any] = {} + if self._capacity_state_path is not None: + try: + payload = read_json(self._capacity_state_path) + if isinstance(payload, dict): + stored = payload.get("accounts") or {} + except (FileNotFoundError, ValueError): + pass + return [ + max(1, int((stored.get(key) or {}).get("knownCap") or default_cap)) + for key in self._account_keys + ] + + def _persist_account_caps(self) -> None: + if self._capacity_state_path is None: + return + with self._state_lock: + payload = { + "version": 1, + "updatedAt": datetime.now().astimezone().isoformat(), + "accounts": { + key: {"accountIndex": index + 1, "knownCap": self._account_caps[index]} + for index, key in enumerate(self._account_keys) + }, + } + write_json(self._capacity_state_path, payload) + + def configure_capacity_probe(self, cycle_number: int) -> None: + cycle_number = max(0, int(cycle_number)) + with self._state_lock: + if cycle_number != self._capacity_probe_cycle: + self._capacity_probe_attempted.clear() + self._capacity_probe_cycle = cycle_number + self._capacity_probe_enabled = ( + self._capacity_probe_interval_cycles > 0 + and cycle_number > 0 + and cycle_number % self._capacity_probe_interval_cycles == 0 + ) + + def account_capacity_limits(self) -> list[int]: + with self._state_lock: + return list(self._account_caps) + + def capacity_probe_enabled(self) -> bool: + with self._state_lock: + return self._capacity_probe_enabled + + def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int: try: - return client.count_active_tasks(max_count=self.active_task_cap, page_size=200) + return client.count_active_tasks(max_count=max_count, page_size=200) except Exception: - return self.active_task_cap + return max(0, max_count - 1) def _safe_search_by_model_id(self, client: ModelHubClient, model_id: str) -> dict[str, Any]: try: @@ -346,8 +415,11 @@ class ModelHubClientPool: if not force and self._counts_are_fresh_locked(now): return + with self._state_lock: + count_limits = [cap + 1 for cap in self._account_caps] + def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]: - return index, self._safe_count_active_tasks(client) + return index, self._safe_count_active_tasks(client, count_limits[index]) results: list[tuple[int, int]] = [] with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor: @@ -363,10 +435,14 @@ class ModelHubClientPool: results.append((index, self.active_task_cap)) refreshed_at = time.monotonic() + caps_changed = False with self._state_lock: for index, count in results: old_remote_count = self._remote_counts[index] - new_remote_count = min(self.active_task_cap, max(0, int(count))) + new_remote_count = max(0, int(count)) + if new_remote_count > self._account_caps[index]: + self._account_caps[index] = new_remote_count + caps_changed = True acknowledged = max(0, new_remote_count - old_remote_count) completed_reservations = sorted( ( @@ -390,12 +466,25 @@ class ModelHubClientPool: self._remote_counts[index] = new_remote_count self._counts_initialized = True self._active_refresh_at = refreshed_at + if caps_changed: + self._persist_account_caps() def _effective_count_locked(self, index: int) -> int: return self._remote_counts[index] + len(self._reservations[index]) def _counts_snapshot_locked(self) -> list[int]: - return [min(self.active_task_cap, self._effective_count_locked(index)) for index in range(len(self.clients))] + return [self._effective_count_locked(index) for index in range(len(self.clients))] + + def _probe_slot_count_locked(self) -> int: + if not self._capacity_probe_enabled: + return 0 + return sum( + 1 + for index in range(len(self.clients)) + if index not in self._capacity_probe_attempted + and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index] + and self._effective_count_locked(index) >= self._account_caps[index] + ) def active_task_counts(self) -> list[int]: self._refresh_active_counts() @@ -405,10 +494,11 @@ class ModelHubClientPool: def available_submit_slots(self) -> int: self._refresh_active_counts() with self._state_lock: - return sum( - max(0, self.active_task_cap - self._effective_count_locked(index)) + normal_slots = sum( + max(0, self._account_caps[index] - self._effective_count_locked(index)) for index in range(len(self.clients)) ) + return normal_slots if normal_slots > 0 else self._probe_slot_count_locked() def list_tasks(self, **kwargs): # noqa: ANN003, ANN001 # Default: use only the reader client to avoid fanout amplification. @@ -480,7 +570,7 @@ class ModelHubClientPool: def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None: with self._state_lock: remaining_by_index = { - index: self.active_task_cap - self._effective_count_locked(index) + index: self._account_caps[index] - self._effective_count_locked(index) for index in range(len(self.clients)) if index not in excluded } @@ -502,6 +592,35 @@ class ModelHubClientPool: } return selected_index, reservation_id + def _reserve_probe_account(self, excluded: set[int]) -> tuple[int, int] | None: + with self._state_lock: + if not self._capacity_probe_enabled: + return None + eligible = [ + index + for index in range(len(self.clients)) + if index not in excluded + and index not in self._capacity_probe_attempted + and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index] + and self._effective_count_locked(index) >= self._account_caps[index] + ] + if not eligible: + return None + selected_index = min( + eligible, + key=lambda index: (index - self._selection_cursor) % len(self.clients), + ) + self._selection_cursor = (selected_index + 1) % len(self.clients) + self._capacity_probe_attempted.add(selected_index) + self._reservation_sequence += 1 + reservation_id = self._reservation_sequence + self._reservations[selected_index][reservation_id] = { + "inflight": True, + "updated_at": time.monotonic(), + "capacity_probe": True, + } + return selected_index, reservation_id + def _finish_reservation(self, index: int, reservation_id: int, *, succeeded: bool) -> None: with self._state_lock: reservation = self._reservations[index].get(reservation_id) @@ -513,11 +632,32 @@ class ModelHubClientPool: reservation["inflight"] = False reservation["updated_at"] = time.monotonic() - def _mark_account_saturated(self, index: int) -> None: + def _mark_account_saturated(self, index: int, *, capacity_probe: bool) -> None: + cap_changed = False with self._state_lock: - self._remote_counts[index] = self.active_task_cap + if capacity_probe: + self._capacity_probe_cooldown_until[index] = ( + self._capacity_probe_cycle + self._capacity_probe_cooldown_cycles + ) + else: + observed_capacity = max(1, self._effective_count_locked(index)) + if observed_capacity < self._account_caps[index]: + self._account_caps[index] = observed_capacity + cap_changed = True + self._remote_counts[index] = self._account_caps[index] self._counts_initialized = True self._active_refresh_at = time.monotonic() + if cap_changed: + self._persist_account_caps() + + def _promote_account_capacity(self, index: int) -> None: + with self._state_lock: + discovered_cap = max(self._account_caps[index] + 1, self._effective_count_locked(index)) + if discovered_cap <= self._account_caps[index]: + return + self._account_caps[index] = discovered_cap + self._persist_account_caps() + print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True) def add_task(self, payload: dict[str, Any]) -> dict[str, Any]: self._refresh_active_counts() @@ -527,6 +667,10 @@ class ModelHubClientPool: while True: reservation = self._reserve_account(attempted_accounts) + capacity_probe = False + if reservation is None: + reservation = self._reserve_probe_account(attempted_accounts) + capacity_probe = reservation is not None if reservation is None: if not forced_refresh_done: self._refresh_active_counts(force=True) @@ -535,7 +679,7 @@ class ModelHubClientPool: if last_capacity_error is not None: raise last_capacity_error raise ModelHubAPIError( - f"当前等待中或运行中的异步模型验证任务数量已达上限（{self.active_task_cap}）" + "当前等待中或运行中的异步模型验证任务数量已达已知上限" ) selected_index, reservation_id = reservation @@ -548,7 +692,7 @@ class ModelHubClientPool: raise # Another process may have filled this account after our count # refresh. Mark it full locally and immediately try another one. - self._mark_account_saturated(selected_index) + self._mark_account_saturated(selected_index, capacity_probe=capacity_probe) attempted_accounts.add(selected_index) last_capacity_error = exc continue @@ -557,6 +701,8 @@ class ModelHubClientPool: raise self._finish_reservation(selected_index, reservation_id, succeeded=True) + if capacity_probe: + self._promote_account_capacity(selected_index) return response def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001 diff --git a/modelhub_submmit_api/poll_runner.py b/modelhub_submmit_api/poll_runner.py index a7fa62db..b1dd5ff5 100644 --- a/modelhub_submmit_api/poll_runner.py +++ b/modelhub_submmit_api/poll_runner.py @@ -10,9 +10,10 @@ 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 gpu_strategy import DEFAULT_GPU_STRATEGY_PATH from hf_discovery import HuggingFaceDiscovery from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir -from modelhub_client import ModelHubClient, ModelHubClientPool +from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker from runner_common import DEFAULT_KEY_PATH, ensure_tokens from submission_claims import DEFAULT_CLAIMS_PATH @@ -58,6 +59,13 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--skip-outcome-sync", action="store_true", help="Skip outcome sync from ModelHub before scanning") parser.add_argument("--skip-history-archive", action="store_true", help="Skip historical task archive download for this run") parser.add_argument("--dry-run", action="store_true", help="Plan the day without creating tasks") + parser.add_argument("--disable-gpu-strategy", action="store_true", help="Disable adaptive 50/30/20 GPU scheduling") + parser.add_argument( + "--gpu-strategy-refresh-submissions", + type=int, + default=200, + help="Recalculate adaptive GPU choices after this many accepted submissions", + ) 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=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS) parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS) @@ -68,6 +76,24 @@ 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( + "--gpu-strategy-state-path", + default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument("--gpu-strategy-recent-window", type=int, default=1000, help=argparse.SUPPRESS) + parser.add_argument("--gpu-strategy-min-long-samples", type=int, default=100, help=argparse.SUPPRESS) + parser.add_argument( + "--capacity-state-path", + default=os.getenv("MODELHUB_CAPACITY_STATE_PATH", str(DEFAULT_CAPACITY_STATE_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--capacity-probe-interval-cycles", + type=int, + default=int(os.getenv("MODELHUB_CAPACITY_PROBE_INTERVAL_CYCLES", "3")), + 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("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS) @@ -97,7 +123,11 @@ def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClientPool: modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token])) token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token] clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values] - return ModelHubClientPool(clients) + return ModelHubClientPool( + clients, + capacity_probe_interval_cycles=max(0, int(getattr(base_args, "capacity_probe_interval_cycles", 3) or 0)), + capacity_state_path=Path(getattr(base_args, "capacity_state_path", DEFAULT_CAPACITY_STATE_PATH)), + ) def run_poll_loop( @@ -140,10 +170,16 @@ def run_poll_loop( break cycles += 1 + if hasattr(modelhub_client, "configure_capacity_probe"): + modelhub_client.configure_capacity_probe(cycles) active_counts = modelhub_client.active_task_counts() if hasattr(modelhub_client, "active_task_counts") else [] + capacity_limits = modelhub_client.account_capacity_limits() if hasattr(modelhub_client, "account_capacity_limits") else [] + capacity_probe = modelhub_client.capacity_probe_enabled() if hasattr(modelhub_client, "capacity_probe_enabled") else False available_slots = modelhub_client.available_submit_slots() if hasattr(modelhub_client, "available_submit_slots") else None log( f"[poll] cycle={cycles} active_counts={','.join(str(count) for count in active_counts) if active_counts else 'n/a'} " + f"capacity_limits={','.join(str(limit) for limit in capacity_limits) if capacity_limits else 'n/a'} " + f"capacity_probe={'on' if capacity_probe else 'off'} " f"available_slots={available_slots if available_slots is not None else 'n/a'}" ) @@ -152,8 +188,10 @@ def run_poll_loop( time.sleep(base_args.poll_interval_seconds) continue + cycle_args = _make_cycle_args(base_args) + cycle_args.capacity_probe_already_configured = True cycle_summary = run_fn( - base_args=_make_cycle_args(base_args), + base_args=cycle_args, now=utc_now(), hf_discovery=hf_discovery, modelhub_client=modelhub_client, diff --git a/modelhub_submmit_api/version.py b/modelhub_submmit_api/version.py index 78403870..0a66a1f2 100644 --- a/modelhub_submmit_api/version.py +++ b/modelhub_submmit_api/version.py @@ -1 +1 @@ -AGENT_VERSION = "2026.08.02.2" +AGENT_VERSION = "2026.08.02.3" diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index 0fa5e634..3df411a2 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -14,7 +14,7 @@ if str(PACKAGE_DIR) in sys.path: sys.path.remove(str(PACKAGE_DIR)) sys.path.insert(0, str(PACKAGE_DIR)) -from common import read_jsonl # noqa: E402 +from common import read_json, read_jsonl # noqa: E402 from main import build_parser, make_run_dir, run_submission, submit_candidate # noqa: E402 from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error # noqa: E402 from models import HFModelSummary, ModelInspection # noqa: E402 @@ -106,6 +106,54 @@ class DuplicateThenSuccessClient: return {"code": 0, "data": {"id": f"task-{self.calls}"}} +class DynamicCapacityClient: + def __init__(self, *, active: int, limit: int, token: str = "dynamic-test") -> None: + self.active = active + self.limit = limit + self.token = token + + def count_active_tasks(self, *, max_count: int, **_kwargs) -> int: # noqa: ANN003 + return min(self.active, max_count) + + def add_task(self, _payload: dict) -> dict: + if self.active >= self.limit: + raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限") + self.active += 1 + return {"code": 0, "data": {"id": self.active}} + + +class AutoStrategyClient: + def __init__(self, available: int) -> None: + self.available = available + self.calls = 0 + + @staticmethod + def begin_cycle() -> None: + return None + + def available_submit_slots(self) -> int: + return self.available + + @staticmethod + def list_tasks(**_kwargs) -> list[dict]: # noqa: ANN003 + return [] + + @staticmethod + def list_tasks_page(**_kwargs) -> dict: # noqa: ANN003 + return {"code": 0, "data": {"records": [], "pages": 0}} + + @staticmethod + def processed_gpus_for_model(_model_id: str) -> set[str]: + return set() + + def add_task(self, _payload: dict) -> dict: + self.calls += 1 + if self.calls == 1: + raise ModelHubAPIError("模型正在验证中，请勿重复提交") + self.available -= 1 + return {"code": 0, "data": {"id": f"strategy-{self.calls}"}} + + def make_candidate(index: int) -> dict: return { "repoId": f"owner/model-{index}", @@ -115,6 +163,42 @@ def make_candidate(index: int) -> dict: class ClientPoolConcurrencyTests(unittest.TestCase): + def test_adaptive_strategy_counts_only_platform_accepted_tasks(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + args = build_parser().parse_args( + [ + "--task-types", + "text-generation", + "--limit", + "4", + "--max-scan-models", + "4", + "--skip-outcome-sync", + "--skip-history-archive", + ] + ) + args.runs_dir = str(root / "runs") + args.ledger_path = str(root / "ledger.jsonl") + args.outcomes_path = str(root / "outcomes.jsonl") + args.claims_path = str(root / "claims.jsonl") + args.history_archive_path = str(root / "history.jsonl") + args.gpu_strategy_state_path = str(root / "strategy.json") + client = AutoStrategyClient(available=3) + + summary = run_submission( + args, + now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc), + hf_discovery=FakeDiscovery(4), # type: ignore[arg-type] + modelhub_client=client, # type: ignore[arg-type] + ) + + state = read_json(root / "strategy.json") + self.assertEqual(3, summary["submittedCount"]) + self.assertEqual(1, summary["duplicateCount"]) + self.assertEqual(3, state["acceptedSinceRefresh"]) + self.assertEqual(3, sum(state["acceptedByCategory"].values())) + def test_duplicate_batch_is_replaced_until_available_slots_are_filled(self) -> None: with tempfile.TemporaryDirectory() as temporary_dir: root = Path(temporary_dir) @@ -231,6 +315,45 @@ class ClientPoolConcurrencyTests(unittest.TestCase): pool.search_by_model_id("owner/model") self.assertEqual(2, client.search_calls) + def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None: + client = DynamicCapacityClient(active=2, limit=3) + pool = ModelHubClientPool( + [client], # type: ignore[arg-type] + active_task_cap=2, + active_counts_ttl=60, + capacity_probe_interval_cycles=3, + instance_id="capacity-growth-test", + ) + + pool.configure_capacity_probe(2) + self.assertEqual(0, pool.available_submit_slots()) + pool.configure_capacity_probe(3) + self.assertEqual(1, pool.available_submit_slots()) + pool.add_task({"model": "probe"}) + + self.assertEqual([3], pool.account_capacity_limits()) + self.assertEqual([3], pool.active_task_counts()) + self.assertEqual(0, pool.available_submit_slots()) + + def test_rejected_capacity_probe_keeps_known_limit_and_enters_cooldown(self) -> None: + client = DynamicCapacityClient(active=2, limit=2) + pool = ModelHubClientPool( + [client], # type: ignore[arg-type] + active_task_cap=2, + active_counts_ttl=60, + capacity_probe_interval_cycles=3, + capacity_probe_cooldown_cycles=3, + instance_id="capacity-rejection-test", + ) + + pool.configure_capacity_probe(3) + self.assertEqual(1, pool.available_submit_slots()) + with self.assertRaises(ModelHubAPIError): + pool.add_task({"model": "probe"}) + + self.assertEqual([2], pool.account_capacity_limits()) + self.assertEqual(0, pool.available_submit_slots()) + class ProcessCoordinationTests(unittest.TestCase): def test_duplicate_claims_are_retained_so_next_batch_moves_forward(self) -> None: diff --git a/tests/test_gpu_strategy.py b/tests/test_gpu_strategy.py new file mode 100644 index 00000000..185c8170 --- /dev/null +++ b/tests/test_gpu_strategy.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import sys +import tempfile +import unittest +from collections import Counter +from datetime import datetime, timedelta, timezone +from pathlib import Path + + +PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api" +if str(PACKAGE_DIR) in sys.path: + sys.path.remove(str(PACKAGE_DIR)) +sys.path.insert(0, str(PACKAGE_DIR)) + +from gpu_strategy import ( # noqa: E402 + ALL_SUPPORTED, + LONG_TERM, + RECENT, + GPUStrategyManager, + build_strategy_snapshot, + choose_next_category, +) + + +def make_history(gpu: str, success: int, failure: int, *, start: datetime) -> list[dict]: + records: list[dict] = [] + for index in range(success + failure): + passed = index < success + records.append( + { + "taskId": f"{gpu}-{index}", + "gpuType": gpu, + "status": "success", + "verifyResult": 1 if passed else -1, + "updateTime": (start + timedelta(seconds=index)).isoformat(), + } + ) + return records + + +class HistoryClient: + def __init__(self, tasks: list[dict]) -> None: + self.tasks = tasks + self.list_calls = 0 + + def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003 + self.list_calls += 1 + return list(self.tasks) + + +class GPUStrategyTests(unittest.TestCase): + def test_long_term_ranking_rejects_tiny_high_rate_sample(self) -> None: + start = datetime(2026, 1, 1, tzinfo=timezone.utc) + supported = ["tiny", "best", "second", "third", "bad"] + tasks = [ + *make_history("tiny", 8, 0, start=start), + *make_history("best", 40, 160, start=start), + *make_history("second", 30, 270, start=start), + *make_history("third", 45, 455, start=start), + *make_history("bad", 5, 495, start=start), + ] + + snapshot = build_strategy_snapshot( + tasks, + supported_gpus=supported, + long_term_min_samples=100, + ) + + self.assertEqual(["best", "second", "third"], snapshot["longTermGpus"]) + self.assertNotIn("tiny", snapshot["longTermGpus"]) + + def test_weighted_category_planner_is_exact_over_200_accepts(self) -> None: + counts = {LONG_TERM: 0, ALL_SUPPORTED: 0, RECENT: 0} + for _ in range(200): + category = choose_next_category(counts) + counts[category] += 1 + + self.assertEqual({LONG_TERM: 100, ALL_SUPPORTED: 60, RECENT: 40}, counts) + + def test_candidate_order_uses_50_30_20_for_first_strategy_generation(self) -> None: + supported = ["gpu-a", "gpu-b", "gpu-c", "gpu-d"] + snapshot = build_strategy_snapshot([], supported_gpus=supported) + manager = GPUStrategyManager("unused.json") + manager.state = snapshot + candidates = [ + {"repoId": f"owner/model-{model}", "targetGpu": gpu} + for model in range(250) + for gpu in supported + ] + + ordered = manager.order_candidates(candidates) + categories = Counter(candidate["strategyCategory"] for candidate in ordered[:200]) + + self.assertEqual(100, categories[LONG_TERM]) + self.assertEqual(60, categories[ALL_SUPPORTED]) + self.assertEqual(40, categories[RECENT]) + + def test_refresh_happens_only_after_200_accepted_submissions(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + path = Path(temporary_dir) / "strategy.json" + client = HistoryClient([]) + manager = GPUStrategyManager(path, refresh_submissions=200) + manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"]) + self.assertEqual(1, client.list_calls) + + manager.record_accepted([{"strategyCategory": LONG_TERM}] * 199) + self.assertEqual(1, manager.submissions_until_refresh) + manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"]) + self.assertEqual(1, client.list_calls) + + manager.record_accepted([{"strategyCategory": RECENT}]) + self.assertEqual(0, manager.submissions_until_refresh) + manager.prepare(client, supported_gpus=["gpu-a", "gpu-b", "gpu-c"]) + self.assertEqual(2, client.list_calls) + self.assertEqual(1, manager.state["generation"]) + self.assertEqual(200, manager.state["acceptedTotal"]) + self.assertEqual(0, manager.state["acceptedSinceRefresh"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modelscope_discovery.py b/tests/test_modelscope_discovery.py index 31953e43..e4ec57a9 100644 --- a/tests/test_modelscope_discovery.py +++ b/tests/test_modelscope_discovery.py @@ -11,6 +11,7 @@ if str(PACKAGE_DIR) in sys.path: sys.path.insert(0, str(PACKAGE_DIR)) from hf_discovery import HuggingFaceDiscovery # noqa: E402 +from http_json import HttpJsonError # noqa: E402 class RecordingHttpClient: @@ -22,6 +23,30 @@ class RecordingHttpClient: return {"success": True, "data": {"models": []}} +class RateLimitedPageClient: + def __init__(self) -> None: + self.page_calls: list[int] = [] + self.page_two_attempts = 0 + + def request_json(self, _method: str, _path: str, *, query: dict) -> dict: + page = int(query["page_number"]) + self.page_calls.append(page) + if page == 2: + self.page_two_attempts += 1 + if self.page_two_attempts == 1: + raise HttpJsonError("rate limited", status_code=429) + start = (page - 1) * 50 + items = [ + { + "id": f"owner/model-{index}", + "downloads": 100, + "last_modified": "2026-01-01T00:00:00Z", + } + for index in range(start, start + 50) + ] + return {"success": True, "data": {"models": items}} + + class ModelScopeDiscoveryTests(unittest.TestCase): def test_openapi_page_size_never_exceeds_platform_limit(self) -> None: http_client = RecordingHttpClient() @@ -36,6 +61,29 @@ class ModelScopeDiscoveryTests(unittest.TestCase): self.assertEqual([], models) self.assertEqual(50, http_client.queries[0]["page_size"]) + def test_rate_limited_page_keeps_previous_pages_and_resumes_from_cache(self) -> None: + http_client = RateLimitedPageClient() + discovery = HuggingFaceDiscovery( + http_client=http_client, # type: ignore[arg-type] + page_interval_seconds=0, + page_cache_ttl_seconds=60, + ) + + partial = discovery.list_recent_models( + pipeline_tags=["text-generation"], + limit=100, + min_downloads=0, + ) + resumed = discovery.list_recent_models( + pipeline_tags=["text-generation"], + limit=100, + min_downloads=0, + ) + + self.assertEqual(50, len(partial)) + self.assertEqual(100, len(resumed)) + self.assertEqual([1, 2, 2], http_client.page_calls) + if __name__ == "__main__": unittest.main()