From 4260793c03bbb9f5689e699a80b9d9d7452f34ad Mon Sep 17 00:00:00 2001 From: CoolBoy Date: Sat, 15 Aug 2026 19:24:26 +0800 Subject: [PATCH] feat: add durable success-first modelhub agent --- .env | 5 +- README.md | 78 +- main.py | 23 + modelhub_submmit_api/config_optimizer.py | 120 ++++ modelhub_submmit_api/daily_runner.py | 56 +- modelhub_submmit_api/hf_discovery.py | 129 +++- modelhub_submmit_api/main.py | 256 ++++++- modelhub_submmit_api/market_intelligence.py | 132 ++++ modelhub_submmit_api/modelhub_client.py | 47 ++ modelhub_submmit_api/models.py | 21 +- modelhub_submmit_api/official_capabilities.py | 176 +++++ modelhub_submmit_api/outcome_tracker.py | 22 + modelhub_submmit_api/poll_runner.py | 139 ++++ modelhub_submmit_api/queue_cleanup.py | 115 ++- modelhub_submmit_api/routing_engine.py | 192 +++++ modelhub_submmit_api/setup.sh | 4 +- modelhub_submmit_api/state_sync.py | 674 ++++++++++++++++++ modelhub_submmit_api/task_registry.py | 56 +- modelhub_submmit_api/version.py | 2 +- requirements.txt | 2 +- tests/test_agent_entrypoint.py | 11 + tests/test_super_agent.py | 246 +++++++ 22 files changed, 2420 insertions(+), 86 deletions(-) create mode 100644 modelhub_submmit_api/config_optimizer.py create mode 100644 modelhub_submmit_api/official_capabilities.py create mode 100644 modelhub_submmit_api/routing_engine.py create mode 100644 modelhub_submmit_api/state_sync.py create mode 100644 tests/test_super_agent.py diff --git a/.env b/.env index cecf19ab..a2a66d5c 100644 --- a/.env +++ b/.env @@ -1,4 +1,7 @@ modelhub = 8726eab3d95922413fc9dfe9dec535d3b6a55cbd xc_token = a14776f6e7ad4c04a1710260613c294c modelscope = ms-b4918c83-7eb3-4034-8635-f154938ed3f0 -dashscope = sk-ws-H.EEMMMLP.i9CD.MEYCIQCXgmgQJ8LfF1m-oBT4ogqc6eD8ahI1BokpJUjD4mlQqAIhAPNH_jFbhJS7fZaufHWCCKCF4Ty8HsP3JmMhvyhSpYhm \ No newline at end of file +dashscope = sk-ws-H.EEMMMLP.i9CD.MEYCIQCXgmgQJ8LfF1m-oBT4ogqc6eD8ahI1BokpJUjD4mlQqAIhAPNH_jFbhJS7fZaufHWCCKCF4Ty8HsP3JmMhvyhSpYhm +modelhub_user_name = CoolBoy +modelhub_user_email = 2269097679@qq.com +modelhub_user_password = Woshixzy666! \ No newline at end of file diff --git a/README.md b/README.md index 075f0aba..33d95525 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ This repository is packaged for the ModelHub XC agent platform. - Root-level `Dockerfile` - Listens on port `8080` -- Exposes `GET /health` +- Exposes `GET /health` for liveness and `GET /ready` for submission readiness - Handles `SIGTERM` - Reads platform-provided `STRATEGY_ID` and attaches it to task submissions as `strategyId` @@ -44,6 +44,8 @@ Optional tuning: - `MODELHUB_AGENT_GPUS` - `MODELHUB_AGENT_EXTRA_ARGS` - `MODELHUB_GPU_STRATEGY_STATE_PATH` default `.modelhub_state/gpu_strategy.json` +- `MODELHUB_ROUTING_STATE_PATH` default `.modelhub_state/routing_intelligence.json` +- `MODELHUB_OFFICIAL_CAPABILITIES_PATH` default `.modelhub_state/official_capabilities.json` - `MODELHUB_MARKET_INTELLIGENCE_PATH` default `.modelhub_state/market_intelligence.json` - `MODELHUB_MARKET_QUEUE_REFRESH_SECONDS` default `600` - `MODELHUB_MARKET_FRAMEWORK_REFRESH_SECONDS` default `21600` @@ -64,19 +66,31 @@ Optional tuning: - `MODELHUB_ARCHITECTURE_BOOTSTRAP_MAX_LOGS` default `0` (unlimited) - `MODELHUB_RECENT_MODEL_RESERVE_SLOTS` default `5` per account - `MODELHUB_RECENT_MODEL_DAYS` default `7` +- `MODELHUB_STATE_SYNC_REMOTE` default `https://dev.modelhub.org.cn/CoolBoy/submmit.git` +- `MODELHUB_STATE_SYNC_BRANCH` default `agent-state` +- `MODELHUB_STATE_SYNC_BATCH_SIZE` default `20` + +Git state synchronization reads `modelhub_user_name`, `modelhub_user_email`, and +`modelhub_user_password` from the tracked private-deployment `.env`. Uppercase +`MODELHUB_GIT_USERNAME`, `MODELHUB_GIT_EMAIL`, and `MODELHUB_GIT_PASSWORD` +override them when supplied by the container. ## Adaptive GPU Strategy -When no explicit GPU override is supplied, the worker uses a success-first 70/30 -strategy generation with no self-funded exploration: +When no explicit GPU override is supplied, the worker uses one deterministic +success-first scorer. The former 70/30 long-term/recent quota no longer controls +traffic and there is no self-funded random exploration. Evidence falls back from +the most specific qualifying cohort to broader community evidence: -- 70%: the three long-term GPUs with the best Wilson lower confidence score and at least 100 terminal samples -- 30%: the top recent GPUs among the latest 1,000 terminal tasks -- 0%: unvetted/all-GPU exploration; community-wide results provide the exploration signal +- local task+GPU+framework+architecture+quantization+load-size bucket (20 samples) +- local task+GPU+framework+architecture+quantization (40 samples) +- local task+GPU+framework (100 samples) +- official community task+GPU+framework aggregate +- a conservative global prior when the platform has no observations -The 70/30 category ratio remains exact across accepted tasks. Inside each -category, weighted fair scheduling combines the category's historical rank with -live public market data: +Each rate is ranked by its Wilson lower confidence bound. Routes within five +percentage points of the best success estimate use expected successful +completions per hour as the tie-breaker: - recent public success quality, scored with a strongly weighted Wilson lower confidence bound - estimated backlog hours (`waiting / recent completions per hour`) as a bounded tie-breaker @@ -86,7 +100,9 @@ live public market data: This optimizes expected successful completions rather than blindly selecting the smallest queue. Queue/throughput data is refreshed every 10 minutes and persisted in `.modelhub_state/market_intelligence.json`. A failed refresh keeps the last good -snapshot, uses a retry backoff, and never blocks normal submissions. +snapshot and uses retry backoff. Critical official capability data is different: +without a usable GPU catalog, task route, framework catalog, build config, or +model/GPU uniqueness response, the entire submission cycle pauses fail-closed. For each compatible model/GPU pair, the worker also ranks the GPU's supported frameworks using ModelHub's public aggregate `modelCount` and `successCount` data, @@ -100,12 +116,12 @@ at least 10%. A new framework is eligible only after the authenticated official build-config endpoint returns a complete config that passes local structure, placeholder, framework-name, and GPU-parallelism validation. Valid official configs are cached and refreshed with -the framework snapshot; local templates remain the fail-safe fallback. +the framework snapshot; production routing never substitutes a hand template +when both the live config and its bounded last-good cache are unavailable. -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 140/60 targets. The active snapshot and -progress are stored in `.modelhub_state/gpu_strategy.json`. +Only platform-accepted tasks count. Routing is recalculated from the current +evidence on every scan; accepted route counters and lifetime sufficient statistics +are stored in `.modelhub_state/routing_intelligence.json`. Five consecutive local failures open a 12-hour GPU/framework circuit breaker. A sub-20% success rate over the latest 20 terminal tasks opens a 6-hour breaker. @@ -116,6 +132,28 @@ idle card does not permanently poison otherwise successful evidence. Candidate shortages expand the model search window; they never unlock an unvetted GPU or framework. +## Durable State Branch + +The hosted worker synchronizes an allowlisted runtime snapshot to the orphan +`agent-state` branch of this repository. Before each batch of at most 20 API +submissions it pushes write-ahead intents containing model, GPU, task, framework, +safe parameter vector, and a config fingerprint. It pushes task IDs and results +after the batch. If either push fails, no further submission begins until the +same snapshot can be synchronized. + +At startup the worker verifies the manifest and file checksums, restores local +capacity, outcomes, routing evidence, architecture rules, exclusions, intents, +and active-task context, then reconciles every account against the platform. +Pending intents are held for two hours before being released as unconfirmed. +The branch keeps 30 days of structured events plus lifetime aggregate counters; +raw stdout, credentials, request headers, downloaded archives, and full configs +are never copied. Git authentication uses a temporary `GIT_ASKPASS` helper, so +the password is absent from command arguments, remotes, commits, and logs. + +`GET /health` reports process liveness. `GET /ready` returns HTTP 200 only after +state recovery, state synchronization, and critical official capability checks +are usable; otherwise it returns HTTP 503 without forcing a liveness restart. + Failed-task archives are also classified conservatively. When ModelHub explicitly says that the selected framework does not support the model or architecture, the runner learns an exact GPU + framework + task type + architecture block from the @@ -324,12 +362,18 @@ counts fail closed for older candidates, and no startup or periodic cleanup can cancel a task by date. It also keeps unclassified/ambiguous historical failures neutral in GPU/framework success feedback while preserving deterministic OOM and architecture cleanup. +Version `2026.08.15.1` replaces the 70/30 quota with hierarchical success-first +routing, dynamically gates models through the official GPU/task/framework/config +APIs, enriches ModelScope metadata and model lineage, learns only proven safe +config vectors, and adds crash-safe write-ahead state synchronization to the +`agent-state` branch. It also exposes `/ready` and extends deterministic cleanup +to officially removed waiting GPU/framework routes. ## Deploy Create a tag and submit the repository URL plus tag in "我的适配智能体". ```bash -git tag agent-v24 -git push origin agent-v24 +git tag -a agent-v25 -m "ModelHub agent 2026.08.15.1" +git push origin main agent-v25 ``` diff --git a/main.py b/main.py index e7260a57..10b5b421 100644 --- a/main.py +++ b/main.py @@ -17,6 +17,7 @@ HOST = "0.0.0.0" PORT = int(os.getenv("PORT", "8080")) ROOT = Path(__file__).resolve().parent WORKER_SCRIPT = ROOT / "modelhub_submmit_api" / "poll_runner.py" +READINESS_PATH = ROOT / ".modelhub_state" / "readiness.json" shutdown_requested = False worker: subprocess.Popen | None = None @@ -53,6 +54,7 @@ def _worker_command() -> list[str]: "--post-cycle-cooldown-seconds", os.getenv("MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS", "2"), "--skip-history-archive", + "--state-sync", ] daily_target = os.getenv("MODELHUB_AGENT_DAILY_TARGET", "").strip() @@ -86,6 +88,14 @@ def _config() -> dict[str, object]: } +def _readiness() -> dict[str, object]: + try: + payload = json.loads(READINESS_PATH.read_text(encoding="utf-8")) + except (FileNotFoundError, OSError, ValueError, TypeError): + return {"ready": False, "reason": "readiness_not_reported"} + return payload if isinstance(payload, dict) else {"ready": False, "reason": "readiness_invalid"} + + class Handler(BaseHTTPRequestHandler): def do_GET(self) -> None: if self.path == "/health": @@ -98,6 +108,19 @@ class Handler(BaseHTTPRequestHandler): self._send_json({"status": "ok", "config": _config()}) return + if self.path == "/ready": + readiness = _readiness() + status = 200 if readiness.get("ready") is True else 503 + self._send_json( + { + "status": "ready" if status == 200 else "not_ready", + "readiness": readiness, + "config": _config(), + }, + status=status, + ) + return + if self.path == "/": self._send_json({"name": "modelhub-submmit-agent", "status": "running", "config": _config()}) return diff --git a/modelhub_submmit_api/config_optimizer.py b/modelhub_submmit_api/config_optimizer.py new file mode 100644 index 00000000..d4c08c83 --- /dev/null +++ b/modelhub_submmit_api/config_optimizer.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import re +from collections import defaultdict +from pathlib import Path +from typing import Any + +from common import read_jsonl +from gpu_strategy import wilson_lower_bound + + +class SafeConfigOptimizer: + """Apply only repeatedly successful, sanitized parameter vectors.""" + + def __init__(self, *, intents_path: Path | str, outcomes_path: Path | str) -> None: + self.intents_path = Path(intents_path) + self.outcomes_path = Path(outcomes_path) + self._qualified = self._learn() + + def _learn(self) -> dict[tuple[str, str, str], list[dict[str, Any]]]: + outcomes = read_jsonl(self.outcomes_path) + by_task = { + str(item.get("taskId")): item + for item in outcomes + if item.get("taskId") is not None and item.get("outcome") in {"success", "failed"} + } + groups: dict[tuple[str, str, str, str], dict[str, Any]] = defaultdict( + lambda: {"success": 0, "failure": 0, "models": set(), "vector": {}} + ) + for intent in read_jsonl(self.intents_path): + task_id = str(intent.get("taskId") or "") + outcome = by_task.get(task_id) + vector = intent.get("safeConfigVector") + if outcome is None or not isinstance(vector, dict) or not vector: + continue + route = ( + str(intent.get("taskType") or ""), + str(intent.get("targetGpu") or ""), + str(intent.get("framework") or ""), + str(intent.get("configFingerprint") or ""), + ) + group = groups[route] + group["vector"] = dict(vector) + group["models"].add(str(intent.get("repoId") or "")) + if outcome.get("outcome") == "success": + group["success"] += 1 + elif not outcome.get("platformFailure") and not outcome.get("policyCancelled"): + group["failure"] += 1 + + qualified: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list) + for (task_type, gpu, framework, fingerprint), group in groups.items(): + success = int(group["success"]) + failure = int(group["failure"]) + if success < 5 or failure > 0 or len(group["models"]) < 2: + continue + qualified[(task_type, gpu, framework)].append( + { + "fingerprint": fingerprint, + "success": success, + "total": success + failure, + "lowerBound": wilson_lower_bound(success, success + failure), + "vector": group["vector"], + } + ) + return dict(qualified) + + def optimize( + self, + *, + task_type: str, + target_gpu: str, + framework: str, + official_config: str, + official_lower_bound: float, + ) -> tuple[str, dict[str, Any]]: + choices = self._qualified.get((task_type, target_gpu, framework)) or [] + choices = [ + item + for item in choices + if float(item.get("lowerBound") or 0.0) >= float(official_lower_bound) + 0.05 + ] + if not choices: + return official_config, {"source": "official", "applied": False} + choice = max(choices, key=lambda item: (float(item["lowerBound"]), int(item["success"]))) + patched = official_config + vector = choice["vector"] + substitutions = { + "gpuNum": (r"(\bgpu_num\s*:\s*['\"]?)\d+", r"\g<1>{}"), + "tensorParallel": ( + r"((?:--tensor-parallel-size|-tp)\s*[, ]?\s*['\"]?)\d+", + r"\g<1>{}", + ), + "maxModelLen": ( + r"((?:--max-model-len|max_model_len|max_seq_len)\s*[: ,]+\s*['\"]?)\d+", + r"\g<1>{}", + ), + "gpuMemoryUtilization": ( + r"((?:--gpu-memory-utilization|gpu_memory_utilization)\s*[: ,]+\s*['\"]?)[0-9.]+", + r"\g<1>{}", + ), + } + for key, (pattern, replacement) in substitutions.items(): + if key in vector: + patched = re.sub(pattern, replacement.format(vector[key]), patched, flags=re.IGNORECASE) + for key, option in (("dtype", "dtype"), ("quantization", "quantization"), ("loadFormat", "load[_-]format")): + if key in vector: + patched = re.sub( + rf"((?:--{option}|{option})\s*[: ,]+\s*['\"]?)[A-Za-z0-9_-]+", + rf"\g<1>{vector[key]}", + patched, + flags=re.IGNORECASE, + ) + return patched, { + "source": "official_plus_learned_patch", + "applied": patched != official_config, + "evidenceSuccess": choice["success"], + "evidenceLowerBound": choice["lowerBound"], + "fingerprint": choice["fingerprint"], + } + diff --git a/modelhub_submmit_api/daily_runner.py b/modelhub_submmit_api/daily_runner.py index e46cf298..a0434ffb 100644 --- a/modelhub_submmit_api/daily_runner.py +++ b/modelhub_submmit_api/daily_runner.py @@ -20,6 +20,7 @@ from market_intelligence import ( ) from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool from outcome_tracker import OutcomeTracker +from official_capabilities import OfficialCapabilityUnavailable from runner_common import DEFAULT_KEY_PATH, ensure_tokens from submission_claims import DEFAULT_CLAIMS_PATH from template_selector import TemplateSelector @@ -220,6 +221,12 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar disable_market_intelligence=getattr(base_args, "disable_market_intelligence", 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)), + routing_state_path=getattr(base_args, "routing_state_path", ".modelhub_state/routing_intelligence.json"), + official_capabilities_path=getattr( + base_args, + "official_capabilities_path", + ".modelhub_state/official_capabilities.json", + ), 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), market_intelligence_state_path=getattr( @@ -279,6 +286,7 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar hf_base_url=base_args.hf_base_url, modelhub_base_url=base_args.modelhub_base_url, modelhub_token=base_args.modelhub_token, + _state_sync_manager=getattr(base_args, "_state_sync_manager", None), ) @@ -337,14 +345,33 @@ def run_daily_batches( f"tasks={','.join(wave.task_types)} gpus={wave.gpus or 'auto'} " f"limit={wave.limit} since_hours={wave.since_hours}" ) - summary = run_fn( - wave_args, - now=utc_now(), - hf_discovery=hf_discovery, - modelhub_client=modelhub_client, - template_selector=template_selector, - outcome_tracker=outcome_tracker, - ) + try: + summary = run_fn( + wave_args, + now=utc_now(), + hf_discovery=hf_discovery, + modelhub_client=modelhub_client, + template_selector=template_selector, + outcome_tracker=outcome_tracker, + ) + except OfficialCapabilityUnavailable as exc: + reason = str(exc) or "critical_official_signal_unavailable" + log(f"[cycle] paused reason=critical_official_signal_unavailable detail={reason}") + summary = { + "candidateCount": 0, + "plannedSubmitCount": 0, + "submittedCount": 0, + "skippedCount": 0, + "duplicateCount": 0, + "modelGpuUniquenessRejectedCount": 0, + "failedCount": 0, + "skipReasonCounts": {}, + "remainingDailyQuotaBeforeRun": None, + "platformAvailableSlotsBeforeRun": None, + "paused": True, + "pauseReason": "critical_official_signal_unavailable", + "pauseDetail": reason, + } wave_result = { "round": round_index, "wave": asdict(wave), @@ -382,7 +409,11 @@ def run_daily_batches( stopped_reason=stopped_reason, ) - if base_args.daily_target > 0 and summary["remainingDailyQuotaBeforeRun"] <= 0: + if ( + not summary.get("paused") + and base_args.daily_target > 0 + and summary["remainingDailyQuotaBeforeRun"] <= 0 + ): stopped_reason = "daily_target_already_reached" log(f"[daily] stop={stopped_reason}") return finalize_daily_run( @@ -394,7 +425,12 @@ def run_daily_batches( attempted_waves=attempted_waves, stopped_reason=stopped_reason, ) - if base_args.daily_target > 0 and not base_args.dry_run and summary["plannedSubmitCount"] <= 0: + if ( + not summary.get("paused") + and base_args.daily_target > 0 + and not base_args.dry_run + and summary["plannedSubmitCount"] <= 0 + ): stopped_reason = "daily_target_reached" log(f"[daily] stop={stopped_reason}") return finalize_daily_run( diff --git a/modelhub_submmit_api/hf_discovery.py b/modelhub_submmit_api/hf_discovery.py index 1f0ff3d5..f17e236d 100644 --- a/modelhub_submmit_api/hf_discovery.py +++ b/modelhub_submmit_api/hf_discovery.py @@ -1,8 +1,10 @@ from __future__ import annotations import os +import re import threading import time +from dataclasses import replace from datetime import datetime, timezone from pathlib import PurePosixPath from typing import Any @@ -92,6 +94,9 @@ class HuggingFaceDiscovery: ), ) self._last_model_page_request_at = 0.0 + self._unsupported_task_filters: dict[str, float] = {} + self._model_card_cache: dict[str, dict[str, Any]] = {} + self._model_card_lock = threading.Lock() def list_recent_models( self, @@ -108,10 +113,12 @@ class HuggingFaceDiscovery: deduped: dict[str, HFModelSummary] = {} del read_concurrency + per_tag_limit = limit if len(pipeline_tags) <= 1 else max(10, (limit + len(pipeline_tags) - 1) // len(pipeline_tags)) + for pipeline_tag in pipeline_tags: for model in self._query_recent_models( pipeline_tag=pipeline_tag, - limit=limit, + limit=per_tag_limit, min_downloads=min_downloads, updated_after=updated_after, ): @@ -122,7 +129,7 @@ class HuggingFaceDiscovery: 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 + return models[: max(1, int(limit))] def _query_recent_models( self, @@ -136,9 +143,10 @@ class HuggingFaceDiscovery: page_size = min(max(1, limit), 50) max_items = min(max(1, limit), 3000) task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag) + filter_disabled = self._unsupported_task_filters.get(task_tag, 0.0) > time.monotonic() models: list[HFModelSummary] = [] for page_number in range(1, (max_items + page_size - 1) // page_size + 1): - cache_key = (task_tag, page_number, page_size) + cache_key = ("*" if filter_disabled else 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]) @@ -147,19 +155,34 @@ class HuggingFaceDiscovery: if self._last_model_page_request_at > 0 and elapsed < self._page_interval_seconds: time.sleep(self._page_interval_seconds - elapsed) try: + query = { + "page_number": page_number, + "page_size": page_size, + "sort": "last_modified", + } + if not filter_disabled: + query["filter.task"] = task_tag payload = self.http_client.request_json( "GET", "/models", - query={ - "page_number": page_number, - "page_size": page_size, - "sort": "last_modified", - "filter.task": task_tag, - }, + query=query, ) self._last_model_page_request_at = time.monotonic() except HttpJsonError as exc: self._last_model_page_request_at = time.monotonic() + if exc.status_code == 400 and not filter_disabled: + self._unsupported_task_filters[task_tag] = time.monotonic() + 86_400 + print( + f"[modelscope] task_filter_unsupported task={task_tag} " + "fallback=unfiltered ttl=86400s", + flush=True, + ) + return self._query_recent_models( + pipeline_tag=pipeline_tag, + limit=limit, + min_downloads=min_downloads, + updated_after=updated_after, + ) print( f"[modelscope] list_models_error task={task_tag} page={page_number} " f"partial_models={len(models)} retry_next_cycle=true error={exc}", @@ -204,8 +227,9 @@ class HuggingFaceDiscovery: entries = self.list_repo_tree(model.repo_id) inspection = inspect_repo_tree(model.repo_id, entries) if not inspection.has_root_config: - return inspection + return replace(inspection, published_size_bytes=model.file_size) model_config, config_error = self.get_model_config(model.repo_id) + model_card_metadata = self.get_model_card_metadata(model.repo_id) return ModelInspection( repo_id=inspection.repo_id, file_paths=inspection.file_paths, @@ -216,8 +240,31 @@ class HuggingFaceDiscovery: onnx_files=inspection.onnx_files, model_config=model_config, config_fetch_error=config_error, + model_card_metadata=model_card_metadata, + published_size_bytes=model.file_size, ) + def get_model_card_metadata(self, repo_id: str) -> dict[str, Any]: + with self._model_card_lock: + cached = self._model_card_cache.get(repo_id) + if cached is not None: + return dict(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}") + data = payload.get("Data") if isinstance(payload, dict) else None + if not isinstance(data, dict) and isinstance(payload, dict): + data = payload.get("data") + readme = "" + if isinstance(data, dict): + readme = str(data.get("ReadMe") or data.get("readme") or data.get("README") or "") + result = parse_model_card_front_matter(readme) + except Exception: + result = {} + with self._model_card_lock: + self._model_card_cache[repo_id] = dict(result) + return result + def get_model_config(self, repo_id: str) -> tuple[dict[str, Any], str | None]: with self._model_config_lock: cached = self._model_config_cache.get(repo_id) @@ -356,9 +403,71 @@ class HuggingFaceDiscovery: last_modified=last_modified, pipeline_tag=pipeline_tag, created_at=parse_datetime(item.get("created_at") or item.get("CreatedAt")), + params=_optional_int(item.get("params") or item.get("Params") or item.get("parameter_count")), + file_size=_optional_int(item.get("file_size") or item.get("FileSize") or item.get("size")), + tags=_string_tuple(item.get("tags") or item.get("Tags")), + tasks=_string_tuple(item.get("tasks") or item.get("Tasks")), + license=str(item.get("license") or item.get("License") or "").strip() or None, + gated=bool(item.get("gated") or item.get("Gated")), + private=bool(item.get("private") or item.get("Private")), + likes=int(item.get("likes") or item.get("Likes") or 0), ) +def _optional_int(value: Any) -> int | None: + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed >= 0 else None + + +def _string_tuple(value: Any) -> tuple[str, ...]: + if isinstance(value, str): + values = [part.strip() for part in value.split(",")] + elif isinstance(value, (list, tuple, set)): + values = [str(part).strip() for part in value] + else: + values = [] + return tuple(dict.fromkeys(part for part in values if part)) + + +def parse_model_card_front_matter(readme: str) -> dict[str, Any]: + text = str(readme or "")[:65_536] + if not text.startswith("---"): + return {} + match = re.match(r"^---\s*\n(.*?)\n---(?:\s*\n|$)", text, flags=re.DOTALL) + if match is None: + return {} + front_matter = match.group(1) + try: + import yaml # type: ignore + + parsed = yaml.safe_load(front_matter) + data = parsed if isinstance(parsed, dict) else {} + except (ImportError, ValueError, TypeError): + data = {} + current_list: str | None = None + for raw_line in front_matter.splitlines(): + if re.match(r"^\s+-\s+", raw_line) and current_list: + value = re.sub(r"^\s+-\s+", "", raw_line).strip().strip('"\'') + data.setdefault(current_list, []).append(value) + continue + if ":" not in raw_line or raw_line[:1].isspace(): + continue + key, value = raw_line.split(":", 1) + key = key.strip() + value = value.strip().strip('"\'') + if not value: + data[key] = [] + current_list = key + else: + data[key] = value + current_list = None + allowed = ("base_model", "base_model_relation", "frameworks", "tasks", "new_version") + return {key: data[key] for key in allowed if key in data} + + def inspect_repo_tree(repo_id: str, entries: list[dict[str, Any]]) -> ModelInspection: file_paths: list[str] = [] file_sizes: dict[str, int] = {} diff --git a/modelhub_submmit_api/main.py b/modelhub_submmit_api/main.py index 4ef6129a..3d59fe9c 100644 --- a/modelhub_submmit_api/main.py +++ b/modelhub_submmit_api/main.py @@ -1,6 +1,7 @@ from __future__ import annotations import argparse +import hashlib import os from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed @@ -10,7 +11,8 @@ from typing import Any from common import parse_datetime, runtime_instance_id, utc_now, write_json, write_jsonl from candidate_preflight import CandidatePreflightAdvisor -from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH, GPUStrategyManager +from config_optimizer import SafeConfigOptimizer +from gpu_strategy import DEFAULT_GPU_STRATEGY_PATH from hf_discovery import HuggingFaceDiscovery from history_stats import ( append_ledger_entry, @@ -38,10 +40,23 @@ from modelhub_client import ( is_model_uniqueness_error, ) from models import CandidateModel, HFModelSummary, ModelInspection +from official_capabilities import ( + DEFAULT_OFFICIAL_CAPABILITIES_PATH, + OfficialCapabilityRegistry, + OfficialCapabilityUnavailable, +) from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker +from routing_engine import DEFAULT_ROUTING_STATE_PATH, SuccessFirstRoutingEngine from submission_claims import DEFAULT_CLAIMS_PATH, SubmissionClaimStore, candidate_key, diversify_candidates from submission_exclusions import DEFAULT_SUBMISSION_EXCLUSIONS_PATH, SubmissionExclusionStore -from task_registry import TASK_SPEC_BY_TYPE, all_task_types, compatible_frameworks_for_task, pipeline_tags_for_task_types, task_specs_for_model +from task_registry import ( + TASK_SPEC_BY_TYPE, + all_task_types, + compatible_frameworks_for_task, + pipeline_tags_for_task_types, + register_dynamic_task_route, + task_specs_for_model, +) from template_selector import TemplateSelector @@ -212,6 +227,19 @@ def build_parser() -> argparse.ArgumentParser: default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)), help=argparse.SUPPRESS, ) + parser.add_argument( + "--routing-state-path", + default=os.getenv("MODELHUB_ROUTING_STATE_PATH", str(DEFAULT_ROUTING_STATE_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--official-capabilities-path", + default=os.getenv( + "MODELHUB_OFFICIAL_CAPABILITIES_PATH", + str(DEFAULT_OFFICIAL_CAPABILITIES_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( @@ -311,6 +339,8 @@ def choose_candidate_for_gpu( target_gpu: str, market_intelligence: MarketIntelligenceManager | None = None, preflight_advisor: CandidatePreflightAdvisor | None = None, + require_official_config: bool = False, + config_optimizer: SafeConfigOptimizer | None = None, ) -> CandidateModel | None: candidate, _reason = choose_candidate_for_gpu_detailed( model=model, @@ -320,6 +350,8 @@ def choose_candidate_for_gpu( target_gpu=target_gpu, market_intelligence=market_intelligence, preflight_advisor=preflight_advisor, + require_official_config=require_official_config, + config_optimizer=config_optimizer, ) return candidate @@ -333,6 +365,8 @@ def choose_candidate_for_gpu_detailed( target_gpu: str, market_intelligence: MarketIntelligenceManager | None = None, preflight_advisor: CandidatePreflightAdvisor | None = None, + require_official_config: bool = False, + config_optimizer: SafeConfigOptimizer | None = None, ) -> tuple[CandidateModel | None, str | None]: last_rejection_reason: str | None = None for task_type in task_types: @@ -371,6 +405,11 @@ def choose_candidate_for_gpu_detailed( ) if config_params is not None: template_id = f"modelhub-live-{task_type}-{framework}-{target_gpu}".lower().replace("_", "-") + elif market_intelligence is not None and require_official_config: + # A live-discovered route without a valid official config is a + # critical capability gap, not permission to use a stale hand template. + last_rejection_reason = "official_build_config_unavailable" + continue else: try: template = template_selector.select_template(task_type, framework, target_gpu) @@ -388,6 +427,19 @@ def choose_candidate_for_gpu_detailed( ) score = float(framework_metadata.get("frameworkCombinedScore") or 0.0) warnings: list[str] = [] + config_optimization: dict[str, Any] = {"source": "official", "applied": False} + if config_optimizer is not None and config_params is not None: + config_params, config_optimization = config_optimizer.optimize( + task_type=task_type, + target_gpu=target_gpu, + framework=framework, + official_config=config_params, + official_lower_bound=float( + framework_metadata.get("frameworkMarketWilsonLowerBound") or 0.0 + ), + ) + if config_optimization.get("applied"): + warnings.append("validated_success_config_patch_applied") if bool(framework_metadata.get("frameworkEvidenceQualified", False)): warnings.append("framework_selected_from_success_evidence") if config_params is not None and bool(framework_metadata.get("frameworkOfficialConfigValid", False)): @@ -407,6 +459,20 @@ def choose_candidate_for_gpu_detailed( config_params = assessment.config_params warnings.extend(assessment.warnings) preflight_metadata = assessment.metadata + preflight_metadata = { + **preflight_metadata, + "configOptimization": config_optimization, + "modelscopeParams": model.params, + "modelscopeFileSize": model.file_size, + "modelscopeTasks": list(model.tasks), + "modelscopeTags": list(model.tags), + "modelscopeLicense": model.license, + "modelCard": dict(inspection.model_card_metadata), + } + repository_size = inspection.repository_size_bytes + if model.file_size and repository_size and model.file_size > repository_size: + warnings.append("modelscope_published_size_used_as_conservative_upper_bound") + preflight_metadata["conservativeRepositorySizeBytes"] = model.file_size spec = TASK_SPEC_BY_TYPE[task_type] return CandidateModel( repo_id=model.repo_id, @@ -471,9 +537,15 @@ def process_model_for_candidates( submission_exclusion_store: SubmissionExclusionStore | None = None, market_intelligence: MarketIntelligenceManager | None = None, preflight_advisor: CandidatePreflightAdvisor | None = None, + official_registry: OfficialCapabilityRegistry | None = None, + config_optimizer: SafeConfigOptimizer | None = None, + allow_dynamic_tasks: bool = False, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]: + if model.private or model.gated: + reason = "modelscope_private" if model.private else "modelscope_gated" + return [], [{"repoId": model.repo_id, "reason": reason}], [] specs = [spec for spec in task_specs_for_model(model) if spec.task_type in allowed_task_types] - if not specs: + if not specs and official_registry is None: return [], [{"repoId": model.repo_id, "reason": f"unsupported_pipeline_tag:{model.pipeline_tag or 'unknown'}"}], [] try: @@ -494,7 +566,8 @@ def process_model_for_candidates( allowed_task_types_set = [spec.task_type for spec in specs] pending_task_types_by_gpu: list[tuple[str, list[str]]] = [] - for target_gpu in target_gpus: + route_target_gpus = target_gpus[:6] if official_registry is not None else target_gpus + for target_gpu in route_target_gpus: if submission_exclusion_store is not None and submission_exclusion_store.is_blocked(model.repo_id, target_gpu): skipped.append( {"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "model_gpu_uniqueness_blocklist"} @@ -507,11 +580,34 @@ def process_model_for_candidates( skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "local_failure_cooldown_24h"}) continue - compatible_task_types = [ - task_type - for task_type in allowed_task_types_set - if template_selector.supported_frameworks_for_auto(task_type, target_gpu) - ] + exact_task_types = list(allowed_task_types_set) + if official_registry is not None: + exact_task_types = official_registry.task_types_for( + modelhub_client, + model_address=model.model_address, + model_last_modified=model.last_modified.isoformat() if model.last_modified else None, + gpu=target_gpu, + ) + route_task_types = ( + list(exact_task_types) + if allow_dynamic_tasks + else [task_type for task_type in allowed_task_types_set if task_type in exact_task_types] + ) + compatible_task_types = [] + for task_type in route_task_types: + register_dynamic_task_route(task_type, model.pipeline_tag or task_type) + if market_intelligence is not None and not market_intelligence.has_framework_route(task_type, target_gpu): + try: + market_intelligence.ensure_framework_route(modelhub_client, task_type, target_gpu) + except Exception as exc: + raise OfficialCapabilityUnavailable( + f"framework capability unavailable for {task_type}/{target_gpu}: {type(exc).__name__}: {exc}" + ) from exc + if template_selector.supported_frameworks_for_auto(task_type, target_gpu): + compatible_task_types.append(task_type) + continue + if market_intelligence is not None and market_intelligence.has_framework_route(task_type, target_gpu): + compatible_task_types.append(task_type) if not compatible_task_types: skipped.append({"repoId": model.repo_id, "targetGpu": target_gpu, "reason": "no_compatible_auto_template_or_framework"}) continue @@ -534,6 +630,8 @@ def process_model_for_candidates( target_gpu=target_gpu, market_intelligence=market_intelligence, preflight_advisor=preflight_advisor, + require_official_config=official_registry is not None, + config_optimizer=config_optimizer, ) if best is None: reason = ( @@ -550,6 +648,9 @@ def process_model_for_candidates( if market_intelligence is not None: record.update(market_intelligence.gpu_metadata(target_gpu)) record.update(market_intelligence.framework_metadata(best.task_type, target_gpu, best.framework)) + optimization = (record.get("preflightMetadata") or {}).get("configOptimization") or {} + if optimization.get("applied"): + record["frameworkConfigSource"] = "official_plus_learned_patch" candidates.append(record) return candidates, skipped, failed @@ -640,6 +741,9 @@ def collect_candidates_from_models( read_concurrency: int, market_intelligence: MarketIntelligenceManager | None = None, preflight_advisor: CandidatePreflightAdvisor | None = None, + official_registry: OfficialCapabilityRegistry | None = None, + config_optimizer: SafeConfigOptimizer | None = None, + allow_dynamic_tasks: bool = False, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]], int]: candidates: list[dict[str, Any]] = [] skipped: list[dict[str, Any]] = [] @@ -670,6 +774,9 @@ def collect_candidates_from_models( submission_exclusion_store=submission_exclusion_store, market_intelligence=market_intelligence, preflight_advisor=preflight_advisor, + official_registry=official_registry, + config_optimizer=config_optimizer, + allow_dynamic_tasks=allow_dynamic_tasks, ): index for index, model in enumerate(chunk) } @@ -679,6 +786,8 @@ def collect_candidates_from_models( model = chunk[index] try: ordered_results[index] = future.result() + except OfficialCapabilityUnavailable: + raise except Exception as exc: ordered_results[index] = ([], [], [{"repoId": model.repo_id, "reason": str(exc)}]) @@ -830,10 +939,9 @@ def run_submission( now = now or utc_now() template_selector = template_selector or TemplateSelector() selected_task_types = resolve_task_types(args) - target_gpus = resolve_target_gpus(args, template_selector, selected_task_types) - if not target_gpus: + configured_target_gpus = resolve_target_gpus(args, template_selector, selected_task_types) + if not configured_target_gpus: raise RuntimeError("No auto-submittable GPUs are available for the selected task types") - submission_target_gpus = list(target_gpus) hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=args.hf_base_url) if modelhub_client is None: @@ -848,6 +956,42 @@ def run_submission( recent_model_days=max(1, int(getattr(args, "recent_model_days", 7) or 7)), ) + official_registry: OfficialCapabilityRegistry | None = None + official_summary: dict[str, Any] = {"enabled": False} + official_capable = all( + hasattr(modelhub_client, name) + for name in ("list_machine_info", "list_task_levels", "list_model_task_types") + ) + if official_capable: + official_registry = OfficialCapabilityRegistry( + Path( + getattr( + args, + "official_capabilities_path", + DEFAULT_OFFICIAL_CAPABILITIES_PATH, + ) + ) + ) + official_registry.prepare( + modelhub_client, + fallback_gpus=configured_target_gpus, + task_types=selected_task_types, + now=now, + ) + if not official_registry.ready: + raise OfficialCapabilityUnavailable( + official_registry.pause_reason or "critical_official_signal_unavailable" + ) + requested_gpus = configured_target_gpus if (getattr(args, "gpu", None) or getattr(args, "gpus", None)) else None + target_gpus = official_registry.eligible_gpus(requested_gpus) + official_summary = official_registry.summary() + else: + # Compatibility path for injected test clients and offline dry-runs. + target_gpus = list(configured_target_gpus) + if not target_gpus: + raise OfficialCapabilityUnavailable("official catalog contains no eligible GPU") + submission_target_gpus = list(target_gpus) + if hasattr(modelhub_client, "begin_cycle"): modelhub_client.begin_cycle() @@ -874,7 +1018,11 @@ def run_submission( preflight_summary: dict[str, Any] = ( preflight_advisor.summary() if preflight_advisor is not None else {"enabled": False} ) - strategy_manager: GPUStrategyManager | None = None + config_optimizer = SafeConfigOptimizer( + intents_path=Path(".modelhub_state/recovery_intents.jsonl"), + outcomes_path=Path(args.outcomes_path), + ) + strategy_manager: SuccessFirstRoutingEngine | None = None strategy_summary: dict[str, Any] = {"enabled": False} market_intelligence: MarketIntelligenceManager | None = None market_summary: dict[str, Any] = {"enabled": False} @@ -1042,6 +1190,14 @@ def run_submission( except Exception: market_intelligence.set_local_outcome_stats(None) market_summary = market_intelligence.summary() + framework_state = (market_intelligence.state or {}).get("frameworkStats") or {} + framework_updated_at = parse_datetime(market_summary.get("frameworkUpdatedAt")) + framework_cache_expired = ( + framework_updated_at is None + or now - framework_updated_at > timedelta(hours=24) + ) + if market_summary.get("frameworkError") and (not framework_state or framework_cache_expired): + raise OfficialCapabilityUnavailable("critical official framework catalog is unavailable") submission_target_gpus = market_intelligence.eligible_gpus(target_gpus) market_summary["eligibleGpus"] = submission_target_gpus market_summary["shadowOnlyGpus"] = [gpu for gpu in target_gpus if gpu not in submission_target_gpus] @@ -1052,19 +1208,20 @@ def run_submission( ) 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)), - market_intelligence=market_intelligence, - ) - strategy_history_records = outcome_tracker.get_strategy_history_records() - strategy_manager.prepare( - modelhub_client, - supported_gpus=target_gpus, - now=now, - history_records=strategy_history_records, + routing_path = getattr(args, "routing_state_path", None) + legacy_strategy_path = getattr(args, "gpu_strategy_state_path", None) + if ( + not routing_path + or ( + str(routing_path) == str(DEFAULT_ROUTING_STATE_PATH) + and legacy_strategy_path + and str(legacy_strategy_path) != str(DEFAULT_GPU_STRATEGY_PATH) + ) + ): + routing_path = legacy_strategy_path or DEFAULT_ROUTING_STATE_PATH + strategy_manager = SuccessFirstRoutingEngine( + Path(routing_path), + outcome_stats=outcome_tracker.get_stats_report(), ) strategy_summary = strategy_manager.summary() @@ -1117,6 +1274,11 @@ def run_submission( candidate_goal = max(submission_goal, submission_goal * attempt_multiplier) pipeline_tags = pipeline_tags_for_task_types(selected_task_types) + if official_registry is not None and not getattr(args, "task_types", None): + for task_name in official_registry.state.get("discoveredTaskTypes") or []: + normalized = str(task_name or "").strip() + if normalized and normalized not in pipeline_tags: + pipeline_tags.append(normalized) explicit_scan_cap = max(0, int(getattr(args, "max_scan_models", 0) or 0)) recent_model_days = max(1, int(getattr(args, "recent_model_days", 7) or 7)) recent_model_reserve_slots = max(0, int(getattr(args, "recent_model_reserve_slots", 5) or 0)) @@ -1197,6 +1359,9 @@ def run_submission( read_concurrency=max(1, args.read_concurrency), market_intelligence=market_intelligence, preflight_advisor=preflight_advisor, + official_registry=official_registry, + config_optimizer=config_optimizer, + allow_dynamic_tasks=not bool(getattr(args, "task_types", None)), ) candidates.extend(stage_candidates) skipped.extend(stage_skipped) @@ -1239,6 +1404,8 @@ def run_submission( claim_store: SubmissionClaimStore | None = None attempted_candidates: list[dict[str, Any]] = [] submit_workers = 1 + state_sync = getattr(args, "_state_sync_manager", None) + state_sync_paused = False if args.dry_run: attempted_candidates = diversified_candidates[:target_submit_count] @@ -1270,12 +1437,23 @@ def run_submission( target_submit_count - len(submitted), max_submit_attempts - len(attempted_candidates), ) + if state_sync is not None: + desired_count = min(desired_count, int(state_sync.batch_size)) batch_candidates = claim_store.claim(remaining_candidates, limit=desired_count) if not batch_candidates: break attempted_candidates.extend(batch_candidates) attempted_keys.update(candidate_key(candidate) for candidate in batch_candidates) + state_batch_id: str | None = None + if state_sync is not None: + state_batch_id = state_sync.begin_batch(batch_candidates) + if state_batch_id is None: + claim_store.release(batch_candidates) + state_sync_paused = True + print("[cycle] paused reason=state_sync_unhealthy phase=intent", flush=True) + break + batch_workers = resolve_submit_concurrency( args, modelhub_client=modelhub_client, @@ -1389,7 +1567,13 @@ def run_submission( task_type=candidate["taskType"], task_id=result["taskId"], submit_time=result["submitTime"], - model_profile=candidate.get("preflightMetadata") or {}, + model_profile={ + **dict(candidate.get("preflightMetadata") or {}), + "configFingerprint": hashlib.sha256( + str(candidate.get("configParams") or "").encode("utf-8") + ).hexdigest(), + "configSource": candidate.get("frameworkConfigSource") or "official", + }, ) claim_store.mark_submitted( @@ -1398,6 +1582,17 @@ def run_submission( claim_store.release([*batch_failed_candidates, *batch_policy_skipped_candidates]) if strategy_manager is not None: strategy_manager.record_accepted(batch_submitted_candidates) + outcome_tracker.save() + if state_sync is not None and state_batch_id is not None: + ordered_batch_results = [ + ordered_results[index] + for index in range(len(batch_candidates)) + if index in ordered_results + ] + if not state_sync.finish_batch(state_batch_id, ordered_batch_results): + state_sync_paused = True + print("[cycle] paused reason=state_sync_unhealthy phase=result", flush=True) + break if hasattr(modelhub_client, "available_submit_slots") and modelhub_client.available_submit_slots() <= 0: break @@ -1435,6 +1630,7 @@ def run_submission( "historyArchivePath": str(history_archive_path), "historyArchiveRecordCount": len(archived_history), "gpuStrategy": strategy_summary, + "officialCapabilities": official_summary, "marketIntelligence": market_summary, "candidatePreflight": preflight_summary, "agePolicy": { @@ -1471,6 +1667,12 @@ def run_submission( "warnings": report.get("warnings", []), "runDir": str(run_dir), "outcomeSyncCount": synced_count, + "stateSync": { + "enabled": state_sync is not None, + "healthy": bool(getattr(state_sync, "healthy", True)), + "paused": state_sync_paused, + "generation": int(getattr(state_sync, "generation", 0) or 0), + }, } write_json(run_dir / "summary.json", summary) return summary diff --git a/modelhub_submmit_api/market_intelligence.py b/modelhub_submmit_api/market_intelligence.py index 5e5ed186..9ae52597 100644 --- a/modelhub_submmit_api/market_intelligence.py +++ b/modelhub_submmit_api/market_intelligence.py @@ -3,6 +3,7 @@ from __future__ import annotations import math import re import statistics +import threading from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timedelta from pathlib import Path @@ -20,6 +21,9 @@ DEFAULT_FETCH_WORKERS = 4 DEFAULT_FRAMEWORK_MIN_SAMPLES = 300 DEFAULT_GPU_MIN_RECENT_TERMINALS = 20 DEFAULT_FRAMEWORK_MIN_WILSON = 0.05 +DEFAULT_COMMUNITY_WINDOW_DAYS = 30 +DEFAULT_COMMUNITY_MAX_MODELS = 2000 +DEFAULT_COMMUNITY_HYDRATE_PER_REFRESH = 50 NEW_FRAMEWORK_PROMOTION_MARGIN = 1.10 ERROR_RETRY_SECONDS = 300 @@ -143,6 +147,7 @@ class MarketIntelligenceManager: self.state: dict[str, Any] | None = None self.local_outcome_stats: dict[str, Any] = {} self._last_framework_request_errors: set[tuple[str, str]] = set() + self._route_refresh_lock = threading.Lock() def set_local_outcome_stats(self, report: dict[str, Any] | None) -> None: self.local_outcome_stats = report if isinstance(report, dict) else {} @@ -214,6 +219,9 @@ class MarketIntelligenceManager: "frameworkStats": {}, "queueError": None, "frameworkError": None, + "communitySample": {}, + "communityUpdatedAt": None, + "communityError": None, } def prepare( @@ -233,6 +241,9 @@ class MarketIntelligenceManager: elif int((loaded or {}).get("version") or 0) == 2 and self._catalog_compatible(loaded, gpus, tasks): state = dict(loaded or {}) state["version"] = MARKET_STATE_VERSION + state.setdefault("communitySample", {}) + state.setdefault("communityUpdatedAt", now.isoformat()) + state.setdefault("communityError", None) self._rescore_cached_gpu_stats(state.get("gpuStats") or {}) else: state = self._base_state(gpus, tasks, now) @@ -314,6 +325,31 @@ class MarketIntelligenceManager: state["frameworkStats"] = {} self.log(f"[market] framework_refresh_error reason={state['frameworkError']}") + community_due = not _fresh( + state.get("communityUpdatedAt"), + now=now, + ttl_seconds=900, + ) + if state.get("communityError") and _fresh( + state.get("communityAttemptedAt"), + now=now, + ttl_seconds=300, + ): + community_due = False + if community_due: + state["communityAttemptedAt"] = now.isoformat() + try: + state["communitySample"] = self._refresh_community_sample( + client, + previous=state.get("communitySample") or {}, + now=now, + ) + state["communityUpdatedAt"] = now.isoformat() + state["communityError"] = None + except Exception as exc: + state["communityError"] = f"{type(exc).__name__}: {exc}" + self.log(f"[market] community_refresh_error reason={state['communityError']}") + state["generatedAt"] = now.isoformat() state["throughputWindowHours"] = self.throughput_window_hours self.state = state @@ -321,6 +357,72 @@ class MarketIntelligenceManager: self._log_snapshot("refreshed" if queue_due or framework_due else "loaded") return state + def _refresh_community_sample( + self, + client: Any, + *, + previous: dict[str, Any], + now: datetime, + ) -> dict[str, Any]: + cutoff = now - timedelta(days=DEFAULT_COMMUNITY_WINDOW_DAYS) + sample = { + str(model_id): dict(item) + for model_id, item in previous.items() + if isinstance(item, dict) + and (parse_datetime(item.get("updateTime")) or cutoff) >= cutoff + } + payload = client.list_tasks_page( + current=1, + page_size=200, + only_mine=False, + begin_time=cutoff, + end_time=now, + ) + records = ((payload.get("data") or {}).get("records") or []) if isinstance(payload, dict) else [] + candidates: list[tuple[str, dict[str, Any]]] = [] + for record in records: + if not isinstance(record, dict): + continue + model_id = str(record.get("modelId") or "").strip() + if model_id and model_id not in sample: + candidates.append((model_id, record)) + if len(candidates) >= DEFAULT_COMMUNITY_HYDRATE_PER_REFRESH: + break + for model_id, record in candidates: + evidence: list[dict[str, Any]] = [] + try: + verify = client.get_verify_result_map(model_id) + except Exception: + verify = {} + if isinstance(verify, dict): + for gpu, gpu_result in verify.items(): + rows = gpu_result.get("records") if isinstance(gpu_result, dict) else [] + for row in rows or []: + if not isinstance(row, dict): + continue + evidence.append( + { + "gpu": gpu, + "framework": row.get("framework"), + "verifyResult": row.get("verifyResult"), + "taskId": row.get("taskId") or row.get("contestTaskId"), + } + ) + sample[model_id] = { + "modelId": model_id, + "updateTime": record.get("updateTime") or now.isoformat(), + "taskType": record.get("taskType") or record.get("modelTaskLevel"), + "evidence": evidence, + } + if len(sample) > DEFAULT_COMMUNITY_MAX_MODELS: + ordered = sorted( + sample.items(), + key=lambda item: parse_datetime(item[1].get("updateTime")) or cutoff, + reverse=True, + ) + sample = dict(ordered[:DEFAULT_COMMUNITY_MAX_MODELS]) + return sample + def _fetch_gpu_stats(self, client: Any, *, gpus: list[str], now: datetime) -> dict[str, dict[str, Any]]: local_end = now.astimezone() local_begin = local_end - timedelta(hours=self.throughput_window_hours) @@ -740,6 +842,33 @@ class MarketIntelligenceManager: compatible.append(framework) return compatible + def has_framework_route(self, task_type: str, target_gpu: str) -> bool: + rows = ( + (((self.state or {}).get("frameworkStats") or {}).get(task_type) or {}).get(target_gpu) + or {} + ) + return any(bool(item.get("officialConfigValid", False)) for item in rows.values()) + + def ensure_framework_route(self, client: Any, task_type: str, target_gpu: str) -> None: + if self.has_framework_route(task_type, target_gpu): + return + with self._route_refresh_lock: + if self.has_framework_route(task_type, target_gpu): + return + fetched = self._fetch_framework_stats( + client, + gpus=[target_gpu], + task_types=[task_type], + ) + rows = ((fetched.get(task_type) or {}).get(target_gpu) or {}) + if not rows: + return + state = self.state or self._base_state([target_gpu], [task_type], utc_now()) + state.setdefault("frameworkStats", {}).setdefault(task_type, {})[target_gpu] = rows + state["frameworkUpdatedAt"] = utc_now().isoformat() + self.state = state + write_json(self.path, state) + def official_config( self, *, @@ -814,5 +943,8 @@ class MarketIntelligenceManager: "throughputWindowHours": int(self.state.get("throughputWindowHours") or self.throughput_window_hours), "queueError": self.state.get("queueError"), "frameworkError": self.state.get("frameworkError"), + "communityUpdatedAt": self.state.get("communityUpdatedAt"), + "communityError": self.state.get("communityError"), + "communityModelCount": len(self.state.get("communitySample") or {}), "gpuStats": dict(self.state.get("gpuStats") or {}), } diff --git a/modelhub_submmit_api/modelhub_client.py b/modelhub_submmit_api/modelhub_client.py index 238567b3..a7e534d8 100644 --- a/modelhub_submmit_api/modelhub_client.py +++ b/modelhub_submmit_api/modelhub_client.py @@ -39,6 +39,8 @@ def parse_model_submission_precheck(payload: Any) -> dict[str, Any]: return { "isInDB": data.get("isInDB") is True, "processedGpus": set(str(gpu) for gpu in verify_result), + "verifyResult": verify_result, + "modelInfo": data.get("modelInfo") if isinstance(data.get("modelInfo"), dict) else {}, } @@ -124,6 +126,16 @@ class ModelHubClient: raise ModelHubAPIError("Machine info response is invalid", payload=payload) return [item for item in data if isinstance(item, dict)] + def list_task_levels(self) -> dict[str, Any]: + return self._request("GET", "/api/computility/task-levels/tree") + + def list_model_task_types(self, target_gpu: str, model_address: str) -> dict[str, Any]: + return self._request( + "GET", + "/api/computility/driver/images/task-types/by-gpu", + query={"gpuTypeName": target_gpu, "modelAddress": model_address}, + ) + def list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict[str, Any]]: payload = self._request( "GET", @@ -942,6 +954,8 @@ class ModelHubClientPool: self._finish_reservation(selected_index, reservation_id, succeeded=True) if capacity_probe: self._promote_account_capacity(selected_index) + if isinstance(response, dict): + response.setdefault("_accountKey", self._account_keys[selected_index]) return response def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001 @@ -951,6 +965,12 @@ class ModelHubClientPool: def list_machine_info(self) -> list[dict[str, Any]]: return self._reader.list_machine_info() + def list_task_levels(self) -> dict[str, Any]: + return self._reader.list_task_levels() + + def list_model_task_types(self, target_gpu: str, model_address: str) -> dict[str, Any]: + return self._reader.list_model_task_types(target_gpu, model_address) + def list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict[str, Any]]: return self._reader.list_framework_stats(task_type, target_gpu) @@ -963,3 +983,30 @@ class ModelHubClientPool: if task_id is not None: return task_id return None + + def list_active_tasks_by_account(self) -> list[dict[str, Any]]: + """Return active tasks with a stable, non-secret account fingerprint.""" + def fetch_account(index: int, client: ModelHubClient) -> tuple[int, list[dict[str, Any]]]: + tasks: list[dict[str, Any]] = [] + seen: set[str] = set() + for status in ("waiting", "running"): + for task in client.list_tasks(page_size=100, only_mine=True, status=status): + task_id = str(task.get("taskId") or "") + if task_id and task_id in seen: + continue + if task_id: + seen.add(task_id) + if isinstance(task, dict) and is_active_task(task): + tasks.append(task) + return index, tasks + + merged: list[dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=min(12, len(self.clients))) as executor: + futures = { + executor.submit(fetch_account, index, client): index + for index, client in enumerate(self.clients) + } + for future in as_completed(futures): + index, tasks = future.result() + merged.extend({**task, "accountKey": self._account_keys[index]} for task in tasks) + return merged diff --git a/modelhub_submmit_api/models.py b/modelhub_submmit_api/models.py index 1c2ec6f6..a1dc2b38 100644 --- a/modelhub_submmit_api/models.py +++ b/modelhub_submmit_api/models.py @@ -25,6 +25,14 @@ class HFModelSummary: last_modified: datetime | None pipeline_tag: str | None created_at: datetime | None = None + params: int | None = None + file_size: int | None = None + tags: tuple[str, ...] = () + tasks: tuple[str, ...] = () + license: str | None = None + gated: bool = False + private: bool = False + likes: int = 0 @property def model_address(self) -> str: @@ -42,6 +50,8 @@ class ModelInspection: onnx_files: list[str] = field(default_factory=list) model_config: dict[str, Any] = field(default_factory=dict) config_fetch_error: str | None = None + model_card_metadata: dict[str, Any] = field(default_factory=dict) + published_size_bytes: int | None = None @property def root_file_names(self) -> set[str]: @@ -132,10 +142,13 @@ class ModelInspection: @property def repository_size_bytes(self) -> int | None: """Return exact recursive on-disk size when every file has a size.""" - if not self.file_paths or any(path not in self.file_sizes for path in self.file_paths): - return None - total = sum(max(0, int(self.file_sizes[path])) for path in self.file_paths) - return total if total > 0 else None + exact: int | None = None + if self.file_paths and not any(path not in self.file_sizes for path in self.file_paths): + total = sum(max(0, int(self.file_sizes[path])) for path in self.file_paths) + exact = total if total > 0 else None + published = int(self.published_size_bytes or 0) or None + values = [value for value in (exact, published) if value is not None] + return max(values) if values else None def estimated_load_bytes(self, framework: str) -> int | None: if framework == "llamacpp": diff --git a/modelhub_submmit_api/official_capabilities.py b/modelhub_submmit_api/official_capabilities.py new file mode 100644 index 00000000..2f2fc43c --- /dev/null +++ b/modelhub_submmit_api/official_capabilities.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +from datetime import datetime, timedelta +from pathlib import Path +import threading +from typing import Any + +from common import parse_datetime, read_json, utc_now, write_json + + +OFFICIAL_CAPABILITY_VERSION = 1 +DEFAULT_OFFICIAL_CAPABILITIES_PATH = Path(".modelhub_state/official_capabilities.json") + + +class OfficialCapabilityUnavailable(RuntimeError): + pass + + +def _fresh(value: Any, now: datetime, seconds: int) -> bool: + parsed = parse_datetime(value) + return parsed is not None and now - parsed <= timedelta(seconds=max(1, seconds)) + + +def _extract_task_names(payload: Any) -> list[str]: + data = payload.get("data") if isinstance(payload, dict) else payload + result: list[str] = [] + + def visit(value: Any) -> None: + if isinstance(value, str): + if value.strip() and value.strip() not in result: + result.append(value.strip()) + return + if isinstance(value, list): + for item in value: + visit(item) + return + if not isinstance(value, dict): + return + for key in ("taskType", "taskTypeName", "value", "code"): + candidate = value.get(key) + if isinstance(candidate, str) and candidate.strip() and candidate.strip() not in result: + result.append(candidate.strip()) + children = value.get("children") + if not children: + candidate = value.get("name") + if isinstance(candidate, str) and candidate.strip() and candidate.strip() not in result: + result.append(candidate.strip()) + for key in ("children", "items", "records", "list"): + visit(value.get(key)) + + visit(data) + return result + + +class OfficialCapabilityRegistry: + def __init__(self, path: Path | str = DEFAULT_OFFICIAL_CAPABILITIES_PATH, *, log_fn=None) -> None: + self.path = Path(path) + self.log = log_fn or (lambda message: print(message, flush=True)) + self.state: dict[str, Any] = {} + self.ready = False + self.pause_reason: str | None = None + self._lock = threading.Lock() + + def _load(self) -> dict[str, Any]: + try: + value = read_json(self.path) + except (FileNotFoundError, ValueError): + return {} + if not isinstance(value, dict) or int(value.get("version") or 0) != OFFICIAL_CAPABILITY_VERSION: + return {} + return value + + def prepare(self, client: Any, *, fallback_gpus: list[str], task_types: list[str], now: datetime | None = None) -> dict[str, Any]: + now = now or utc_now() + state = self._load() + catalog_usable = _fresh(state.get("catalogUpdatedAt"), now, 1800) + task_tree_usable = _fresh(state.get("taskTreeUpdatedAt"), now, 7 * 86400) + errors: list[str] = [] + try: + machines = client.list_machine_info() + catalog: dict[str, dict[str, Any]] = {} + for item in machines: + gpu = str(item.get("gpuType") or item.get("gpuTypeName") or "").strip() + if gpu: + catalog[gpu] = { + "gpu": gpu, + "canVerify": item.get("canVerify"), + "maxConcurrentTasks": item.get("maxConcurrentTasks"), + } + if not catalog: + raise OfficialCapabilityUnavailable("official GPU catalog is empty") + state["gpuCatalog"] = catalog + state["catalogUpdatedAt"] = now.isoformat() + catalog_usable = True + except Exception as exc: + errors.append(f"catalog:{type(exc).__name__}:{exc}") + + try: + payload = client.list_task_levels() + dynamic_tasks = _extract_task_names(payload) + if not dynamic_tasks: + raise OfficialCapabilityUnavailable("official task tree is empty") + state["taskTree"] = payload.get("data") if isinstance(payload, dict) else payload + state["discoveredTaskTypes"] = dynamic_tasks + state["taskTreeUpdatedAt"] = now.isoformat() + task_tree_usable = True + except Exception as exc: + errors.append(f"task_tree:{type(exc).__name__}:{exc}") + + state.update( + { + "version": OFFICIAL_CAPABILITY_VERSION, + "generatedAt": now.isoformat(), + "configuredTaskTypes": list(task_types), + "errors": errors, + } + ) + self.state = state + self.ready = bool(catalog_usable and task_tree_usable) + self.pause_reason = None if self.ready else "critical_official_signal_unavailable" + write_json(self.path, state) + eligible = self.eligible_gpus(fallback_gpus) + self.log( + f"[official] ready={str(self.ready).lower()} catalog={len(state.get('gpuCatalog') or {})} " + f"task_types={len(state.get('discoveredTaskTypes') or [])} eligible={len(eligible)} " + f"errors={len(errors)}" + ) + return state + + def eligible_gpus(self, requested: list[str] | None = None) -> list[str]: + catalog = self.state.get("gpuCatalog") or {} + names = list(requested or catalog.keys()) + return [ + gpu + for gpu in names + if gpu in catalog + and (catalog[gpu].get("canVerify") is not False) + and int(catalog[gpu].get("maxConcurrentTasks") or 0) > 0 + ] + + def task_types_for(self, client: Any, *, model_address: str, model_last_modified: Any, gpu: str) -> list[str]: + now = utc_now() + key = f"{model_address}|{model_last_modified or ''}|{gpu}" + with self._lock: + cache = self.state.setdefault("modelGpuTaskTypes", {}) + item = dict(cache.get(key) or {}) + ttl = 86400 if item.get("taskTypes") else 3600 + if _fresh(item.get("updatedAt"), now, ttl): + return list(item.get("taskTypes") or []) + try: + payload = client.list_model_task_types(gpu, model_address) + task_types = _extract_task_names(payload) + except Exception as exc: + if _fresh(item.get("updatedAt"), now, 86400): + return list(item.get("taskTypes") or []) + raise OfficialCapabilityUnavailable( + f"model+GPU task capability unavailable for {gpu}: {type(exc).__name__}: {exc}" + ) from exc + with self._lock: + cache = self.state.setdefault("modelGpuTaskTypes", {}) + cache[key] = {"updatedAt": now.isoformat(), "taskTypes": task_types} + self.state["generatedAt"] = now.isoformat() + write_json(self.path, self.state) + return task_types + + def summary(self) -> dict[str, Any]: + return { + "enabled": True, + "ready": self.ready, + "pauseReason": self.pause_reason, + "catalogUpdatedAt": self.state.get("catalogUpdatedAt"), + "taskTreeUpdatedAt": self.state.get("taskTreeUpdatedAt"), + "gpuCount": len(self.state.get("gpuCatalog") or {}), + "discoveredTaskTypeCount": len(self.state.get("discoveredTaskTypes") or []), + "errors": list(self.state.get("errors") or []), + } diff --git a/modelhub_submmit_api/outcome_tracker.py b/modelhub_submmit_api/outcome_tracker.py index 2674e5c1..f5d92b89 100644 --- a/modelhub_submmit_api/outcome_tracker.py +++ b/modelhub_submmit_api/outcome_tracker.py @@ -492,6 +492,7 @@ class OutcomeTracker: framework_groups: dict[str, list[dict[str, Any]]] = defaultdict(list) combo_groups: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list) profile_groups: dict[tuple[str, str, str, str, str], list[dict[str, Any]]] = defaultdict(list) + sized_profile_groups: dict[tuple[str, str, str, str, str, int], list[dict[str, Any]]] = defaultdict(list) for record in terminal: gpu = record.get("targetGpu") or "unknown" @@ -505,6 +506,13 @@ class OutcomeTracker: if model_type: quantization = str(profile.get("quantizationMethod") or "none").strip() profile_groups[(gpu, fw, tt, model_type, quantization)].append(record) + try: + load_bytes = max(1, int(profile.get("estimatedLoadBytes") or 0)) + except (TypeError, ValueError): + load_bytes = 0 + if load_bytes > 0: + size_bucket = int(load_bytes).bit_length() - 1 + sized_profile_groups[(gpu, fw, tt, model_type, quantization, size_bucket)].append(record) gpu_summaries = {gpu: _summarize(records) for gpu, records in gpu_groups.items()} framework_summaries = {fw: _summarize(records) for fw, records in framework_groups.items()} @@ -561,6 +569,19 @@ class OutcomeTracker: "lastTerminalAt": last_terminal_at.isoformat() if last_terminal_at else None, } + sized_profile_combination_stats: dict[str, dict[str, Any]] = {} + for (gpu, fw, tt, model_type, quantization, size_bucket), records in sized_profile_groups.items(): + key = f"{gpu}|{fw}|{tt}|{model_type}|{quantization}|{size_bucket}" + sized_profile_combination_stats[key] = { + "targetGpu": gpu, + "framework": fw, + "taskType": tt, + "modelType": model_type, + "quantizationMethod": quantization, + "loadSizeLog2Bucket": size_bucket, + **_summarize(records), + } + warnings: list[str] = [] for gpu, summary in gpu_summaries.items(): if summary["decisionTotal"] >= 4 and summary["decisionFailureRate"] >= 0.5: @@ -607,6 +628,7 @@ class OutcomeTracker: "recentCombinationStats": recent_combination_stats, "profileCombinationStats": profile_combination_stats, "recentProfileCombinationStats": recent_profile_combination_stats, + "sizedProfileCombinationStats": sized_profile_combination_stats, "architectureCompatibilityBlocks": architecture_blocks, "architectureCompatibilitySummary": { "activeBlockCount": len(architecture_blocks), diff --git a/modelhub_submmit_api/poll_runner.py b/modelhub_submmit_api/poll_runner.py index f9e62bc1..683a3c39 100644 --- a/modelhub_submmit_api/poll_runner.py +++ b/modelhub_submmit_api/poll_runner.py @@ -25,8 +25,17 @@ from market_intelligence import ( ) from modelhub_client import DEFAULT_CAPACITY_STATE_PATH, ModelHubClient, ModelHubClientPool from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker +from official_capabilities import DEFAULT_OFFICIAL_CAPABILITIES_PATH from queue_cleanup import cleanup_certain_oom_tasks +from routing_engine import DEFAULT_ROUTING_STATE_PATH from runner_common import DEFAULT_KEY_PATH, ensure_tokens +from state_sync import ( + DEFAULT_BATCH_SIZE, + DEFAULT_BRANCH, + DEFAULT_REMOTE, + StateGitSync, + load_state_git_credentials, +) from submission_claims import DEFAULT_CLAIMS_PATH from template_selector import TemplateSelector from version import AGENT_VERSION @@ -132,6 +141,33 @@ def build_parser() -> argparse.ArgumentParser: default=os.getenv("MODELHUB_GPU_STRATEGY_STATE_PATH", str(DEFAULT_GPU_STRATEGY_PATH)), help=argparse.SUPPRESS, ) + parser.add_argument( + "--routing-state-path", + default=os.getenv("MODELHUB_ROUTING_STATE_PATH", str(DEFAULT_ROUTING_STATE_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--official-capabilities-path", + default=os.getenv("MODELHUB_OFFICIAL_CAPABILITIES_PATH", str(DEFAULT_OFFICIAL_CAPABILITIES_PATH)), + help=argparse.SUPPRESS, + ) + parser.add_argument("--state-sync", action="store_true", help=argparse.SUPPRESS) + parser.add_argument( + "--state-sync-remote", + default=os.getenv("MODELHUB_STATE_SYNC_REMOTE", DEFAULT_REMOTE), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--state-sync-branch", + default=os.getenv("MODELHUB_STATE_SYNC_BRANCH", DEFAULT_BRANCH), + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--state-sync-batch-size", + type=int, + default=int(os.getenv("MODELHUB_STATE_SYNC_BATCH_SIZE", str(DEFAULT_BATCH_SIZE))), + 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( @@ -471,10 +507,53 @@ def run_poll_loop( outcome_tracker: OutcomeTracker | None = None, ) -> dict[str, Any]: now = now or utc_now() + state_sync: StateGitSync | None = getattr(base_args, "_state_sync_manager", None) + if state_sync is None and bool(getattr(base_args, "state_sync", False)): + state_sync = StateGitSync( + project_root=Path(__file__).resolve().parent.parent, + credentials=load_state_git_credentials(), + remote=str(getattr(base_args, "state_sync_remote", DEFAULT_REMOTE)), + branch=str(getattr(base_args, "state_sync_branch", DEFAULT_BRANCH)), + batch_size=max(1, int(getattr(base_args, "state_sync_batch_size", DEFAULT_BATCH_SIZE) or DEFAULT_BATCH_SIZE)), + log_fn=log, + ) + try: + state_sync.acquire_process_lock() + except Exception as exc: + state_sync.last_error = str(exc) + state_sync.healthy = False + else: + state_sync.restore() + base_args._state_sync_manager = state_sync + if state_sync is not None: + state_sync.write_readiness( + ready=False, + reason="startup_recovery" if state_sync.healthy else "state_sync_unhealthy", + ) hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url) modelhub_client = modelhub_client or _build_modelhub_client(base_args) template_selector = template_selector or TemplateSelector() + if state_sync is not None and state_sync.healthy: + try: + active_tasks = ( + modelhub_client.list_active_tasks_by_account() + if hasattr(modelhub_client, "list_active_tasks_by_account") + else [] + ) + recovery = state_sync.reconcile_active_tasks(active_tasks) + if not state_sync.sync("startup"): + log("[cycle] paused reason=state_sync_unhealthy") + else: + log( + f"[state-recovery] active={recovery['active']} " + f"reconciled={recovery['reconciled']} unresolved={recovery['unresolved']}" + ) + except Exception as exc: + state_sync.healthy = False + state_sync.last_error = str(exc) + log(f"[state-recovery] active_scan_failed reason={type(exc).__name__}: {exc}") + poll_runs_dir = Path(base_args.poll_runs_dir) poll_runs_dir.mkdir(parents=True, exist_ok=True) poll_run_dir = make_run_dir(poll_runs_dir, now) @@ -535,6 +614,17 @@ def run_poll_loop( break cycles += 1 + if state_sync is not None and not state_sync.healthy: + recovered = ( + state_sync.sync("retry") + if state_sync._workspace is not None + else state_sync.retry_restore() + ) + if not recovered: + state_sync.write_readiness(ready=False, reason="state_sync_unhealthy") + log(f"[poll] cycle={cycles} sleep={base_args.idle_interval_seconds}s reason=state_sync_unhealthy") + time.sleep(base_args.idle_interval_seconds) + continue if hasattr(modelhub_client, "configure_capacity_probe"): modelhub_client.configure_capacity_probe(cycles) @@ -678,6 +768,10 @@ def run_poll_loop( "architectureIncompatibleCount": cleanup_summary[ "architectureIncompatibleCount" ], + "officialCapabilityInvalidCount": cleanup_summary.get( + "officialCapabilityInvalidCount", + 0, + ), "oldOverflowCount": cleanup_summary["oldOverflowCount"], "cancelledCount": cleanup_summary["cancelledCount"], "policyCancelledRecorded": policy_cancelled_recorded, @@ -701,6 +795,20 @@ def run_poll_loop( ) if available_slots is not None and available_slots <= 0: + if state_sync is not None: + try: + if cycles % 3 == 0 and hasattr(modelhub_client, "list_active_tasks_by_account"): + state_sync.reconcile_active_tasks(modelhub_client.list_active_tasks_by_account()) + sync_ok = state_sync.sync("cycle_no_slots") + state_sync.write_readiness( + ready=sync_ok, + reason=None if sync_ok else "state_sync_unhealthy", + extra={"cycle": cycles}, + ) + except Exception as exc: + state_sync.healthy = False + state_sync.last_error = str(exc) + state_sync.write_readiness(ready=False, reason="state_sync_unhealthy") log(f"[poll] cycle={cycles} sleep={base_args.poll_interval_seconds}s reason=no_available_slots") time.sleep(base_args.poll_interval_seconds) continue @@ -726,6 +834,34 @@ def run_poll_loop( f"stop={cycle_summary['stoppedReason']}" ) + if state_sync is not None: + try: + if cycle_summary.get("submittedTotal", 0) > 0 or cycles % 3 == 0: + active_tasks = ( + modelhub_client.list_active_tasks_by_account() + if hasattr(modelhub_client, "list_active_tasks_by_account") + else [] + ) + state_sync.reconcile_active_tasks(active_tasks) + sync_ok = state_sync.sync("cycle") + official_paused = any( + bool((wave_result.get("summary") or {}).get("paused")) + for wave_result in (cycle_summary.get("waveResults") or []) + ) + state_sync.write_readiness( + ready=sync_ok and not official_paused, + reason=( + "critical_official_signal_unavailable" + if official_paused + else (None if sync_ok else "state_sync_unhealthy") + ), + extra={"cycle": cycles}, + ) + except Exception as exc: + state_sync.healthy = False + state_sync.last_error = str(exc) + state_sync.write_readiness(ready=False, reason="state_sync_unhealthy") + if base_args.daily_target > 0 and remaining_before_run is not None and remaining_before_run <= 0: stopped_reason = "daily_target_already_reached" break @@ -791,6 +927,9 @@ def run_poll_loop( } write_json(poll_run_dir / "summary.json", summary) log(f"[poll] finished submitted_total={submitted_total} cycles={cycles} stopped_reason={stopped_reason}") + if state_sync is not None: + state_sync.write_readiness(ready=False, reason="worker_stopped") + state_sync.close() return summary diff --git a/modelhub_submmit_api/queue_cleanup.py b/modelhub_submmit_api/queue_cleanup.py index bbc7c953..e9dd27c1 100644 --- a/modelhub_submmit_api/queue_cleanup.py +++ b/modelhub_submmit_api/queue_cleanup.py @@ -262,6 +262,7 @@ def _load_framework_catalog( f"frameworks={sum(len(items) for items in catalog.values())} " f"unknown={len(errors)}" ) + return catalog, errors @@ -475,6 +476,37 @@ def cleanup_certain_oom_tasks( f"framework_exposed={sum(1 for task in tasks if task.framework)}" ) + capability_decisions: list[dict[str, Any]] = [] + capability_catalog_error: str | None = None + if not architecture_only: + try: + machine_rows = modelhub.list_machine_info() + machine_catalog = { + str(item.get("gpuType") or item.get("gpuTypeName") or "").strip().casefold(): item + for item in machine_rows + if isinstance(item, dict) + and str(item.get("gpuType") or item.get("gpuTypeName") or "").strip() + } + if not machine_catalog: + raise RuntimeError("official GPU catalog is empty") + for task in tasks: + machine = machine_catalog.get(task.gpu_type.casefold()) + if task.status != "waiting": + continue + if machine is None or machine.get("canVerify") is False: + capability_decisions.append( + { + "accountIndex": task.account_index + 1, + "taskId": task.task_id, + "modelId": task.model_id, + "gpuType": task.gpu_type, + "status": task.status, + "reason": "official_gpu_unavailable", + } + ) + except Exception as exc: + capability_catalog_error = f"{type(exc).__name__}: {exc}" + observed_active_counts: list[int | None] = [0 for _ in clients] for task in tasks: current_count = observed_active_counts[task.account_index] @@ -576,6 +608,8 @@ def cleanup_certain_oom_tasks( exhaustive_relevant = bool(not framework and gpu_task_pair in block_gpu_task_pairs) if exhaustive_relevant: framework_catalog_combinations.add((task.gpu_type, task_type)) + if not architecture_only and task.status == "waiting" and framework and task_type: + framework_catalog_combinations.add((task.gpu_type, task_type)) if (exact_relevant or exhaustive_relevant) and not ( profile.get("modelType") or profile.get("architectures") ): @@ -599,6 +633,43 @@ def cleanup_certain_oom_tasks( model_configs=model_configs, framework_catalog=framework_catalog, ) + if not architecture_only: + existing_capability_keys = { + (int(item["accountIndex"]), int(item["taskId"])) + for item in capability_decisions + } + for task in tasks: + if task.status != "waiting": + continue + context = task_contexts.get(str(task.task_id)) + context = context if isinstance(context, dict) else {} + framework = str(context.get("framework") or task.framework or "").strip() + task_type = str(context.get("taskType") or task.task_type or "").strip() + if not framework or not task_type: + continue + catalog_key = (task.gpu_type.casefold(), task_type.casefold()) + available = framework_catalog.get(catalog_key) + if available is None or framework in available: + continue + decision_key = (task.account_index + 1, task.task_id) + if decision_key in existing_capability_keys: + continue + capability_decisions.append( + { + "accountIndex": task.account_index + 1, + "taskId": task.task_id, + "modelId": task.model_id, + "gpuType": task.gpu_type, + "framework": framework, + "taskType": task_type, + "status": task.status, + "reason": "official_framework_unavailable", + } + ) + log( + f"[queue-cleanup] official_capability_invalid={len(capability_decisions)} " + f"catalog_error={capability_catalog_error or 'none'}" + ) log( f"[queue-cleanup] architecture_incompatible={len(architecture_decisions)} " f"blocks={len(architecture_blocks)} " @@ -612,7 +683,7 @@ def cleanup_certain_oom_tasks( ) decisions_by_key: dict[tuple[int, int], dict[str, Any]] = {} - for decision in [*oom_decisions, *architecture_decisions]: + for decision in [*oom_decisions, *architecture_decisions, *capability_decisions]: enriched = dict(decision) enriched["cleanupReasons"] = [decision["reason"]] key = (int(decision["accountIndex"]), int(decision["taskId"])) @@ -665,12 +736,11 @@ def cleanup_certain_oom_tasks( disappeared.append(decision) continue cleanup_reasons = set(decision.get("cleanupReasons") or [decision.get("reason")]) - architecture_without_oom = bool( - "known_framework_architecture_incompatible" in cleanup_reasons - and "certain_oom_repository_size_exceeds_gpu_capacity" not in cleanup_reasons + protected_without_oom = bool( + "certain_oom_repository_size_exceeds_gpu_capacity" not in cleanup_reasons ) current_status = active_status_by_account.get(account_index, {}).get(int(decision["taskId"])) - if architecture_without_oom and current_status != "waiting": + if protected_without_oom and current_status != "waiting": policy_no_longer_applies.append( { **decision, @@ -687,7 +757,7 @@ def cleanup_certain_oom_tasks( if stop_failed: break decisions_by_id = {int(item["taskId"]): item for item in by_account[account_index]} - for cleanup_phase in ("oom", "architecture"): + for cleanup_phase in ("oom", "architecture", "capability"): task_ids = sorted( task_id for task_id, decision in decisions_by_id.items() @@ -703,6 +773,15 @@ def cleanup_certain_oom_tasks( and "known_framework_architecture_incompatible" in decision["cleanupReasons"] ) + or ( + cleanup_phase == "capability" + and "certain_oom_repository_size_exceeds_gpu_capacity" + not in decision["cleanupReasons"] + and any( + reason in decision["cleanupReasons"] + for reason in ("official_gpu_unavailable", "official_framework_unavailable") + ) + ) ) if cleanup_phase != "oom" and task_ids: # A waiting architecture task can start running after the @@ -737,12 +816,23 @@ def cleanup_certain_oom_tasks( decisions_by_id[task_id].get("cleanupReasons") or [decisions_by_id[task_id].get("reason")] ) - architecture_applies = bool( - cleanup_phase == "architecture" - and "known_framework_architecture_incompatible" in reasons - and current_task.status == "waiting" + protected_policy_applies = bool( + current_task.status == "waiting" + and ( + ( + cleanup_phase == "architecture" + and "known_framework_architecture_incompatible" in reasons + ) + or ( + cleanup_phase == "capability" + and any( + reason in reasons + for reason in ("official_gpu_unavailable", "official_framework_unavailable") + ) + ) + ) ) - if not architecture_applies: + if not protected_policy_applies: policy_no_longer_applies.append( { **decisions_by_id[task_id], @@ -799,6 +889,9 @@ def cleanup_certain_oom_tasks( "architectureBlockCount": len(architecture_blocks), "architectureIncompatibleCount": len(architecture_decisions), "architectureIncompatibleTasks": architecture_decisions, + "officialCapabilityInvalidCount": len(capability_decisions), + "officialCapabilityInvalidTasks": capability_decisions, + "officialCapabilityCatalogError": capability_catalog_error, "architectureModelConfigsComplete": len(model_configs), "architectureModelConfigErrors": model_config_errors, "architectureFrameworkCatalog": { diff --git a/modelhub_submmit_api/routing_engine.py b/modelhub_submmit_api/routing_engine.py new file mode 100644 index 00000000..ac4ef7fe --- /dev/null +++ b/modelhub_submmit_api/routing_engine.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import math +from collections import defaultdict +from pathlib import Path +from typing import Any + +from common import read_json, utc_now, write_json +from gpu_strategy import wilson_lower_bound +from submission_claims import candidate_key + + +ROUTING_STATE_VERSION = 1 +DEFAULT_ROUTING_STATE_PATH = Path(".modelhub_state/routing_intelligence.json") + + +class SuccessFirstRoutingEngine: + """Deterministic success-first routing; queue speed only breaks close races.""" + + def __init__( + self, + path: Path | str = DEFAULT_ROUTING_STATE_PATH, + *, + outcome_stats: dict[str, Any] | None = None, + log_fn=None, + ) -> None: + self.path = Path(path) + self.log = log_fn or (lambda message: print(message, flush=True)) + self.state: dict[str, Any] = self._load() + self.outcome_stats = outcome_stats if isinstance(outcome_stats, dict) else {} + + def _load(self) -> dict[str, Any]: + try: + value = read_json(self.path) + except (FileNotFoundError, ValueError): + value = {} + if not isinstance(value, dict) or int(value.get("version") or 0) != ROUTING_STATE_VERSION: + value = { + "version": ROUTING_STATE_VERSION, + "generatedAt": utc_now().isoformat(), + "acceptedTotal": 0, + "acceptedSinceRefresh": 0, + "acceptedByCategory": {"unified_success_first": 0}, + "acceptedByRoute": {}, + } + return value + + def _evidence(self, candidate: dict[str, Any]) -> tuple[str, int, int, float]: + profile = candidate.get("preflightMetadata") or {} + gpu = str(candidate.get("targetGpu") or "") + framework = str(candidate.get("framework") or "") + task_type = str(candidate.get("taskType") or "") + model_type = str(profile.get("modelType") or "").strip() + quantization = str(profile.get("quantizationMethod") or "none").strip() + try: + load_bytes = max(1, int(profile.get("estimatedLoadBytes") or 0)) + except (TypeError, ValueError): + load_bytes = 0 + if model_type and load_bytes > 0: + size_bucket = int(load_bytes).bit_length() - 1 + key = f"{gpu}|{framework}|{task_type}|{model_type}|{quantization}|{size_bucket}" + item = (self.outcome_stats.get("sizedProfileCombinationStats") or {}).get(key) or {} + total = int(item.get("decisionTotal") or 0) + if total >= 20: + success = int(item.get("successCount") or 0) + return "local_profile_size", success, total, wilson_lower_bound(success, total) + if model_type: + key = f"{gpu}|{framework}|{task_type}|{model_type}|{quantization}" + item = (self.outcome_stats.get("profileCombinationStats") or {}).get(key) or {} + total = int(item.get("decisionTotal") or 0) + if total >= 40: + success = int(item.get("successCount") or 0) + return "local_profile", success, total, wilson_lower_bound(success, total) + public_samples = max(0, int(candidate.get("frameworkMarketSamples") or 0)) + public_rate = candidate.get("frameworkMarketSuccessRate") + if public_rate is None: + public_success = 0 + else: + public_success = max(0, min(public_samples, round(float(public_rate) * public_samples))) + local_samples = max(0, int(candidate.get("frameworkLocalSamples") or 0)) + local_rate = candidate.get("frameworkLocalSuccessRate") + local_success = 0 if local_rate is None else max(0, min(local_samples, round(float(local_rate) * local_samples))) + + if local_samples >= 100: + total = local_samples + success = local_success + level = "local_task_gpu_framework" + elif public_samples >= 300: + total = public_samples + success = public_success + level = "official_task_gpu_framework" + elif public_samples > 0: + total = public_samples + success = public_success + level = "official_low_sample" + else: + total = 2 + success = 1 + level = "global_conservative_prior" + return level, success, total, wilson_lower_bound(success, total) + + def annotate(self, candidate: dict[str, Any]) -> dict[str, Any]: + annotated = dict(candidate) + evidence, success, total, lower = self._evidence(candidate) + raw_eta = candidate.get("queueBacklogHours") + try: + eta = max(0.25, float(raw_eta)) + if not math.isfinite(eta) or eta >= 9999: + raise ValueError + eta_unknown = False + except (TypeError, ValueError): + eta = 36.0 + eta_unknown = True + success_band = int(lower / 0.05) + expected = (lower * lower) / eta + annotated.update( + { + "routingEvidenceLevel": evidence, + "routingSuccesses": success, + "routingSamples": total, + "routingSuccessLowerBound": lower, + "routingSuccessBand": success_band, + "routingEtaHours": eta, + "routingEtaUnknown": eta_unknown, + "routingExpectedSuccessPerHour": expected, + "strategyCategory": "unified_success_first", + } + ) + return annotated + + def order_candidates(self, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not candidates: + return [] + annotated = [self.annotate(candidate) for candidate in candidates] + by_model: dict[str, list[dict[str, Any]]] = defaultdict(list) + for candidate in annotated: + by_model[str(candidate.get("repoId") or candidate.get("modelAddress") or "")].append(candidate) + + selected: list[dict[str, Any]] = [] + for routes in by_model.values(): + best = max(float(item["routingSuccessLowerBound"]) for item in routes) + close = [item for item in routes if best - float(item["routingSuccessLowerBound"]) <= 0.05] + close.sort( + key=lambda item: ( + -int(item["routingSuccessBand"]), + -float(item["routingExpectedSuccessPerHour"]), + -int(item["routingSamples"]), + str(item.get("targetGpu") or ""), + str(item.get("framework") or ""), + ) + ) + selected.append(close[0]) + + selected.sort( + key=lambda item: ( + -int(item["routingSuccessBand"]), + -float(item["routingExpectedSuccessPerHour"]), + -int(item["routingSamples"]), + -int(item.get("downloads") or 0), + str(item.get("repoId") or ""), + ) + ) + return selected + + @property + def submissions_until_refresh(self) -> int: + return 1_000_000_000 + + def record_accepted(self, candidates: list[dict[str, Any]]) -> dict[str, Any]: + routes = self.state.setdefault("acceptedByRoute", {}) + for candidate in candidates: + key = "|".join( + str(candidate.get(name) or "") + for name in ("taskType", "targetGpu", "framework") + ) + routes[key] = int(routes.get(key) or 0) + 1 + self.state["acceptedTotal"] = int(self.state.get("acceptedTotal") or 0) + len(candidates) + self.state["acceptedSinceRefresh"] = int(self.state.get("acceptedSinceRefresh") or 0) + len(candidates) + categories = self.state.setdefault("acceptedByCategory", {"unified_success_first": 0}) + categories["unified_success_first"] = int(categories.get("unified_success_first") or 0) + len(candidates) + self.state["generatedAt"] = utc_now().isoformat() + write_json(self.path, self.state) + return self.state + + def summary(self) -> dict[str, Any]: + return { + "enabled": True, + "mode": "unified_success_first", + "statePath": str(self.path), + "acceptedTotal": int(self.state.get("acceptedTotal") or 0), + "routeCount": len(self.state.get("acceptedByRoute") or {}), + } diff --git a/modelhub_submmit_api/setup.sh b/modelhub_submmit_api/setup.sh index 5ae4660c..184b0b14 100644 --- a/modelhub_submmit_api/setup.sh +++ b/modelhub_submmit_api/setup.sh @@ -6,8 +6,8 @@ echo "=== ModelHub Submmit Setup ===" python3 -c 'import sys; sys.exit(0 if sys.version_info >= (3,9) else 1)' || { echo "ERROR: Python 3.9+ required"; exit 1; } mkdir -p runs ledger outcomes history daily_runs poll_runs logs python3 -c ' -import json,argparse,sys,time,pathlib,typing,concurrent.futures,threading,http.client,urllib,socket,dataclasses,collections,datetime,os,re -print("Python stdlib OK (zero external deps)") +import json,argparse,sys,time,pathlib,typing,concurrent.futures,threading,http.client,urllib,socket,dataclasses,collections,datetime,os,re,yaml +print("Python runtime dependencies OK") ' wc -l templates/public_submit/adapt_task_templates.jsonl grep -q 'XC_TOKEN\s*=\s*[a-z0-9]\{32\}' KEY.md 2>/dev/null && echo "XC_TOKEN found" || echo "WARNING: Set XC_TOKEN in KEY.md" diff --git a/modelhub_submmit_api/state_sync.py b/modelhub_submmit_api/state_sync.py new file mode 100644 index 00000000..de1f514e --- /dev/null +++ b/modelhub_submmit_api/state_sync.py @@ -0,0 +1,674 @@ +from __future__ import annotations + +import fcntl +import hashlib +import json +import os +import shutil +import stat +import subprocess +import tempfile +import threading +import uuid +import re +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Iterable + +from common import read_json, read_jsonl, write_json, write_jsonl +from runner_common import load_key_files +from version import AGENT_VERSION + + +STATE_SCHEMA_VERSION = 1 +DEFAULT_REMOTE = "https://dev.modelhub.org.cn/CoolBoy/submmit.git" +DEFAULT_BRANCH = "agent-state" +DEFAULT_BATCH_SIZE = 20 +DEFAULT_RETENTION_DAYS = 30 +DEFAULT_HISTORY_DEPTH = 20 + +# Only these runtime files may cross the trust boundary into the state branch. +# Credentials, raw stdout, downloaded archives and run directories are excluded. +STATE_ALLOWLIST = ( + ".modelhub_state/account_capacity.json", + ".modelhub_state/architecture_compatibility_blacklist.json", + ".modelhub_state/gpu_strategy.json", + ".modelhub_state/market_intelligence.json", + ".modelhub_state/official_capabilities.json", + ".modelhub_state/queue_cleanup_latest.json", + ".modelhub_state/recovery_active_tasks.jsonl", + ".modelhub_state/recovery_intents.jsonl", + ".modelhub_state/routing_intelligence.json", + ".modelhub_state/submission_exclusions.jsonl", + "ledger/submissions.jsonl", + "outcomes/submissions.jsonl", +) + + +class StateSyncError(RuntimeError): + pass + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _safe_text(value: Any, limit: int = 2000) -> str: + text = str(value or "") + for marker in ("xc-token", "authorization", "password", "token="): + if marker in text.lower(): + return "" + return text[:limit] + + +def _safe_config_vector(config: str) -> dict[str, Any]: + """Extract only bounded tuning values; never persist the original config.""" + vector: dict[str, Any] = {} + patterns = { + "gpuNum": r"\bgpu_num\s*:\s*['\"]?(\d+)", + "tensorParallel": r"(?:--tensor-parallel-size|-tp)\s*[, ]?\s*['\"]?(\d+)", + "maxModelLen": r"(?:--max-model-len|max_model_len|max_seq_len)\s*[: ,]+\s*['\"]?(\d+)", + "gpuMemoryUtilization": r"(?:--gpu-memory-utilization|gpu_memory_utilization)\s*[: ,]+\s*['\"]?([0-9.]+)", + } + for key, pattern in patterns.items(): + values = re.findall(pattern, config, flags=re.IGNORECASE) + if not values: + continue + try: + parsed = float(values[-1]) if "." in values[-1] else int(values[-1]) + except ValueError: + continue + vector[key] = parsed + for key, pattern in { + "dtype": r"(?:--dtype|dtype)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)", + "quantization": r"(?:--quantization|quantization)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)", + "loadFormat": r"(?:--load-format|load_format)\s*[: ,]+\s*['\"]?([A-Za-z0-9_-]+)", + }.items(): + values = re.findall(pattern, config, flags=re.IGNORECASE) + if values: + vector[key] = values[-1][:64] + return vector + + +def load_state_git_credentials(*dotenv_paths: Path) -> dict[str, str]: + paths = dotenv_paths or (Path(".env"), Path(__file__).resolve().parent.parent / ".env") + values = load_key_files(*paths) + + def choose(env_name: str, dotenv_name: str) -> str: + raw = os.getenv(env_name) or values.get(dotenv_name) or "" + return str(raw).strip().strip('"').strip("'") + + return { + "username": choose("MODELHUB_GIT_USERNAME", "modelhub_user_name"), + "email": choose("MODELHUB_GIT_EMAIL", "modelhub_user_email"), + "password": choose("MODELHUB_GIT_PASSWORD", "modelhub_user_password"), + } + + +def state_git_credentials_present(*dotenv_paths: Path) -> bool: + credentials = load_state_git_credentials(*dotenv_paths) + return all(credentials.values()) + + +class StateGitSync: + def __init__( + self, + *, + project_root: Path | str, + credentials: dict[str, str], + remote: str = DEFAULT_REMOTE, + branch: str = DEFAULT_BRANCH, + batch_size: int = DEFAULT_BATCH_SIZE, + retention_days: int = DEFAULT_RETENTION_DAYS, + history_depth: int = DEFAULT_HISTORY_DEPTH, + log_fn=None, + ) -> None: + self.project_root = Path(project_root).resolve() + self.credentials = dict(credentials) + self.remote = remote + self.branch = branch + self.batch_size = max(1, min(100, int(batch_size))) + self.retention_days = max(1, int(retention_days)) + self.history_depth = max(2, int(history_depth)) + self.log = log_fn or (lambda message: print(message, flush=True)) + self.writer_id = uuid.uuid4().hex + self.generation = 0 + self.healthy = False + self.last_error: str | None = None + self.last_sync_at: str | None = None + self._lock_handle = None + self._mutex = threading.Lock() + self._workspace: Path | None = None + self._askpass_dir: Path | None = None + self._askpass_path: Path | None = None + self._expected_remote_oid: str | None = None + + @property + def state_dir(self) -> Path: + return self.project_root / ".modelhub_state" + + @property + def intents_path(self) -> Path: + return self.state_dir / "recovery_intents.jsonl" + + @property + def active_tasks_path(self) -> Path: + return self.state_dir / "recovery_active_tasks.jsonl" + + @property + def readiness_path(self) -> Path: + return self.state_dir / "readiness.json" + + def acquire_process_lock(self) -> None: + self.state_dir.mkdir(parents=True, exist_ok=True) + path = self.state_dir / "state_sync.lock" + handle = path.open("a+", encoding="utf-8") + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + handle.close() + raise StateSyncError("another state-sync writer is already running") from exc + self._lock_handle = handle + + def close(self) -> None: + if self._lock_handle is not None: + try: + fcntl.flock(self._lock_handle.fileno(), fcntl.LOCK_UN) + finally: + self._lock_handle.close() + self._lock_handle = None + if self._workspace is not None: + shutil.rmtree(self._workspace.parent, ignore_errors=True) + self._workspace = None + if self._askpass_dir is not None: + shutil.rmtree(self._askpass_dir, ignore_errors=True) + self._askpass_dir = None + self._askpass_path = None + + def _git_environment(self) -> dict[str, str]: + if not all(self.credentials.get(key) for key in ("username", "email", "password")): + raise StateSyncError("missing ModelHub Git username, email or password") + if self._askpass_path is None: + directory = Path(tempfile.mkdtemp(prefix="modelhub-state-askpass-")) + script = directory / "askpass.py" + script.write_text( + "#!/usr/bin/env python3\n" + "import os, sys\n" + "prompt = ' '.join(sys.argv[1:]).lower()\n" + "key = 'MODELHUB_STATE_GIT_USERNAME' if 'username' in prompt else 'MODELHUB_STATE_GIT_PASSWORD'\n" + "print(os.environ.get(key, ''))\n", + encoding="utf-8", + ) + script.chmod(stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) + self._askpass_dir = directory + self._askpass_path = script + env = os.environ.copy() + env.update( + { + "GIT_ASKPASS": str(self._askpass_path), + "GIT_TERMINAL_PROMPT": "0", + "MODELHUB_STATE_GIT_USERNAME": self.credentials["username"], + "MODELHUB_STATE_GIT_PASSWORD": self.credentials["password"], + } + ) + return env + + def _git(self, *args: str, cwd: Path | None = None, check: bool = True) -> subprocess.CompletedProcess[str]: + result = subprocess.run( + ["git", *args], + cwd=str(cwd or self.project_root), + env=self._git_environment(), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=45, + check=False, + ) + if check and result.returncode != 0: + reason = _safe_text(result.stderr.strip() or result.stdout.strip() or f"git exited {result.returncode}") + raise StateSyncError(reason) + return result + + def _remote_oid(self) -> str | None: + result = self._git("ls-remote", "--heads", self.remote, self.branch) + line = result.stdout.strip().splitlines() + return line[0].split()[0] if line else None + + def _create_workspace(self, remote_oid: str | None) -> None: + parent = Path(tempfile.mkdtemp(prefix="modelhub-agent-state-")) + workspace = parent / "state" + if remote_oid: + self._git( + "clone", + "--quiet", + "--single-branch", + "--branch", + self.branch, + self.remote, + str(workspace), + cwd=parent, + ) + else: + workspace.mkdir(parents=True) + self._git("init", "--quiet", cwd=workspace) + self._git("checkout", "--orphan", self.branch, cwd=workspace) + self._git("remote", "add", "origin", self.remote, cwd=workspace) + self._git("config", "user.name", self.credentials["username"], cwd=workspace) + self._git("config", "user.email", self.credentials["email"], cwd=workspace) + self._workspace = workspace + self._expected_remote_oid = remote_oid + + def _validate_manifest(self, root: Path) -> dict[str, Any]: + manifest = read_json(root / "manifest.json") + if not isinstance(manifest, dict) or int(manifest.get("schemaVersion") or 0) != STATE_SCHEMA_VERSION: + raise StateSyncError("unsupported or missing state manifest") + checksums = manifest.get("checksums") or {} + if not isinstance(checksums, dict): + raise StateSyncError("invalid state manifest checksums") + for relative, expected in checksums.items(): + if relative not in STATE_ALLOWLIST and not str(relative).startswith("events/"): + raise StateSyncError(f"state manifest contains a non-allowlisted path: {relative}") + path = root / str(relative) + if not path.is_file() or _sha256_file(path) != str(expected): + raise StateSyncError(f"state checksum mismatch: {relative}") + return manifest + + def restore(self) -> bool: + with self._mutex: + try: + remote_oid = self._remote_oid() + self._create_workspace(remote_oid) + if remote_oid is not None: + try: + manifest = self._validate_manifest(self._workspace) + except Exception as current_error: + manifest = None + commits = self._git( + "rev-list", + f"--max-count={self.history_depth}", + "HEAD", + cwd=self._workspace, + ).stdout.splitlines()[1:] + for commit in commits: + self._git("checkout", "--quiet", "--detach", commit, cwd=self._workspace) + try: + manifest = self._validate_manifest(self._workspace) + self.log(f"[state-recovery] fallback_commit={commit[:12]} reason=checksum_recovery") + break + except Exception: + continue + if manifest is None: + raise current_error + self.generation = max(0, int(manifest.get("generation") or 0)) + for relative in STATE_ALLOWLIST: + source = self._workspace / relative + if not source.is_file(): + continue + destination = self.project_root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_name(f".{destination.name}.restore-{uuid.uuid4().hex}") + shutil.copy2(source, temporary) + os.replace(temporary, destination) + self.log( + f"[state-recovery] generation={self.generation} source=remote branch={self.branch} status=ok" + ) + else: + self.log(f"[state-recovery] source=local_bootstrap branch={self.branch} status=ok") + self.healthy = True + self.last_error = None + return True + except Exception as exc: + self.healthy = False + self.last_error = _safe_text(exc) + self.log(f"[state-recovery] status=failed reason={self.last_error}") + return False + + def retry_restore(self) -> bool: + if self._workspace is not None: + shutil.rmtree(self._workspace.parent, ignore_errors=True) + self._workspace = None + self._expected_remote_oid = None + return self.restore() + + def _event_files(self) -> list[Path]: + event_dir = self.state_dir / "events" + if not event_dir.exists(): + return [] + cutoff = (_utc_now() - timedelta(days=self.retention_days)).date() + result: list[Path] = [] + for path in sorted(event_dir.glob("*.jsonl")): + try: + event_day = datetime.strptime(path.stem, "%Y-%m-%d").date() + except ValueError: + continue + if event_day >= cutoff: + result.append(path) + else: + path.unlink(missing_ok=True) + return result + + def _append_event(self, event: dict[str, Any]) -> None: + now = _utc_now() + path = self.state_dir / "events" / f"{now.date().isoformat()}.jsonl" + existing = read_jsonl(path) + sanitized = {key: value for key, value in event.items() if key not in {"configParams", "token", "password"}} + sanitized["at"] = sanitized.get("at") or now.isoformat() + sanitized["eventId"] = sanitized.get("eventId") or uuid.uuid4().hex + if "reason" in sanitized: + sanitized["reason"] = _safe_text(sanitized["reason"]) + existing.append(sanitized) + write_jsonl(path, existing) + + @staticmethod + def _intent(candidate: dict[str, Any], batch_id: str) -> dict[str, Any]: + config = str(candidate.get("configParams") or "") + return { + "intentId": uuid.uuid4().hex, + "batchId": batch_id, + "status": "pending", + "createdAt": _utc_now().isoformat(), + "repoId": candidate.get("repoId"), + "modelAddress": candidate.get("modelAddress"), + "targetGpu": candidate.get("targetGpu"), + "taskType": candidate.get("taskType"), + "framework": candidate.get("framework"), + "lastModified": candidate.get("lastModified"), + "configSource": candidate.get("frameworkConfigSource") or "official", + "configFingerprint": hashlib.sha256(config.encode("utf-8")).hexdigest(), + "safeConfigVector": _safe_config_vector(config), + } + + def begin_batch(self, candidates: list[dict[str, Any]]) -> str | None: + if not self.healthy: + return None + batch_id = uuid.uuid4().hex + records = read_jsonl(self.intents_path) + intents = [self._intent(candidate, batch_id) for candidate in candidates] + records.extend(intents) + write_jsonl(self.intents_path, records) + for intent in intents: + self._append_event({**intent, "event": "submission_intent"}) + if not self.sync("intent"): + return None + return batch_id + + def finish_batch(self, batch_id: str, results: Iterable[dict[str, Any]]) -> bool: + records = read_jsonl(self.intents_path) + by_key = { + ( + str(item.get("repoId") or ""), + str(item.get("targetGpu") or ""), + str(item.get("framework") or ""), + ): item + for item in records + if item.get("batchId") == batch_id + } + for result in results: + candidate = result.get("candidate") or {} + key = ( + str(candidate.get("repoId") or ""), + str(candidate.get("targetGpu") or ""), + str(candidate.get("framework") or ""), + ) + intent = by_key.get(key) + if intent is None: + continue + intent["status"] = str(result.get("outcome") or "unknown") + intent["completedAt"] = _utc_now().isoformat() + intent["taskId"] = result.get("taskId") + intent["reason"] = _safe_text(result.get("reason")) if result.get("reason") else None + self._append_event({**intent, "event": "submission_result"}) + write_jsonl(self.intents_path, records) + return self.sync("result") + + def record_active_tasks(self, tasks: list[dict[str, Any]]) -> None: + sanitized: list[dict[str, Any]] = [] + for task in tasks: + sanitized.append( + { + key: task.get(key) + for key in ( + "accountKey", + "taskId", + "modelId", + "modelAddress", + "gpuType", + "targetGpu", + "taskType", + "modelTaskLevel", + "modelTaskLevelId", + "framework", + "configFingerprint", + "status", + "createTime", + "updateTime", + ) + if task.get(key) is not None + } + ) + write_jsonl(self.active_tasks_path, sanitized) + + def reconcile_active_tasks(self, tasks: list[dict[str, Any]]) -> dict[str, int]: + intents = read_jsonl(self.intents_path) + recoverable = [ + item + for item in intents + if str(item.get("status") or "") in {"pending", "submitted", "recovered_active"} + ] + by_route = { + ( + str(item.get("repoId") or ""), + str(item.get("targetGpu") or ""), + ): item + for item in recoverable + } + reconciled = 0 + enriched: list[dict[str, Any]] = [] + active_routes: set[tuple[str, str]] = set() + for task in tasks: + route = ( + str(task.get("modelId") or task.get("repoId") or ""), + str(task.get("gpuType") or task.get("targetGpu") or ""), + ) + active_routes.add(route) + matched = by_route.get(route) + copied = dict(task) + if matched is not None: + if not copied.get("framework"): + copied["framework"] = matched.get("framework") + copied["configFingerprint"] = matched.get("configFingerprint") + matched["status"] = "recovered_active" + matched["reconciledAt"] = _utc_now().isoformat() + if copied.get("taskId") is not None: + matched["taskId"] = copied.get("taskId") + reconciled += 1 + enriched.append(copied) + + now = _utc_now() + unresolved = 0 + outcome_by_task = { + str(item.get("taskId")): item + for item in read_jsonl(self.project_root / "outcomes/submissions.jsonl") + if item.get("taskId") is not None and item.get("outcome") in {"success", "failed", "policy_cancelled"} + } + for intent in recoverable: + route = (str(intent.get("repoId") or ""), str(intent.get("targetGpu") or "")) + if route in active_routes: + continue + terminal = outcome_by_task.get(str(intent.get("taskId") or "")) + if terminal is not None: + intent["status"] = str(terminal.get("outcome") or "terminal") + intent["completedAt"] = now.isoformat() + continue + created = str(intent.get("createdAt") or "") + try: + created_at = datetime.fromisoformat(created.replace("Z", "+00:00")) + except ValueError: + created_at = now + if created_at.tzinfo is None: + created_at = created_at.replace(tzinfo=timezone.utc) + if intent.get("status") == "pending" and now - created_at >= timedelta(hours=2): + intent["status"] = "orphan_unconfirmed" + intent["completedAt"] = now.isoformat() + elif intent.get("status") == "pending": + unresolved += 1 + retention_cutoff = now - timedelta(days=self.retention_days) + retained: list[dict[str, Any]] = [] + for intent in intents: + completed_text = intent.get("completedAt") + if not completed_text: + retained.append(intent) + continue + try: + completed_at = datetime.fromisoformat(str(completed_text).replace("Z", "+00:00")) + except ValueError: + retained.append(intent) + continue + if completed_at.tzinfo is None: + completed_at = completed_at.replace(tzinfo=timezone.utc) + if completed_at >= retention_cutoff: + retained.append(intent) + write_jsonl(self.intents_path, retained) + self.record_active_tasks(enriched) + return {"active": len(enriched), "reconciled": reconciled, "unresolved": unresolved} + + def _copy_snapshot(self) -> dict[str, str]: + assert self._workspace is not None + checksums: dict[str, str] = {} + for relative in STATE_ALLOWLIST: + source = self.project_root / relative + destination = self._workspace / relative + if not source.is_file(): + destination.unlink(missing_ok=True) + continue + destination.parent.mkdir(parents=True, exist_ok=True) + if relative == ".modelhub_state/market_intelligence.json": + payload = read_json(source) + + def strip_configs(value: Any) -> Any: + if isinstance(value, dict): + return { + key: strip_configs(item) + for key, item in value.items() + if key not in {"officialConfig", "configParams"} + } + if isinstance(value, list): + return [strip_configs(item) for item in value] + return value + + sanitized_market = strip_configs(payload) + if isinstance(sanitized_market, dict): + # A restored snapshot must refresh official configs before routing. + sanitized_market["frameworkUpdatedAt"] = None + write_json(destination, sanitized_market) + elif relative == "outcomes/submissions.jsonl": + sanitized_outcomes: list[dict[str, Any]] = [] + for row in read_jsonl(source): + sanitized_outcomes.append( + { + key: value + for key, value in row.items() + if not any( + marker in key.casefold() + for marker in ("url", "token", "cookie", "authorization", "configparams") + ) + } + ) + write_jsonl(destination, sanitized_outcomes) + else: + shutil.copy2(source, destination) + checksums[relative] = _sha256_file(destination) + for source in self._event_files(): + relative = f"events/{source.name}" + destination = self._workspace / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + checksums[relative] = _sha256_file(destination) + remote_event_dir = self._workspace / "events" + if remote_event_dir.exists(): + allowed = {Path(path).name for path in checksums if path.startswith("events/")} + for path in remote_event_dir.glob("*.jsonl"): + if path.name not in allowed: + path.unlink(missing_ok=True) + return checksums + + def sync(self, phase: str) -> bool: + with self._mutex: + try: + if self._workspace is None: + raise StateSyncError("state workspace is not initialized") + checksums = self._copy_snapshot() + self.generation += 1 + manifest = { + "schemaVersion": STATE_SCHEMA_VERSION, + "generation": self.generation, + "updatedAt": _utc_now().isoformat(), + "agentVersion": AGENT_VERSION, + "writerId": self.writer_id, + "phase": phase, + "checksums": checksums, + } + write_json(self._workspace / "manifest.json", manifest) + self._git("add", "--all", cwd=self._workspace) + diff = self._git("diff", "--cached", "--quiet", cwd=self._workspace, check=False) + if diff.returncode == 0: + self.healthy = True + return True + self._git("commit", "--quiet", "-m", f"state: generation {self.generation} ({phase})", cwd=self._workspace) + commit_count = int( + self._git("rev-list", "--count", "HEAD", cwd=self._workspace).stdout.strip() or "0" + ) + if commit_count > self.history_depth: + compact_branch = f"state-compact-{uuid.uuid4().hex[:8]}" + self._git("checkout", "--quiet", "--orphan", compact_branch, cwd=self._workspace) + self._git("rm", "--quiet", "--cached", "-r", ".", cwd=self._workspace, check=False) + self._git("add", "--all", cwd=self._workspace) + self._git( + "commit", + "--quiet", + "-m", + f"state: compacted generation {self.generation}", + cwd=self._workspace, + ) + lease = ( + f"--force-with-lease=refs/heads/{self.branch}:{self._expected_remote_oid}" + if self._expected_remote_oid + else f"--force-with-lease=refs/heads/{self.branch}:" + ) + self._git("push", "--quiet", lease, "origin", f"HEAD:refs/heads/{self.branch}", cwd=self._workspace) + self._expected_remote_oid = self._git("rev-parse", "HEAD", cwd=self._workspace).stdout.strip() + self.last_sync_at = manifest["updatedAt"] + self.last_error = None + self.healthy = True + self.log(f"[state-sync] generation={self.generation} phase={phase} status=ok") + return True + except Exception as exc: + self.healthy = False + self.last_error = _safe_text(exc) + self.log(f"[state-sync] phase={phase} status=failed reason={self.last_error}") + return False + + def write_readiness(self, *, ready: bool, reason: str | None = None, extra: dict[str, Any] | None = None) -> None: + body = { + "ready": bool(ready), + "reason": reason, + "updatedAt": _utc_now().isoformat(), + "pid": os.getpid(), + "stateSync": { + "healthy": self.healthy, + "generation": self.generation, + "lastSyncAt": self.last_sync_at, + "lastError": self.last_error, + }, + } + if extra: + body.update(extra) + write_json(self.readiness_path, body) diff --git a/modelhub_submmit_api/task_registry.py b/modelhub_submmit_api/task_registry.py index 0d286c5f..d8c297ce 100644 --- a/modelhub_submmit_api/task_registry.py +++ b/modelhub_submmit_api/task_registry.py @@ -42,6 +42,7 @@ TASK_SPECS: tuple[TaskSpec, ...] = ( TASK_SPEC_BY_TYPE = {task.task_type: task for task in TASK_SPECS} +DYNAMIC_TASK_TYPES: list[str] = [] TASK_TYPE_BY_MODEL_TASK_LEVEL_ID = { @@ -91,7 +92,55 @@ def task_type_from_history_task(task: dict) -> str | None: def all_task_types() -> list[str]: - return [task.task_type for task in TASK_SPECS] + return [task.task_type for task in TASK_SPECS] + list(DYNAMIC_TASK_TYPES) + + +def register_dynamic_task_types(task_types: list[str]) -> list[str]: + """Register API task identifiers that can be sourced by the same ModelScope tag.""" + added: list[str] = [] + next_priority = max(spec.priority for spec in TASK_SPEC_BY_TYPE.values()) + 10 + for raw in task_types: + task_type = str(raw or "").strip() + if not task_type or task_type in TASK_SPEC_BY_TYPE: + continue + if any(not (character.isascii() and (character.isalnum() or character in "_-")) for character in task_type): + continue + TASK_SPEC_BY_TYPE[task_type] = TaskSpec( + task_type=task_type, + modality="generic", + pipeline_tags=(task_type,), + priority=next_priority, + ) + next_priority += 10 + DYNAMIC_TASK_TYPES.append(task_type) + added.append(task_type) + return added + + +def register_dynamic_task_route(task_type: str, pipeline_tag: str) -> None: + task_type = str(task_type or "").strip() + pipeline_tag = str(pipeline_tag or "").strip().lower() + if not task_type or not pipeline_tag: + return + existing = TASK_SPEC_BY_TYPE.get(task_type) + if existing is not None: + if pipeline_tag not in existing.pipeline_tags: + TASK_SPEC_BY_TYPE[task_type] = TaskSpec( + task_type=existing.task_type, + modality=existing.modality, + pipeline_tags=(*existing.pipeline_tags, pipeline_tag), + priority=existing.priority, + ) + return + register_dynamic_task_types([task_type]) + created = TASK_SPEC_BY_TYPE.get(task_type) + if created is not None: + TASK_SPEC_BY_TYPE[task_type] = TaskSpec( + task_type=created.task_type, + modality=created.modality, + pipeline_tags=(pipeline_tag,), + priority=created.priority, + ) def pipeline_tags_for_task_types(task_types: list[str]) -> list[str]: @@ -106,7 +155,7 @@ def pipeline_tags_for_task_types(task_types: list[str]) -> list[str]: def task_specs_for_model(model: HFModelSummary) -> list[TaskSpec]: pipeline_tag = (model.pipeline_tag or "").strip().lower() - return [task for task in TASK_SPECS if pipeline_tag in task.pipeline_tags] + return [task for task in TASK_SPEC_BY_TYPE.values() if pipeline_tag in task.pipeline_tags] def compatible_text_generation_frameworks( @@ -170,6 +219,9 @@ def compatible_frameworks_for_task( return ["diffusers"] raise ValueError(f"No compatible diffusers template found for {target_gpu}") + if inspection.has_standard_weights and supported_frameworks: + return sorted(supported_frameworks) + raise ValueError(f"Unsupported task type for auto framework selection: {task_type}") diff --git a/modelhub_submmit_api/version.py b/modelhub_submmit_api/version.py index 3bb782ab..45744503 100644 --- a/modelhub_submmit_api/version.py +++ b/modelhub_submmit_api/version.py @@ -1 +1 @@ -AGENT_VERSION = "2026.08.12.6" +AGENT_VERSION = "2026.08.15.1" diff --git a/requirements.txt b/requirements.txt index 59567f62..e0d33d1a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1 @@ -# The runner currently uses only the Python standard library. +PyYAML>=6.0.2,<7 diff --git a/tests/test_agent_entrypoint.py b/tests/test_agent_entrypoint.py index 636c34b1..cba11fd7 100644 --- a/tests/test_agent_entrypoint.py +++ b/tests/test_agent_entrypoint.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util import os import unittest +import tempfile from pathlib import Path from unittest.mock import patch @@ -28,6 +29,16 @@ class HostedAgentEntrypointTests(unittest.TestCase): command = ENTRYPOINT._worker_command() self.assertEqual(["--max-submits-per-run", "0"], command[-2:]) + self.assertIn("--state-sync", command) + + def test_readiness_file_is_separate_from_liveness(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + path = Path(temporary_dir) / "readiness.json" + path.write_text('{"ready": false, "reason": "state_sync_unhealthy"}', encoding="utf-8") + with patch.object(ENTRYPOINT, "READINESS_PATH", path): + readiness = ENTRYPOINT._readiness() + self.assertFalse(readiness["ready"]) + self.assertEqual("state_sync_unhealthy", readiness["reason"]) if __name__ == "__main__": diff --git a/tests/test_super_agent.py b/tests/test_super_agent.py new file mode 100644 index 00000000..f3d1e843 --- /dev/null +++ b/tests/test_super_agent.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import subprocess +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import patch + +import sys + +ROOT = Path(__file__).resolve().parents[1] +MODULE_ROOT = ROOT / "modelhub_submmit_api" +if str(MODULE_ROOT) not in sys.path: + sys.path.insert(0, str(MODULE_ROOT)) + +from common import read_jsonl, write_json, write_jsonl # noqa: E402 +from config_optimizer import SafeConfigOptimizer # noqa: E402 +from hf_discovery import HuggingFaceDiscovery, parse_model_card_front_matter # noqa: E402 +from official_capabilities import OfficialCapabilityRegistry # noqa: E402 +from routing_engine import SuccessFirstRoutingEngine # noqa: E402 +from state_sync import StateGitSync # noqa: E402 + + +class OfficialClient: + def __init__(self, *, fail_catalog: bool = False) -> None: + self.fail_catalog = fail_catalog + + def list_machine_info(self): # noqa: ANN201 + if self.fail_catalog: + raise RuntimeError("offline") + return [ + {"gpuType": "gpu-fast", "canVerify": True, "maxConcurrentTasks": 2}, + {"gpuType": "gpu-disabled", "canVerify": False, "maxConcurrentTasks": 8}, + ] + + def list_task_levels(self): # noqa: ANN201 + return {"data": [{"taskType": "text-generation"}, {"taskType": "new-task"}]} + + def list_model_task_types(self, target_gpu, model_address): # noqa: ANN001, ANN201 + del target_gpu, model_address + return {"data": [{"taskType": "text-generation"}]} + + +class SuperAgentTests(unittest.TestCase): + def test_official_registry_discovers_catalog_and_exact_model_routes(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None) + client = OfficialClient() + registry.prepare( + client, + fallback_gpus=["legacy"], + task_types=["text-generation"], + now=datetime(2026, 8, 15, tzinfo=timezone.utc), + ) + self.assertTrue(registry.ready) + self.assertEqual(["gpu-fast"], registry.eligible_gpus()) + self.assertEqual( + ["text-generation"], + registry.task_types_for( + client, + model_address="https://modelscope.cn/models/owner/model", + model_last_modified="2026-08-15T00:00:00+00:00", + gpu="gpu-fast", + ), + ) + + def test_official_registry_fails_closed_without_catalog_cache(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + registry = OfficialCapabilityRegistry(Path(temporary_dir) / "official.json", log_fn=lambda _: None) + registry.prepare( + OfficialClient(fail_catalog=True), + fallback_gpus=["legacy"], + task_types=["text-generation"], + now=datetime(2026, 8, 15, tzinfo=timezone.utc), + ) + self.assertFalse(registry.ready) + self.assertEqual("critical_official_signal_unavailable", registry.pause_reason) + + def test_success_band_beats_shorter_queue(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + engine = SuccessFirstRoutingEngine(Path(temporary_dir) / "routing.json", log_fn=lambda _: None) + candidates = [ + { + "repoId": "owner/model-a", + "targetGpu": "reliable", + "framework": "vllm", + "taskType": "text-generation", + "frameworkMarketSamples": 1000, + "frameworkMarketSuccessRate": 0.9, + "queueBacklogHours": 12, + }, + { + "repoId": "owner/model-b", + "targetGpu": "fast", + "framework": "vllm", + "taskType": "text-generation", + "frameworkMarketSamples": 1000, + "frameworkMarketSuccessRate": 0.4, + "queueBacklogHours": 0.25, + }, + ] + ordered = engine.order_candidates(candidates) + self.assertEqual("reliable", ordered[0]["targetGpu"]) + self.assertGreater(ordered[0]["routingSuccessBand"], ordered[1]["routingSuccessBand"]) + + def test_modelscope_metadata_and_model_card_lineage_are_structured(self) -> None: + item = { + "id": "owner/model", + "downloads": 123, + "params": 7_000_000_000, + "file_size": 14_000_000_000, + "tags": ["qwen", "chat"], + "tasks": ["text-generation"], + "license": "apache-2.0", + "likes": 9, + } + model = HuggingFaceDiscovery._parse_model( + item, + fallback_pipeline_tag="text-generation", + min_downloads=0, + ) + self.assertIsNotNone(model) + assert model is not None + self.assertEqual(7_000_000_000, model.params) + self.assertEqual(("qwen", "chat"), model.tags) + metadata = parse_model_card_front_matter( + "---\nbase_model: Qwen/base\nframeworks:\n - transformers\ntasks:\n - text-generation\n---\nbody" + ) + self.assertEqual("Qwen/base", metadata["base_model"]) + self.assertEqual(["transformers"], metadata["frameworks"]) + + def test_state_branch_round_trip_persists_intent_without_config_or_secret(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + remote = root / "remote.git" + project = root / "project" + restored_project = root / "restored" + project.mkdir() + restored_project.mkdir() + subprocess.run(["git", "init", "--bare", str(remote)], check=True, stdout=subprocess.DEVNULL) + write_json(project / ".modelhub_state" / "account_capacity.json", {"version": 1}) + credentials = {"username": "tester", "email": "tester@example.com", "password": "secret-value"} + manager = StateGitSync( + project_root=project, + credentials=credentials, + remote=str(remote), + log_fn=lambda _: None, + ) + manager.acquire_process_lock() + self.assertTrue(manager.restore()) + batch_id = manager.begin_batch( + [ + { + "repoId": "owner/model", + "modelAddress": "https://modelscope.cn/models/owner/model", + "targetGpu": "gpu-a", + "taskType": "text-generation", + "framework": "vllm", + "configParams": "password: must-not-be-copied", + } + ] + ) + self.assertIsNotNone(batch_id) + manager.close() + + restored = StateGitSync( + project_root=restored_project, + credentials=credentials, + remote=str(remote), + log_fn=lambda _: None, + ) + restored.acquire_process_lock() + self.assertTrue(restored.restore()) + intents = read_jsonl(restored_project / ".modelhub_state" / "recovery_intents.jsonl") + self.assertEqual("owner/model", intents[0]["repoId"]) + state_text = "\n".join( + path.read_text(encoding="utf-8") + for path in restored._workspace.rglob("*") + if path.is_file() and ".git" not in path.parts + ) + self.assertNotIn("secret-value", state_text) + self.assertNotIn("must-not-be-copied", state_text) + restored.close() + + def test_failed_intent_push_returns_no_batch_id(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + manager = StateGitSync( + project_root=Path(temporary_dir), + credentials={"username": "u", "email": "e@example.com", "password": "p"}, + remote="unused", + log_fn=lambda _: None, + ) + manager.healthy = True + with patch.object(manager, "sync", return_value=False): + self.assertIsNone( + manager.begin_batch( + [ + { + "repoId": "owner/model", + "targetGpu": "gpu", + "taskType": "text-generation", + "framework": "vllm", + "configParams": "safe", + } + ] + ) + ) + + def test_config_patch_requires_repeated_cross_model_success(self) -> None: + with tempfile.TemporaryDirectory() as temporary_dir: + root = Path(temporary_dir) + intents = [] + outcomes = [] + for index in range(5): + intents.append( + { + "taskId": str(index), + "taskType": "text-generation", + "targetGpu": "gpu-a", + "framework": "vllm", + "repoId": f"owner/model-{index % 2}", + "configFingerprint": "proven", + "safeConfigVector": {"gpuNum": 1, "tensorParallel": 1}, + } + ) + outcomes.append({"taskId": str(index), "outcome": "success"}) + write_jsonl(root / "intents.jsonl", intents) + write_jsonl(root / "outcomes.jsonl", outcomes) + optimizer = SafeConfigOptimizer( + intents_path=root / "intents.jsonl", + outcomes_path=root / "outcomes.jsonl", + ) + config, metadata = optimizer.optimize( + task_type="text-generation", + target_gpu="gpu-a", + framework="vllm", + official_config="framework: vllm\nsut_config:\n gpu_num: 2\nref_config:\n gpu_num: 2\n", + official_lower_bound=0.40, + ) + self.assertTrue(metadata["applied"]) + self.assertNotIn("gpu_num: 2", config) + + +if __name__ == "__main__": + unittest.main()