237 lines
9.3 KiB
Python
237 lines
9.3 KiB
Python
from __future__ import annotations
|
||
|
||
from collections import defaultdict
|
||
from datetime import datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from common import append_jsonl, parse_datetime, read_jsonl, utc_now, write_jsonl
|
||
from history_stats import classify_failure, is_failure, is_success
|
||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||
|
||
|
||
DEFAULT_OUTCOMES_PATH = Path("outcomes/submissions.jsonl")
|
||
|
||
|
||
def _now_iso() -> str:
|
||
return utc_now().isoformat()
|
||
|
||
|
||
class OutcomeTracker:
|
||
def __init__(self, path: Path | str) -> None:
|
||
self.path = Path(path)
|
||
self._records: list[dict[str, Any]] = []
|
||
self._by_task_id: dict[str, dict[str, Any]] = {}
|
||
self._by_model_gpu: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
|
||
self._failed_model_gpus: set[tuple[str, str]] = set()
|
||
|
||
loaded = read_jsonl(self.path)
|
||
for record in loaded:
|
||
self._records.append(record)
|
||
task_id = record.get("taskId")
|
||
if task_id:
|
||
self._by_task_id[str(task_id)] = record
|
||
model_id = record.get("modelId") or ""
|
||
target_gpu = record.get("targetGpu") or ""
|
||
self._by_model_gpu[(model_id, target_gpu)].append(record)
|
||
|
||
self._rebuild_failed_index()
|
||
|
||
last_sync_times = [
|
||
parse_datetime(record.get("lastSyncTime"))
|
||
for record in self._records
|
||
if record.get("lastSyncTime")
|
||
]
|
||
self._last_sync_time: datetime = max(last_sync_times) if last_sync_times else utc_now() - timedelta(days=7)
|
||
|
||
def record_submission(
|
||
self,
|
||
model_id: str,
|
||
target_gpu: str,
|
||
framework: str,
|
||
task_type: str,
|
||
task_id: str | None,
|
||
submit_time: str,
|
||
) -> None:
|
||
record: dict[str, Any] = {
|
||
"modelId": model_id,
|
||
"targetGpu": target_gpu,
|
||
"framework": framework,
|
||
"taskType": task_type,
|
||
"taskId": task_id,
|
||
"submitTime": submit_time,
|
||
"lastSyncTime": None,
|
||
"status": "pending",
|
||
"verifyResult": None,
|
||
"outcome": "pending",
|
||
"failReason": None,
|
||
}
|
||
self._records.append(record)
|
||
if task_id:
|
||
self._by_task_id[task_id] = record
|
||
self._by_model_gpu[(model_id, target_gpu)].append(record)
|
||
append_jsonl(self.path, record)
|
||
|
||
def sync_from_api(self, client: ModelHubClient | ModelHubClientPool) -> int:
|
||
try:
|
||
begin = self._last_sync_time
|
||
end = utc_now()
|
||
# Use fanout for pools to sync outcomes across all accounts
|
||
list_kwargs: dict[str, Any] = {"begin_time": begin, "end_time": end, "page_size": 100, "only_mine": True}
|
||
if isinstance(client, ModelHubClientPool):
|
||
list_kwargs["_fanout_all"] = True
|
||
tasks = client.list_tasks(**list_kwargs)
|
||
except Exception:
|
||
return 0
|
||
|
||
updated_count = 0
|
||
for task in tasks:
|
||
task_id = str(task.get("taskId")) if task.get("taskId") is not None else None
|
||
if not task_id:
|
||
continue
|
||
|
||
existing = self._by_task_id.get(task_id)
|
||
if existing is not None:
|
||
if existing.get("outcome") == "pending":
|
||
self._update_record_from_task(existing, task)
|
||
updated_count += 1
|
||
else:
|
||
status = str(task.get("status") or "").lower()
|
||
if status in {"success", "failed", "error", "cancelled", "completed"}:
|
||
record = self._create_record_from_task(task)
|
||
self._records.append(record)
|
||
self._by_task_id[task_id] = record
|
||
model_id = record["modelId"]
|
||
target_gpu = record["targetGpu"]
|
||
self._by_model_gpu[(model_id, target_gpu)].append(record)
|
||
updated_count += 1
|
||
|
||
if updated_count:
|
||
self._last_sync_time = end
|
||
self._rebuild_failed_index()
|
||
self.save()
|
||
|
||
return updated_count
|
||
|
||
def is_model_gpu_failed(self, model_id: str, target_gpu: str) -> bool:
|
||
return (model_id, target_gpu) in self._failed_model_gpus
|
||
|
||
def get_stats_report(self) -> dict[str, Any]:
|
||
now = _now_iso()
|
||
terminal = [r for r in self._records if r.get("outcome") in {"success", "failed"}]
|
||
|
||
gpu_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
framework_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
combo_groups: dict[tuple[str, str, str], list[dict[str, Any]]] = defaultdict(list)
|
||
|
||
for record in terminal:
|
||
gpu = record.get("targetGpu") or "unknown"
|
||
fw = record.get("framework") or "unknown"
|
||
tt = record.get("taskType") or "unknown"
|
||
gpu_groups[gpu].append(record)
|
||
framework_groups[f"{fw}"].append(record)
|
||
combo_groups[(gpu, fw, tt)].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()}
|
||
combination_stats = {
|
||
f"{gpu}|{fw}|{tt}": {"targetGpu": gpu, "framework": fw, "taskType": tt, **_summarize(records)}
|
||
for (gpu, fw, tt), records in combo_groups.items()
|
||
}
|
||
|
||
warnings: list[str] = []
|
||
for gpu, summary in gpu_summaries.items():
|
||
if summary["total"] >= 4 and summary["failureRate"] >= 0.5:
|
||
warnings.append(f"GPU {gpu} 本地统计失败率偏高（≥50%），建议重点关注。")
|
||
for key, stat in combination_stats.items():
|
||
if stat["total"] >= 3 and stat["failureRate"] >= 0.6:
|
||
warnings.append(f"组合 {key} 近期失败集中，建议降低该 GPU+框架的提交优先级。")
|
||
|
||
pending_count = sum(1 for r in self._records if r.get("outcome") == "pending")
|
||
|
||
return {
|
||
"generatedAt": now,
|
||
"totalRecords": len(self._records),
|
||
"pendingRecords": pending_count,
|
||
"terminalRecords": len(terminal),
|
||
"gpuSummaries": gpu_summaries,
|
||
"frameworkSummaries": framework_summaries,
|
||
"combinationStats": combination_stats,
|
||
"totals": _summarize(terminal),
|
||
"warnings": warnings,
|
||
}
|
||
|
||
def save(self) -> None:
|
||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||
write_jsonl(self.path, self._records)
|
||
|
||
def _rebuild_failed_index(self) -> None:
|
||
self._failed_model_gpus.clear()
|
||
for record in self._records:
|
||
if record.get("outcome") == "failed":
|
||
model_id = record.get("modelId") or ""
|
||
target_gpu = record.get("targetGpu") or ""
|
||
if model_id and target_gpu:
|
||
self._failed_model_gpus.add((model_id, target_gpu))
|
||
|
||
@staticmethod
|
||
def _update_record_from_task(record: dict[str, Any], task: dict[str, Any]) -> None:
|
||
record["status"] = task.get("status")
|
||
record["verifyResult"] = task.get("verifyResult")
|
||
record["lastSyncTime"] = _now_iso()
|
||
if is_success(task):
|
||
record["outcome"] = "success"
|
||
record["failReason"] = None
|
||
elif is_failure(task):
|
||
record["outcome"] = "failed"
|
||
record["failReason"] = classify_failure(task)
|
||
else:
|
||
record["outcome"] = "pending"
|
||
|
||
@staticmethod
|
||
def _create_record_from_task(task: dict[str, Any]) -> dict[str, Any]:
|
||
create_time = parse_datetime(task.get("createTime"))
|
||
record: dict[str, Any] = {
|
||
"modelId": task.get("modelId") or task.get("model_id") or "",
|
||
"targetGpu": task.get("gpuType") or task.get("targetGpu") or "",
|
||
"framework": task.get("framework") or "",
|
||
"taskType": task.get("taskType") or "",
|
||
"taskId": str(task.get("taskId")) if task.get("taskId") is not None else None,
|
||
"submitTime": create_time.isoformat() if create_time else _now_iso(),
|
||
"lastSyncTime": _now_iso(),
|
||
"status": task.get("status"),
|
||
"verifyResult": task.get("verifyResult"),
|
||
"outcome": "pending",
|
||
"failReason": None,
|
||
}
|
||
if is_success(task):
|
||
record["outcome"] = "success"
|
||
elif is_failure(task):
|
||
record["outcome"] = "failed"
|
||
record["failReason"] = classify_failure(task)
|
||
return record
|
||
|
||
|
||
def _summarize(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||
total = len(records)
|
||
success_count = sum(1 for r in records if r.get("outcome") == "success")
|
||
failure_count = sum(1 for r in records if r.get("outcome") == "failed")
|
||
pending_count = sum(1 for r in records if r.get("outcome") == "pending")
|
||
|
||
failure_breakdown: dict[str, int] = defaultdict(int)
|
||
for r in records:
|
||
reason = r.get("failReason")
|
||
if reason:
|
||
failure_breakdown[reason] += 1
|
||
|
||
return {
|
||
"total": total,
|
||
"successCount": success_count,
|
||
"failureCount": failure_count,
|
||
"pendingCount": pending_count,
|
||
"successRate": round(success_count / total, 4) if total > 0 else 0.0,
|
||
"failureRate": round(failure_count / total, 4) if total > 0 else 0.0,
|
||
"pendingRate": round(pending_count / total, 4) if total > 0 else 0.0,
|
||
"failureBreakdown": dict(failure_breakdown),
|
||
}
|