submmit/modelhub_submmit_api/outcome_tracker.py

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),
    }