social-data-attribution-w2/src/dolma/quality/benchmark.py

"""Benchmark helpers for quality sidecar jobs."""

from __future__ import annotations

from typing import Sequence


def summarize_cpu_results(
    cpu: int,
    results: Sequence[dict[str, object]],
    *,
    total_pending_groups: int | None = None,
) -> dict[str, object]:
    docs = 0
    wall_time = 0.0
    for result in results:
        docs += int(result.get("docs_classified", 0))
        wall_time += float(result.get("wall_time_seconds", 0.0))
    docs_per_second = docs / wall_time if wall_time > 0 else 0.0
    summary = {
        "cpu": cpu,
        "group_count": len(results),
        "docs_classified": docs,
        "wall_time_seconds": wall_time,
        "docs_per_second": docs_per_second,
    }
    if total_pending_groups is not None:
        projected_seconds = (
            total_pending_groups * wall_time / len(results) if results else 0.0
        )
        summary["projected_full_runtime_seconds"] = projected_seconds
    return summary


def select_recommended_shape(
    summaries: Sequence[dict[str, object]],
    *,
    tolerance: float = 0.10,
) -> dict[str, object]:
    if not summaries:
        raise ValueError("No summaries provided")
    best_docs_per_second = max(float(item["docs_per_second"]) for item in summaries)
    threshold = best_docs_per_second * (1.0 - tolerance)
    candidates = [
        item for item in summaries if float(item["docs_per_second"]) >= threshold
    ]
    return min(candidates, key=lambda item: int(item["cpu"]))


__all__ = [
    "select_recommended_shape",
    "summarize_cpu_results",
]