feat: add durable success-first modelhub agent
This commit is contained in:
CoolBoy
@@ -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": {
|
||||
|
||||
Reference in New Issue
Block a user