qwen_test_runner/vision/throughput.py

"""
throughput.py — Pure GPU-hours / cost model (CPU-testable, no torch).

The labeler verdict trades accuracy against throughput: a model that is 94% as
good but 3× faster is the better choice for labeling a million images. This
module converts measured decode speed into samples/hour and GPU-hours/$ per
million labels.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class ThroughputEstimate:
    model: str
    tokens_per_sec: float
    mean_output_tokens: float
    samples_per_hour: float
    gpu_hours_per_million: float
    est_cost_per_million_usd: float


def estimate(model: str, tokens_per_sec: float, mean_output_tokens: float,
             prefill_overhead_s: float = 0.0, gpu_hourly_rate: float = 2.0) -> ThroughputEstimate:
    """samples_per_hour = 3600 / (prefill + output_tokens / tok_per_sec).

    `prefill_overhead_s` is the per-sample vision-encoder + image-token cost
    (measured during the run, not guessed). `gpu_hourly_rate` is a config rate
    printed alongside the result so the dollar figure is transparent.
    """
    if tokens_per_sec <= 0 or mean_output_tokens <= 0:
        return ThroughputEstimate(model, tokens_per_sec, mean_output_tokens, 0.0, float("inf"), float("inf"))
    per_sample_s = prefill_overhead_s + mean_output_tokens / tokens_per_sec
    samples_per_hour = 3600.0 / per_sample_s
    gpu_hours_per_million = 1_000_000.0 / samples_per_hour
    cost = gpu_hours_per_million * gpu_hourly_rate
    return ThroughputEstimate(
        model=model, tokens_per_sec=tokens_per_sec, mean_output_tokens=mean_output_tokens,
        samples_per_hour=samples_per_hour, gpu_hours_per_million=gpu_hours_per_million,
        est_cost_per_million_usd=cost,
    )


def fleet_score(labeler: float, samples_per_hour: float, saturate_at: float = 50_000.0) -> float:
    """Fold throughput into the labeler score for the 'label 1M images' goal.

    Throughput weight saturates (diminishing returns past `saturate_at`), so a
    tiny-but-inaccurate model can't win on speed alone.
    """
    if labeler is None:
        return 0.0
    import math
    w = math.log1p(max(0.0, samples_per_hour)) / math.log1p(saturate_at)
    return labeler * min(1.0, w)