966 lines
40 KiB
Python
966 lines
40 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import os
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from common import format_modelhub_datetime, parse_datetime, read_json, runtime_instance_id, utc_now, write_json
|
||
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
||
from http_json import HttpJsonError, JsonHttpClient
|
||
|
||
|
||
class ModelHubAPIError(RuntimeError):
|
||
def __init__(self, message: str, *, code: int | None = None, payload: Any = None) -> None:
|
||
super().__init__(message)
|
||
self.code = code
|
||
self.payload = payload
|
||
|
||
|
||
class OldModelQueuePolicyError(ModelHubAPIError):
|
||
"""No account may accept an older model under the per-account queue policy."""
|
||
|
||
|
||
def parse_model_submission_precheck(payload: Any) -> dict[str, Any]:
|
||
"""Validate the community lookup response instead of failing open."""
|
||
if not isinstance(payload, dict):
|
||
raise ModelHubAPIError("Community model precheck returned a non-object response", payload=payload)
|
||
data = payload.get("data")
|
||
if not isinstance(data, dict) or "verifyResult" not in data:
|
||
raise ModelHubAPIError("Community model precheck response is incomplete", payload=payload)
|
||
verify_result = data.get("verifyResult") or {}
|
||
if not isinstance(verify_result, dict):
|
||
raise ModelHubAPIError("Community model precheck verifyResult is invalid", payload=payload)
|
||
|
||
return {
|
||
"isInDB": data.get("isInDB") is True,
|
||
"processedGpus": set(str(gpu) for gpu in verify_result),
|
||
}
|
||
|
||
|
||
class ModelHubClient:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
token: str | None = None,
|
||
base_url: str = "https://modelhub.org.cn",
|
||
timeout: int = 30,
|
||
retries: int = 2,
|
||
http_client: JsonHttpClient | None = None,
|
||
) -> None:
|
||
self.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or EMBEDDED_MODELHUB_XC_TOKEN
|
||
if not self.token and http_client is None:
|
||
raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN or XC_TOKEN.")
|
||
default_headers = {"Xc-Token": self.token} if self.token else {}
|
||
self.http_client = http_client or JsonHttpClient(
|
||
base_url=base_url,
|
||
default_headers=default_headers,
|
||
timeout=timeout,
|
||
retries=retries,
|
||
)
|
||
|
||
def search_by_model_id(self, model_id: str, *, force_refresh: bool = False) -> dict[str, Any]:
|
||
del force_refresh
|
||
return self._request("GET", "/api/computility/models/search-by-model-id", query={"modelId": model_id})
|
||
|
||
def model_submission_precheck(self, model_id: str, *, force_refresh: bool = False) -> dict[str, Any]:
|
||
return parse_model_submission_precheck(self.search_by_model_id(model_id, force_refresh=force_refresh))
|
||
|
||
def is_model_processed_for_gpu(self, model_id: str, target_gpu: str) -> bool:
|
||
gpu_result = self.get_verify_result_map(model_id).get(target_gpu)
|
||
if gpu_result is None:
|
||
return False
|
||
return "result" in gpu_result or "records" in gpu_result
|
||
|
||
def get_verify_result_map(self, model_id: str) -> dict[str, Any]:
|
||
payload = self.search_by_model_id(model_id)
|
||
return ((payload.get("data") or {}).get("verifyResult") or {})
|
||
|
||
def processed_gpus_for_model(self, model_id: str) -> set[str]:
|
||
return set(self.model_submission_precheck(model_id)["processedGpus"])
|
||
|
||
def list_tasks_page(
|
||
self,
|
||
*,
|
||
current: int = 1,
|
||
page_size: int = 50,
|
||
only_mine: bool = True,
|
||
begin_time: datetime | str | None = None,
|
||
end_time: datetime | str | None = None,
|
||
gpu_type: str | None = None,
|
||
model_id: str | None = None,
|
||
status: str | None = None,
|
||
verify_result: int | None = None,
|
||
) -> dict[str, Any]:
|
||
def query_time(value: datetime | str | None) -> str | None:
|
||
if isinstance(value, str):
|
||
return value
|
||
return format_modelhub_datetime(value) if value else None
|
||
|
||
return self._request(
|
||
"GET",
|
||
"/api/adapt/task/page",
|
||
query={
|
||
"current": current,
|
||
"pageSize": page_size,
|
||
"onlyMine": str(only_mine).lower(),
|
||
"beginTime": query_time(begin_time),
|
||
"endTime": query_time(end_time),
|
||
"gpuType": gpu_type,
|
||
"modelId": model_id,
|
||
"status": status,
|
||
"verifyResult": verify_result,
|
||
},
|
||
)
|
||
|
||
def list_machine_info(self) -> list[dict[str, Any]]:
|
||
payload = self._request("GET", "/api/computility/power/machine/list/machine-info")
|
||
data = payload.get("data") or []
|
||
if not isinstance(data, list):
|
||
raise ModelHubAPIError("Machine info response is invalid", payload=payload)
|
||
return [item for item in data if isinstance(item, dict)]
|
||
|
||
def list_framework_stats(self, task_type: str, target_gpu: str) -> list[dict[str, Any]]:
|
||
payload = self._request(
|
||
"GET",
|
||
"/api/computility/driver/images/frameworks",
|
||
query={"taskType": task_type, "gpuTypeName": target_gpu},
|
||
)
|
||
data = payload.get("data") or []
|
||
if not isinstance(data, list):
|
||
raise ModelHubAPIError("Framework statistics response is invalid", payload=payload)
|
||
return [item for item in data if isinstance(item, dict)]
|
||
|
||
def get_build_config(self, task_type: str, target_gpu: str, framework: str) -> str:
|
||
payload = self._request(
|
||
"POST",
|
||
"/api/adapt/task/build-config",
|
||
query={"taskType": task_type, "gpuType": target_gpu, "framework": framework},
|
||
)
|
||
data = payload.get("data")
|
||
if not isinstance(data, str) or not data.strip():
|
||
raise ModelHubAPIError("Official build config response is invalid", payload=payload)
|
||
return data
|
||
|
||
def list_tasks(
|
||
self,
|
||
*,
|
||
page_size: int = 50,
|
||
only_mine: bool = True,
|
||
begin_time: datetime | None = None,
|
||
end_time: datetime | None = None,
|
||
gpu_type: str | None = None,
|
||
model_id: str | None = None,
|
||
status: str | None = None,
|
||
verify_result: int | None = None,
|
||
max_records: int = 0,
|
||
) -> list[dict[str, Any]]:
|
||
current = 1
|
||
records: list[dict[str, Any]] = []
|
||
while True:
|
||
page = self.list_tasks_page(
|
||
current=current,
|
||
page_size=page_size,
|
||
only_mine=only_mine,
|
||
begin_time=begin_time,
|
||
end_time=end_time,
|
||
gpu_type=gpu_type,
|
||
model_id=model_id,
|
||
status=status,
|
||
verify_result=verify_result,
|
||
)
|
||
page_data = page.get("data") or {}
|
||
page_records = page_data.get("records") or []
|
||
records.extend(page_records)
|
||
if max_records > 0 and len(records) >= max_records:
|
||
return records[:max_records]
|
||
pages = int(page_data.get("pages") or 0)
|
||
if pages <= current or not page_records:
|
||
break
|
||
current += 1
|
||
return records
|
||
|
||
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||
return self._request("POST", "/api/adapt/task/add", data=payload)
|
||
|
||
def stop_tasks(self, task_ids: list[int | str]) -> dict[str, Any]:
|
||
"""Stop active validation tasks owned by this authenticated account."""
|
||
normalized: list[int] = []
|
||
seen: set[int] = set()
|
||
for task_id in task_ids:
|
||
try:
|
||
numeric_id = int(str(task_id).strip())
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValueError(f"Invalid ModelHub task ID: {task_id!r}") from exc
|
||
if numeric_id <= 0:
|
||
raise ValueError(f"Invalid ModelHub task ID: {task_id!r}")
|
||
if numeric_id not in seen:
|
||
seen.add(numeric_id)
|
||
normalized.append(numeric_id)
|
||
if not normalized:
|
||
raise ValueError("At least one ModelHub task ID is required")
|
||
return self._request(
|
||
"PUT",
|
||
"/api/async/task/stop-create-contest-task",
|
||
data={"taskIds": normalized},
|
||
)
|
||
|
||
def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
|
||
recent_tasks = self.list_tasks(
|
||
page_size=20,
|
||
only_mine=True,
|
||
begin_time=submitted_after - timedelta(hours=1),
|
||
end_time=submitted_after + timedelta(hours=6),
|
||
gpu_type=gpu_type,
|
||
model_id=model_id,
|
||
)
|
||
if not recent_tasks:
|
||
return None
|
||
recent_tasks.sort(key=lambda task: parse_datetime(task.get("updateTime")) or submitted_after, reverse=True)
|
||
return str(recent_tasks[0].get("taskId")) if recent_tasks[0].get("taskId") is not None else None
|
||
|
||
def count_active_tasks(
|
||
self,
|
||
*,
|
||
page_size: int = 100,
|
||
only_mine: bool = True,
|
||
begin_time: datetime | None = None,
|
||
end_time: datetime | None = None,
|
||
max_count: int | None = None,
|
||
) -> int:
|
||
"""Count active tasks with early pagination termination.
|
||
|
||
Unlike list_tasks() which fetches all pages, this method stops
|
||
fetching pages as soon as max_count active tasks are found.
|
||
"""
|
||
max_count_int: int | None = max_count
|
||
if max_count_int is not None:
|
||
max_count_int = max(0, int(max_count_int))
|
||
if max_count_int == 0:
|
||
return 0
|
||
|
||
active_count = 0
|
||
current = 1
|
||
while True:
|
||
page = self.list_tasks_page(
|
||
current=current,
|
||
page_size=page_size,
|
||
only_mine=only_mine,
|
||
begin_time=begin_time,
|
||
end_time=end_time,
|
||
)
|
||
page_data = page.get("data") or {}
|
||
page_records = page_data.get("records") or []
|
||
|
||
for task in page_records:
|
||
if is_active_task(task):
|
||
active_count += 1
|
||
if max_count_int is not None and active_count >= max_count_int:
|
||
return active_count
|
||
|
||
pages = int(page_data.get("pages") or 0)
|
||
if pages <= current or not page_records:
|
||
break
|
||
current += 1
|
||
|
||
return active_count
|
||
|
||
def _request(
|
||
self,
|
||
method: str,
|
||
path: str,
|
||
*,
|
||
query: dict[str, Any] | None = None,
|
||
data: dict[str, Any] | None = None,
|
||
) -> dict[str, Any]:
|
||
try:
|
||
payload = self.http_client.request_json(method, path, query=query, data=data)
|
||
except HttpJsonError as exc:
|
||
if isinstance(exc.payload, dict) and "message" in exc.payload:
|
||
raise ModelHubAPIError(exc.payload["message"], payload=exc.payload, code=exc.status_code) from exc
|
||
raise ModelHubAPIError(str(exc), payload=exc.payload, code=exc.status_code) from exc
|
||
if not isinstance(payload, dict):
|
||
raise ModelHubAPIError("Unexpected API response shape", payload=payload)
|
||
code = payload.get("code")
|
||
if code != 0:
|
||
raise ModelHubAPIError(payload.get("message") or "ModelHub API request failed", code=code, payload=payload)
|
||
return payload
|
||
|
||
|
||
ACTIVE_TASK_STATUSES = {
|
||
"waiting",
|
||
"running",
|
||
"processing",
|
||
"queued",
|
||
"pending",
|
||
"validating",
|
||
"submitting",
|
||
"initializing",
|
||
}
|
||
|
||
TERMINAL_TASK_STATUSES = {
|
||
"success",
|
||
"failed",
|
||
"error",
|
||
"cancelled",
|
||
"canceled",
|
||
"rejected",
|
||
"timeout",
|
||
"completed",
|
||
"complete",
|
||
"done",
|
||
}
|
||
|
||
|
||
def is_active_task(task: dict[str, Any]) -> bool:
|
||
status = str(task.get("status") or "").strip().lower()
|
||
if not status:
|
||
return False
|
||
if status in ACTIVE_TASK_STATUSES:
|
||
return True
|
||
if status in TERMINAL_TASK_STATUSES:
|
||
return False
|
||
return True
|
||
|
||
|
||
CAPACITY_ERROR_MARKERS = (
|
||
"达到上限",
|
||
"达上限",
|
||
"任务数量已达",
|
||
"队列已满",
|
||
"queue is full",
|
||
"queue full",
|
||
"capacity",
|
||
"too many active",
|
||
"active task limit",
|
||
)
|
||
|
||
DUPLICATE_SUBMISSION_MARKERS = (
|
||
"正在验证中",
|
||
"请勿重复提交",
|
||
"重复提交",
|
||
"already validating",
|
||
"already being validated",
|
||
"already in progress",
|
||
)
|
||
|
||
MODEL_UNIQUENESS_ERROR_MARKERS = (
|
||
"模型唯一性检查",
|
||
"唯一性检查没有通过",
|
||
"模型已存在",
|
||
"model uniqueness",
|
||
"uniqueness check",
|
||
"duplicate model",
|
||
"model already exists",
|
||
)
|
||
|
||
DEFAULT_CAPACITY_STATE_PATH = Path(".modelhub_state/account_capacity.json")
|
||
|
||
|
||
def is_capacity_error(error: ModelHubAPIError) -> bool:
|
||
if error.code in {409, 429}:
|
||
return True
|
||
message = str(error).strip().lower()
|
||
if isinstance(error.payload, dict):
|
||
message = f"{message} {error.payload.get('message') or ''}".lower()
|
||
return any(marker in message for marker in CAPACITY_ERROR_MARKERS)
|
||
|
||
|
||
def is_duplicate_submission_error(error: ModelHubAPIError) -> bool:
|
||
message = str(error).strip().lower()
|
||
if isinstance(error.payload, dict):
|
||
message = f"{message} {error.payload.get('message') or ''}".lower()
|
||
return any(marker in message for marker in DUPLICATE_SUBMISSION_MARKERS)
|
||
|
||
|
||
def is_model_uniqueness_error(error: ModelHubAPIError) -> bool:
|
||
message = str(error).strip().lower()
|
||
if isinstance(error.payload, dict):
|
||
message = f"{message} {error.payload.get('message') or ''}".lower()
|
||
return any(marker in message for marker in MODEL_UNIQUENESS_ERROR_MARKERS)
|
||
|
||
|
||
class ModelHubClientPool:
|
||
def __init__(
|
||
self,
|
||
clients: list[ModelHubClient],
|
||
*,
|
||
active_task_cap: int | None = None,
|
||
active_counts_ttl: float | None = None,
|
||
reservation_ttl: float | None = None,
|
||
instance_id: str | None = None,
|
||
capacity_probe_interval_cycles: int = 3,
|
||
capacity_probe_cooldown_cycles: int = 3,
|
||
capacity_state_path: Path | str | None = None,
|
||
recent_model_reserve_slots: int | None = None,
|
||
recent_model_days: int | None = None,
|
||
) -> None:
|
||
if not clients:
|
||
raise ValueError("At least one ModelHub client is required")
|
||
self.clients = clients
|
||
configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
|
||
self.active_task_cap = max(1, int(configured_cap))
|
||
configured_recent_reserve = (
|
||
recent_model_reserve_slots
|
||
if recent_model_reserve_slots is not None
|
||
else os.getenv("MODELHUB_RECENT_MODEL_RESERVE_SLOTS", "5")
|
||
)
|
||
configured_recent_days = (
|
||
recent_model_days
|
||
if recent_model_days is not None
|
||
else os.getenv("MODELHUB_RECENT_MODEL_DAYS", "7")
|
||
)
|
||
self.recent_model_reserve_slots = max(0, int(configured_recent_reserve))
|
||
self.recent_model_days = max(1, int(configured_recent_days))
|
||
self._capacity_state_path = Path(capacity_state_path) if capacity_state_path else None
|
||
self._account_keys = [self._account_key(client, index) for index, client in enumerate(clients)]
|
||
self._account_caps = self._load_account_caps(self.active_task_cap)
|
||
self._capacity_probe_interval_cycles = max(0, int(capacity_probe_interval_cycles))
|
||
self._capacity_probe_cooldown_cycles = max(1, int(capacity_probe_cooldown_cycles))
|
||
self._capacity_probe_cycle = 0
|
||
self._capacity_probe_enabled = False
|
||
self._capacity_probe_attempted: set[int] = set()
|
||
self._capacity_probe_cooldown_until: list[int] = [0 for _ in clients]
|
||
configured_ttl = active_counts_ttl if active_counts_ttl is not None else os.getenv("MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS", "15")
|
||
configured_reservation_ttl = (
|
||
reservation_ttl
|
||
if reservation_ttl is not None
|
||
else os.getenv("MODELHUB_AGENT_RESERVATION_TTL_SECONDS", "120")
|
||
)
|
||
self._active_counts_ttl = max(1.0, float(configured_ttl))
|
||
self._reservation_ttl = max(self._active_counts_ttl * 2, float(configured_reservation_ttl))
|
||
self._remote_counts: list[int] = [0 for _ in clients]
|
||
# Old-model admission is fail-closed per account. A configured or
|
||
# persisted capacity is not enough without a successful active-count
|
||
# read for the current scheduling window.
|
||
self._count_known: list[bool] = [False for _ in clients]
|
||
self._active_refresh_at: float = 0.0
|
||
self._counts_initialized = False
|
||
self._state_lock = threading.Lock()
|
||
self._refresh_lock = threading.Lock()
|
||
self._reservations: list[dict[int, dict[str, Any]]] = [{} for _ in clients]
|
||
self._reservation_sequence = 0
|
||
identity = instance_id or runtime_instance_id()
|
||
identity_hash = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||
self._selection_cursor = int(identity_hash[:12], 16) % len(clients)
|
||
self._verify_cache: dict[str, tuple[float, dict[str, Any]]] = {}
|
||
self._verify_cache_ttl = max(1.0, float(os.getenv("MODELHUB_AGENT_VERIFY_CACHE_TTL_SECONDS", "900")))
|
||
# Single reader client to avoid fanout on read operations
|
||
self._reader = clients[0]
|
||
|
||
@staticmethod
|
||
def _account_key(client: ModelHubClient, index: int) -> str:
|
||
token = str(getattr(client, "token", "") or "")
|
||
identity = token if token else f"account-index:{index}"
|
||
return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20]
|
||
|
||
def _load_account_caps(self, default_cap: int) -> list[int]:
|
||
stored: dict[str, Any] = {}
|
||
if self._capacity_state_path is not None:
|
||
try:
|
||
payload = read_json(self._capacity_state_path)
|
||
if isinstance(payload, dict):
|
||
stored = payload.get("accounts") or {}
|
||
except (FileNotFoundError, ValueError):
|
||
pass
|
||
return [
|
||
max(1, int((stored.get(key) or {}).get("knownCap") or default_cap))
|
||
for key in self._account_keys
|
||
]
|
||
|
||
def _persist_account_caps(self) -> None:
|
||
if self._capacity_state_path is None:
|
||
return
|
||
with self._state_lock:
|
||
payload = {
|
||
"version": 1,
|
||
"updatedAt": datetime.now().astimezone().isoformat(),
|
||
"accounts": {
|
||
key: {"accountIndex": index + 1, "knownCap": self._account_caps[index]}
|
||
for index, key in enumerate(self._account_keys)
|
||
},
|
||
}
|
||
write_json(self._capacity_state_path, payload)
|
||
|
||
def configure_capacity_probe(self, cycle_number: int) -> None:
|
||
cycle_number = max(0, int(cycle_number))
|
||
with self._state_lock:
|
||
if cycle_number != self._capacity_probe_cycle:
|
||
self._capacity_probe_attempted.clear()
|
||
self._capacity_probe_cycle = cycle_number
|
||
self._capacity_probe_enabled = (
|
||
self._capacity_probe_interval_cycles > 0
|
||
and cycle_number > 0
|
||
and cycle_number % self._capacity_probe_interval_cycles == 0
|
||
)
|
||
|
||
def account_capacity_limits(self) -> list[int]:
|
||
with self._state_lock:
|
||
return list(self._account_caps)
|
||
|
||
def observe_capacity_lower_bounds(self, active_counts: list[int | None]) -> list[int]:
|
||
"""Promote known caps from complete account listings without guessing an upper bound."""
|
||
caps_changed = False
|
||
with self._state_lock:
|
||
for index, value in enumerate(active_counts[: len(self._account_caps)]):
|
||
if value is None:
|
||
continue
|
||
observed_count = max(0, int(value))
|
||
if observed_count > self._account_caps[index]:
|
||
self._account_caps[index] = observed_count
|
||
caps_changed = True
|
||
result = list(self._account_caps)
|
||
if caps_changed:
|
||
self._persist_account_caps()
|
||
return result
|
||
|
||
def capacity_probe_enabled(self) -> bool:
|
||
with self._state_lock:
|
||
return self._capacity_probe_enabled
|
||
|
||
def _safe_count_active_tasks(self, client: ModelHubClient, max_count: int) -> int | None:
|
||
try:
|
||
return client.count_active_tasks(max_count=max_count, page_size=200)
|
||
except Exception:
|
||
return None
|
||
|
||
def _safe_list_tasks(self, client: ModelHubClient, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
|
||
try:
|
||
return client.list_tasks(**kwargs)
|
||
except Exception:
|
||
return []
|
||
|
||
def _counts_are_fresh_locked(self, now: float) -> bool:
|
||
return self._counts_initialized and (now - self._active_refresh_at) < self._active_counts_ttl
|
||
|
||
def _refresh_active_counts(self, *, force: bool = False) -> None:
|
||
now = time.monotonic()
|
||
with self._state_lock:
|
||
if not force and self._counts_are_fresh_locked(now):
|
||
return
|
||
|
||
# Do not hold the scheduler lock during remote I/O. One refresher is
|
||
# enough; all submitting threads can continue using their reservations.
|
||
with self._refresh_lock:
|
||
now = time.monotonic()
|
||
with self._state_lock:
|
||
if not force and self._counts_are_fresh_locked(now):
|
||
return
|
||
|
||
with self._state_lock:
|
||
count_limits = [cap + 1 for cap in self._account_caps]
|
||
|
||
def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int | None]:
|
||
return index, self._safe_count_active_tasks(client, count_limits[index])
|
||
|
||
results: list[tuple[int, int | None]] = []
|
||
with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
|
||
futures = {
|
||
executor.submit(_to_indexed_result, index, client): index
|
||
for index, client in enumerate(self.clients)
|
||
}
|
||
for future in as_completed(futures):
|
||
index = futures[future]
|
||
try:
|
||
results.append(future.result())
|
||
except Exception:
|
||
results.append((index, None))
|
||
|
||
refreshed_at = time.monotonic()
|
||
caps_changed = False
|
||
with self._state_lock:
|
||
for index, count in results:
|
||
if count is None:
|
||
self._count_known[index] = False
|
||
continue
|
||
old_remote_count = self._remote_counts[index]
|
||
new_remote_count = max(0, int(count))
|
||
self._count_known[index] = True
|
||
if new_remote_count > self._account_caps[index]:
|
||
self._account_caps[index] = new_remote_count
|
||
caps_changed = True
|
||
acknowledged = max(0, new_remote_count - old_remote_count)
|
||
completed_reservations = sorted(
|
||
(
|
||
(reservation_id, reservation)
|
||
for reservation_id, reservation in self._reservations[index].items()
|
||
if not reservation["inflight"]
|
||
),
|
||
key=lambda item: item[1]["updated_at"],
|
||
)
|
||
for reservation_id, _reservation in completed_reservations[:acknowledged]:
|
||
self._reservations[index].pop(reservation_id, None)
|
||
|
||
# A remote count can stay flat when one old task finishes as
|
||
# one new task appears. Expiry prevents that net-zero update
|
||
# from reserving a slot forever.
|
||
for reservation_id, reservation in list(self._reservations[index].items()):
|
||
if reservation["inflight"]:
|
||
continue
|
||
if refreshed_at - reservation["updated_at"] >= self._reservation_ttl:
|
||
self._reservations[index].pop(reservation_id, None)
|
||
self._remote_counts[index] = new_remote_count
|
||
self._counts_initialized = True
|
||
self._active_refresh_at = refreshed_at
|
||
if caps_changed:
|
||
self._persist_account_caps()
|
||
|
||
def _effective_count_locked(self, index: int) -> int:
|
||
return self._remote_counts[index] + len(self._reservations[index])
|
||
|
||
def _counts_snapshot_locked(self) -> list[int]:
|
||
return [self._effective_count_locked(index) for index in range(len(self.clients))]
|
||
|
||
def _probe_slot_count_locked(self) -> int:
|
||
if not self._capacity_probe_enabled:
|
||
return 0
|
||
return sum(
|
||
1
|
||
for index in range(len(self.clients))
|
||
if index not in self._capacity_probe_attempted
|
||
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
|
||
and self._effective_count_locked(index) >= self._account_caps[index]
|
||
)
|
||
|
||
def active_task_counts(self) -> list[int]:
|
||
self._refresh_active_counts()
|
||
with self._state_lock:
|
||
return self._counts_snapshot_locked()
|
||
|
||
def refresh_active_counts(self) -> list[int]:
|
||
"""Force a remote refresh after out-of-band queue mutations."""
|
||
self._refresh_active_counts(force=True)
|
||
with self._state_lock:
|
||
return self._counts_snapshot_locked()
|
||
|
||
def available_submit_slots(self) -> int:
|
||
self._refresh_active_counts()
|
||
with self._state_lock:
|
||
normal_slots = sum(
|
||
max(0, self._account_caps[index] - self._effective_count_locked(index))
|
||
for index in range(len(self.clients))
|
||
)
|
||
return normal_slots if normal_slots > 0 else self._probe_slot_count_locked()
|
||
|
||
def list_tasks(self, **kwargs): # noqa: ANN003, ANN001
|
||
# Default: use only the reader client to avoid fanout amplification.
|
||
# The fanout_all parameter allows explicit cross-account merging when needed.
|
||
fanout_all = kwargs.pop("_fanout_all", False)
|
||
if not fanout_all:
|
||
return self._safe_list_tasks(self._reader, kwargs)
|
||
|
||
merged: list[dict[str, Any]] = []
|
||
seen: set[str] = set()
|
||
with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
|
||
futures = {executor.submit(self._safe_list_tasks, client, kwargs): index for index, client in enumerate(self.clients)}
|
||
for future in as_completed(futures):
|
||
try:
|
||
tasks = future.result()
|
||
except Exception:
|
||
tasks = []
|
||
for task in tasks:
|
||
task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
|
||
if task_id and task_id in seen:
|
||
continue
|
||
if task_id:
|
||
seen.add(task_id)
|
||
merged.append(task)
|
||
return merged
|
||
|
||
def _read_from_cache(self, model_id: str) -> dict[str, Any] | None:
|
||
cached = self._verify_cache.get(model_id)
|
||
if cached is None:
|
||
return None
|
||
cached_at, payload = cached
|
||
if time.monotonic() - cached_at >= self._verify_cache_ttl:
|
||
self._verify_cache.pop(model_id, None)
|
||
return None
|
||
return dict(payload)
|
||
|
||
def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
|
||
self._verify_cache[model_id] = (time.monotonic(), payload)
|
||
|
||
def begin_cycle(self) -> None:
|
||
"""Keep recent verification results across cycles and prune expired entries."""
|
||
with self._state_lock:
|
||
now = time.monotonic()
|
||
for model_id, (cached_at, _payload) in list(self._verify_cache.items()):
|
||
if now - cached_at >= self._verify_cache_ttl:
|
||
self._verify_cache.pop(model_id, None)
|
||
|
||
def search_by_model_id(self, model_id: str, *, force_refresh: bool = False) -> dict[str, Any]:
|
||
# Reuse recent model verification results across short poll cycles.
|
||
if not force_refresh:
|
||
with self._state_lock:
|
||
cached = self._read_from_cache(model_id)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
# Only query ONE client (the reader) instead of fanning out to all clients.
|
||
# verifyResult is model-specific platform data, not account-specific.
|
||
response = self._reader.search_by_model_id(model_id)
|
||
if not isinstance(response, dict):
|
||
raise ModelHubAPIError("Community model precheck returned a non-object response", payload=response)
|
||
|
||
with self._state_lock:
|
||
self._write_to_cache(model_id, response)
|
||
return response
|
||
|
||
def model_submission_precheck(self, model_id: str, *, force_refresh: bool = False) -> dict[str, Any]:
|
||
return parse_model_submission_precheck(self.search_by_model_id(model_id, force_refresh=force_refresh))
|
||
|
||
def get_verify_result_map(self, model_id: str) -> dict[str, Any]:
|
||
payload = self.search_by_model_id(model_id)
|
||
return ((payload.get("data") or {}).get("verifyResult") or {})
|
||
|
||
def processed_gpus_for_model(self, model_id: str) -> set[str]:
|
||
return set(self.model_submission_precheck(model_id)["processedGpus"])
|
||
|
||
def _reserve_account(
|
||
self,
|
||
excluded: set[int],
|
||
*,
|
||
reserve_capacity_slots: int | None = None,
|
||
) -> tuple[int, int] | None:
|
||
with self._state_lock:
|
||
remaining_by_index = {
|
||
index: self._account_caps[index] - self._effective_count_locked(index)
|
||
for index in range(len(self.clients))
|
||
if index not in excluded
|
||
and (
|
||
reserve_capacity_slots is None
|
||
or (
|
||
self._count_known[index]
|
||
and self._effective_count_locked(index)
|
||
< max(0, self._account_caps[index] - reserve_capacity_slots)
|
||
)
|
||
)
|
||
}
|
||
usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
|
||
if not usable:
|
||
return None
|
||
best_remaining = max(usable.values())
|
||
tied = [index for index, remaining in usable.items() if remaining == best_remaining]
|
||
selected_index = min(
|
||
tied,
|
||
key=lambda index: (index - self._selection_cursor) % len(self.clients),
|
||
)
|
||
self._selection_cursor = (selected_index + 1) % len(self.clients)
|
||
self._reservation_sequence += 1
|
||
reservation_id = self._reservation_sequence
|
||
self._reservations[selected_index][reservation_id] = {
|
||
"inflight": True,
|
||
"updated_at": time.monotonic(),
|
||
}
|
||
return selected_index, reservation_id
|
||
|
||
def _reserve_probe_account(
|
||
self,
|
||
excluded: set[int],
|
||
*,
|
||
reserve_capacity_slots: int | None = None,
|
||
) -> tuple[int, int] | None:
|
||
with self._state_lock:
|
||
if not self._capacity_probe_enabled:
|
||
return None
|
||
eligible = [
|
||
index
|
||
for index in range(len(self.clients))
|
||
if index not in excluded
|
||
and index not in self._capacity_probe_attempted
|
||
and self._capacity_probe_cycle >= self._capacity_probe_cooldown_until[index]
|
||
and self._effective_count_locked(index) >= self._account_caps[index]
|
||
and (
|
||
reserve_capacity_slots is None
|
||
or (
|
||
self._count_known[index]
|
||
and self._effective_count_locked(index)
|
||
< max(0, self._account_caps[index] - reserve_capacity_slots)
|
||
)
|
||
)
|
||
]
|
||
if not eligible:
|
||
return None
|
||
selected_index = min(
|
||
eligible,
|
||
key=lambda index: (index - self._selection_cursor) % len(self.clients),
|
||
)
|
||
self._selection_cursor = (selected_index + 1) % len(self.clients)
|
||
self._capacity_probe_attempted.add(selected_index)
|
||
self._reservation_sequence += 1
|
||
reservation_id = self._reservation_sequence
|
||
self._reservations[selected_index][reservation_id] = {
|
||
"inflight": True,
|
||
"updated_at": time.monotonic(),
|
||
"capacity_probe": True,
|
||
}
|
||
return selected_index, reservation_id
|
||
|
||
def _finish_reservation(self, index: int, reservation_id: int, *, succeeded: bool) -> None:
|
||
with self._state_lock:
|
||
reservation = self._reservations[index].get(reservation_id)
|
||
if reservation is None:
|
||
return
|
||
if not succeeded:
|
||
self._reservations[index].pop(reservation_id, None)
|
||
return
|
||
reservation["inflight"] = False
|
||
reservation["updated_at"] = time.monotonic()
|
||
|
||
def _mark_account_saturated(self, index: int, *, capacity_probe: bool) -> None:
|
||
cap_changed = False
|
||
with self._state_lock:
|
||
if capacity_probe:
|
||
self._capacity_probe_cooldown_until[index] = (
|
||
self._capacity_probe_cycle + self._capacity_probe_cooldown_cycles
|
||
)
|
||
else:
|
||
observed_capacity = max(1, self._effective_count_locked(index))
|
||
if observed_capacity < self._account_caps[index]:
|
||
self._account_caps[index] = observed_capacity
|
||
cap_changed = True
|
||
self._remote_counts[index] = self._account_caps[index]
|
||
self._count_known[index] = True
|
||
self._counts_initialized = True
|
||
self._active_refresh_at = time.monotonic()
|
||
if cap_changed:
|
||
self._persist_account_caps()
|
||
|
||
def _promote_account_capacity(self, index: int) -> None:
|
||
with self._state_lock:
|
||
discovered_cap = max(self._account_caps[index] + 1, self._effective_count_locked(index))
|
||
if discovered_cap <= self._account_caps[index]:
|
||
return
|
||
self._account_caps[index] = discovered_cap
|
||
self._persist_account_caps()
|
||
print(f"[capacity] account={index + 1:02d} discovered_cap={discovered_cap}", flush=True)
|
||
|
||
def old_model_submit_slots(self) -> int:
|
||
"""Return queue positions where models older than the recent window are allowed."""
|
||
self._refresh_active_counts()
|
||
with self._state_lock:
|
||
return sum(
|
||
max(
|
||
0,
|
||
self._old_model_count_limit_locked(index) - self._effective_count_locked(index),
|
||
)
|
||
for index in range(len(self.clients))
|
||
if self._count_known[index]
|
||
)
|
||
|
||
def _old_model_count_limit_locked(self, index: int) -> int:
|
||
return max(0, self._account_caps[index] - self.recent_model_reserve_slots)
|
||
|
||
def old_model_queue_thresholds(self) -> list[int | None]:
|
||
"""Return thresholds, or None when old-model admission must fail closed."""
|
||
with self._state_lock:
|
||
return [
|
||
self._old_model_count_limit_locked(index) if self._count_known[index] else None
|
||
for index in range(len(self.clients))
|
||
]
|
||
|
||
def can_submit_old_models(self) -> bool:
|
||
return self.old_model_submit_slots() > 0
|
||
|
||
def add_task_for_model(
|
||
self,
|
||
payload: dict[str, Any],
|
||
*,
|
||
model_last_modified: datetime | str | None,
|
||
submitted_at: datetime | str | None = None,
|
||
) -> dict[str, Any]:
|
||
reference_time = parse_datetime(submitted_at) or utc_now()
|
||
last_modified = parse_datetime(model_last_modified)
|
||
is_old_model = (
|
||
last_modified is None
|
||
or last_modified < reference_time - timedelta(days=self.recent_model_days)
|
||
)
|
||
return self._add_task(payload, old_model_only=is_old_model)
|
||
|
||
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||
return self._add_task(payload, old_model_only=False)
|
||
|
||
def _add_task(self, payload: dict[str, Any], *, old_model_only: bool) -> dict[str, Any]:
|
||
self._refresh_active_counts()
|
||
attempted_accounts: set[int] = set()
|
||
forced_refresh_done = False
|
||
last_capacity_error: ModelHubAPIError | None = None
|
||
reserved_capacity_slots = self.recent_model_reserve_slots if old_model_only else None
|
||
|
||
while True:
|
||
reservation = self._reserve_account(
|
||
attempted_accounts,
|
||
reserve_capacity_slots=reserved_capacity_slots,
|
||
)
|
||
capacity_probe = False
|
||
if reservation is None:
|
||
reservation = self._reserve_probe_account(
|
||
attempted_accounts,
|
||
reserve_capacity_slots=reserved_capacity_slots,
|
||
)
|
||
capacity_probe = reservation is not None
|
||
if reservation is None:
|
||
if not forced_refresh_done:
|
||
self._refresh_active_counts(force=True)
|
||
forced_refresh_done = True
|
||
continue
|
||
if last_capacity_error is not None:
|
||
raise last_capacity_error
|
||
if old_model_only:
|
||
thresholds = self.old_model_queue_thresholds()
|
||
raise OldModelQueuePolicyError(
|
||
f"所有账号均已达到动态旧模型阈值 {thresholds}（账号上限减 "
|
||
f"{self.recent_model_reserve_slots}），"
|
||
f"仅允许提交最近 {self.recent_model_days} 天内更新的模型"
|
||
)
|
||
raise ModelHubAPIError(
|
||
"当前等待中或运行中的异步模型验证任务数量已达已知上限"
|
||
)
|
||
|
||
selected_index, reservation_id = reservation
|
||
selected_client = self.clients[selected_index]
|
||
try:
|
||
response = selected_client.add_task(payload)
|
||
except ModelHubAPIError as exc:
|
||
self._finish_reservation(selected_index, reservation_id, succeeded=False)
|
||
if not is_capacity_error(exc):
|
||
raise
|
||
# Another process may have filled this account after our count
|
||
# refresh. Mark it full locally and immediately try another one.
|
||
self._mark_account_saturated(selected_index, capacity_probe=capacity_probe)
|
||
attempted_accounts.add(selected_index)
|
||
last_capacity_error = exc
|
||
continue
|
||
except Exception:
|
||
self._finish_reservation(selected_index, reservation_id, succeeded=False)
|
||
raise
|
||
|
||
self._finish_reservation(selected_index, reservation_id, succeeded=True)
|
||
if capacity_probe:
|
||
self._promote_account_capacity(selected_index)
|
||
return response
|
||
|
||
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001
|
||
"""Single-page task listing via the reader client (no fanout)."""
|
||
return self._reader.list_tasks_page(**kwargs)
|
||
|
||
def list_machine_info(self) -> list[dict[str, Any]]:
|
||
return self._reader.list_machine_info()
|
||
|
||
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)
|
||
|
||
def get_build_config(self, task_type: str, target_gpu: str, framework: str) -> str:
|
||
return self._reader.get_build_config(task_type, target_gpu, framework)
|
||
|
||
def find_recent_task_id(self, model_id: str, gpu_type: str, submitted_after: datetime) -> str | None:
|
||
for client in self.clients:
|
||
task_id = client.find_recent_task_id(model_id, gpu_type, submitted_after)
|
||
if task_id is not None:
|
||
return task_id
|
||
return None
|