social-data-attribution-w2/tests/quality/test_benchmark.py

from __future__ import annotations

from dolma.quality.benchmark import select_recommended_shape, summarize_cpu_results


def test_summarize_cpu_results_projects_full_runtime() -> None:
    summary = summarize_cpu_results(
        4,
        [
            {"docs_classified": 200, "wall_time_seconds": 10.0},
            {"docs_classified": 100, "wall_time_seconds": 5.0},
        ],
        total_pending_groups=30,
    )
    assert summary["cpu"] == 4
    assert summary["docs_classified"] == 300
    assert summary["docs_per_second"] == 20.0
    assert summary["projected_full_runtime_seconds"] == 225.0


def test_select_recommended_shape_chooses_lowest_cpu_within_tolerance() -> None:
    recommended = select_recommended_shape(
        [
            {"cpu": 4, "docs_per_second": 90.0},
            {"cpu": 8, "docs_per_second": 100.0},
            {"cpu": 16, "docs_per_second": 102.0},
        ]
    )
    assert recommended["cpu"] == 8