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, update_jsonl, utc_now
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: dict[tuple[str, str], datetime] = {}

        self._records = read_jsonl(self.path)
        self._rebuild_indexes()

        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 _rebuild_indexes(self) -> None:
        self._by_task_id.clear()
        self._by_model_gpu.clear()
        for record in self._records:
            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()

    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,
        *,
        cooldown_hours: int = 24,
        now: datetime | None = None,
    ) -> bool:
        failed_at = self._failed_model_gpus.get((model_id, target_gpu))
        if failed_at is None:
            return False
        now = now or utc_now()
        return failed_at >= now - timedelta(hours=max(0, int(cooldown_hours)))

    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:
        local_records = list(self._records)

        def merge(existing: list[dict[str, Any]]) -> list[dict[str, Any]]:
            merged = list(existing)
            index_by_key = {_outcome_record_key(record): index for index, record in enumerate(merged)}
            for record in local_records:
                key = _outcome_record_key(record)
                existing_index = index_by_key.get(key)
                if existing_index is None:
                    index_by_key[key] = len(merged)
                    merged.append(record)
                    continue
                merged[existing_index] = _prefer_newer_outcome(merged[existing_index], record)
            return merged

        self._records = update_jsonl(self.path, merge)
        self._rebuild_indexes()

    def _rebuild_failed_index(self) -> None:
        self._failed_model_gpus.clear()
        latest_by_combo: dict[tuple[str, str], tuple[datetime, dict[str, Any]]] = {}
        for record in self._records:
            model_id = record.get("modelId") or ""
            target_gpu = record.get("targetGpu") or ""
            event_time = parse_datetime(record.get("lastSyncTime")) or parse_datetime(record.get("submitTime"))
            if not model_id or not target_gpu or event_time is None:
                continue
            key = (model_id, target_gpu)
            current = latest_by_combo.get(key)
            if current is None or event_time >= current[0]:
                latest_by_combo[key] = (event_time, record)

        for key, (event_time, record) in latest_by_combo.items():
            if record.get("outcome") == "failed":
                self._failed_model_gpus[key] = event_time

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


def _outcome_record_key(record: dict[str, Any]) -> str:
    task_id = record.get("taskId")
    if task_id is not None:
        return f"task:{task_id}"
    return "fallback:{model}|{gpu}|{time}".format(
        model=record.get("modelId") or "",
        gpu=record.get("targetGpu") or "",
        time=record.get("submitTime") or "",
    )


def _outcome_version(record: dict[str, Any]) -> tuple[int, float, float]:
    last_sync = parse_datetime(record.get("lastSyncTime"))
    submit_time = parse_datetime(record.get("submitTime"))
    outcome_rank = 1 if record.get("outcome") in {"success", "failed"} else 0
    return (
        outcome_rank,
        last_sync.timestamp() if last_sync else 0.0,
        submit_time.timestamp() if submit_time else 0.0,
    )


def _prefer_newer_outcome(existing: dict[str, Any], candidate: dict[str, Any]) -> dict[str, Any]:
    if _outcome_version(candidate) >= _outcome_version(existing):
        return candidate
    return existing