feat: add durable success-first modelhub agent
This commit is contained in:
5
.env
5
.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
|
||||
dashscope = sk-ws-H.EEMMMLP.i9CD.MEYCIQCXgmgQJ8LfF1m-oBT4ogqc6eD8ahI1BokpJUjD4mlQqAIhAPNH_jFbhJS7fZaufHWCCKCF4Ty8HsP3JmMhvyhSpYhm
|
||||
modelhub_user_name = CoolBoy
|
||||
modelhub_user_email = 2269097679@qq.com
|
||||
modelhub_user_password = Woshixzy666!
|
||||
78
README.md
78
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
|
||||
```
|
||||
|
||||
23
main.py
23
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
|
||||
|
||||
120
modelhub_submmit_api/config_optimizer.py
Normal file
120
modelhub_submmit_api/config_optimizer.py
Normal file
@@ -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"],
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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] = {}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {}),
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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":
|
||||
|
||||
176
modelhub_submmit_api/official_capabilities.py
Normal file
176
modelhub_submmit_api/official_capabilities.py
Normal file
@@ -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 []),
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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": {
|
||||
|
||||
192
modelhub_submmit_api/routing_engine.py
Normal file
192
modelhub_submmit_api/routing_engine.py
Normal file
@@ -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 {}),
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
674
modelhub_submmit_api/state_sync.py
Normal file
674
modelhub_submmit_api/state_sync.py
Normal file
@@ -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 "<redacted>"
|
||||
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)
|
||||
@@ -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}")
|
||||
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
AGENT_VERSION = "2026.08.12.6"
|
||||
AGENT_VERSION = "2026.08.15.1"
|
||||
|
||||
@@ -1 +1 @@
|
||||
# The runner currently uses only the Python standard library.
|
||||
PyYAML>=6.0.2,<7
|
||||
|
||||
@@ -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__":
|
||||
|
||||
246
tests/test_super_agent.py
Normal file
246
tests/test_super_agent.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user