src/inferscale/measurements.py

from __future__ import annotations

import json
import math
import random
from copy import deepcopy
from statistics import mean, median
from typing import Any

from .models import SimulationConfig
from .simulator import run_simulation
from .validation import validate_cases


def _first_number(row: dict[str, Any], keys: tuple[str, ...]) -> float | None:
    for key in keys:
        value = row.get(key)
        if isinstance(value, (int, float)) and math.isfinite(float(value)):
            return float(value)
    return None


def _percentile_from_pairs(value: Any, target: float) -> float | None:
    if not isinstance(value, list):
        return None
    pairs: list[tuple[float, float]] = []
    for item in value:
        if isinstance(item, (list, tuple)) and len(item) >= 2:
            try:
                pairs.append((float(item[0]), float(item[1])))
            except (TypeError, ValueError):
                continue
    if not pairs:
        return None
    exact = [metric for percentile, metric in pairs if abs(percentile - target) < 1e-9]
    if exact:
        return exact[0]
    nearest = min(pairs, key=lambda pair: abs(pair[0] - target))
    return nearest[1] if abs(nearest[0] - target) <= 1.0 else None


def _metric(row: dict[str, Any], kind: str, percentile: int) -> float | None:
    aliases = {
        "ttft": (f"p{percentile}_ttft_ms", f"ttft_p{percentile}_ms"),
        "e2e": (f"p{percentile}_e2e_latency_ms", f"p{percentile}_e2el_ms", f"p{percentile}_e2e_ms"),
    }
    direct = _first_number(row, aliases[kind])
    if direct is not None:
        return direct
    if kind == "ttft":
        return _percentile_from_pairs(row.get("percentiles_ttft_ms"), percentile)
    return _percentile_from_pairs(row.get("percentiles_e2el_ms"), percentile)


def _parse_content(content: str) -> list[dict[str, Any]]:
    text = str(content).strip()
    if not text:
        raise ValueError("Measurement file is empty")
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return [dict(item) for item in parsed if isinstance(item, dict)]
        if isinstance(parsed, dict):
            if isinstance(parsed.get("cases"), list):
                return [dict(item) for item in parsed["cases"] if isinstance(item, dict)]
            return [parsed]
    except json.JSONDecodeError:
        pass
    rows: list[dict[str, Any]] = []
    for line_no, line in enumerate(text.splitlines(), 1):
        line = line.strip()
        if not line:
            continue
        try:
            parsed = json.loads(line)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid JSON/JSONL on line {line_no}: {exc.msg}") from exc
        if not isinstance(parsed, dict):
            raise ValueError(f"JSONL line {line_no} must contain an object")
        rows.append(parsed)
    if not rows:
        raise ValueError("No measurement records found")
    return rows


def _detect_source(row: dict[str, Any]) -> str:
    backend = str(row.get("backend", "")).lower()
    if "sglang" in backend or "input_throughput" in row or "total_output_tokens_retokenized" in row:
        return "sglang"
    if "request_goodput" in row or "percentiles_ttft_ms" in row or "label" in row:
        return "vllm"
    if "config" in row and "measured" in row:
        return "case_bundle"
    return "generic"


def _normalized_case(
    row: dict[str, Any],
    source: str,
    base_config: dict[str, Any],
    index: int,
) -> dict[str, Any]:
    if source == "case_bundle" and "config" in row and "measured" in row:
        return deepcopy(row)

    cfg = SimulationConfig.from_dict(base_config).to_dict()
    input_len = _first_number(row, ("input_len", "random_input_len", "random_input", "sharegpt_input_len"))
    output_len = _first_number(row, ("output_len", "random_output_len", "random_output", "sharegpt_output_len"))
    request_rate = _first_number(row, ("request_rate",))
    if input_len is not None and input_len > 0:
        cfg["prompt_tokens_mean"] = int(round(input_len))
        cfg["prompt_tokens_cv"] = 0.0
    if output_len is not None and output_len > 0:
        cfg["output_tokens_mean"] = int(round(output_len))
        cfg["output_tokens_cv"] = 0.0
    if request_rate is not None and math.isfinite(request_rate) and request_rate > 0:
        cfg["request_rate_rps"] = request_rate

    measured: dict[str, float] = {}
    for percentile in (95, 99):
        ttft = _metric(row, "ttft", percentile)
        e2e = _metric(row, "e2e", percentile)
        if ttft is not None:
            measured[f"p{percentile}_ttft_ms"] = ttft
        if e2e is not None:
            measured[f"p{percentile}_e2e_ms"] = e2e
    goodput = _first_number(row, ("request_goodput", "goodput_rps"))
    throughput = _first_number(row, ("request_throughput", "request_throughput_rps"))
    if goodput is not None:
        measured["goodput_rps"] = goodput
    elif throughput is not None:
        measured["request_throughput_rps"] = throughput

    if not any(key.startswith("p95_") or key.startswith("p99_") for key in measured):
        raise ValueError(
            f"Measurement record {index + 1} does not contain a supported TTFT/E2E percentile. "
            "Prefer p95/p99 output from the serving benchmark."
        )

    return {
        "name": str(row.get("label") or row.get("name") or f"{source}-case-{index + 1}"),
        "config": cfg,
        "measured": measured,
        "measurement_metadata": {
            "source": source,
            "backend": row.get("backend"),
            "model": row.get("model") or row.get("model_id"),
            "dataset_name": row.get("dataset_name"),
            "max_concurrency": row.get("max_concurrency"),
            "completed": row.get("completed"),
            "raw_request_rate": row.get("request_rate"),
        },
    }


def import_measurements(content: str, source: str = "auto", base_config: dict[str, Any] | None = None) -> dict[str, Any]:
    rows = _parse_content(content)
    base_config = base_config or SimulationConfig().to_dict()
    cases: list[dict[str, Any]] = []
    detected: list[str] = []
    for index, row in enumerate(rows):
        row_source = _detect_source(row) if source == "auto" else source
        detected.append(row_source)
        cases.append(_normalized_case(row, row_source, base_config, index))
    return {
        "case_count": len(cases),
        "source": source,
        "detected_sources": sorted(set(detected)),
        "cases": cases,
        "provenance": "external-serving-benchmark-import",
        "note": (
            "Imported benchmark artifacts are normalized into InferScale validation cases. Model/accelerator/precision "
            "fall back to the current Serving Lab configuration unless the case bundle supplies an explicit config."
        ),
    }


def _ratio_median(values: list[float]) -> float:
    cleaned = [value for value in values if math.isfinite(value) and value > 0]
    return min(max(median(cleaned), 0.20), 5.0) if cleaned else 1.0


def _predict_metric(result: dict[str, Any], metric: str) -> float:
    if metric.startswith("p95_ttft"):
        return float(result["latency"]["ttft_ms"]["p95"])
    if metric.startswith("p99_ttft"):
        return float(result["latency"]["ttft_ms"]["p99"])
    if metric.startswith("p95_e2e"):
        return float(result["latency"]["e2e_ms"]["p95"])
    if metric.startswith("p99_e2e"):
        return float(result["latency"]["e2e_ms"]["p99"])
    raise KeyError(metric)


def _fit_scales(cases: list[dict[str, Any]]) -> dict[str, float]:
    prefill_ratios: list[float] = []
    decode_ratios: list[float] = []
    for case in cases:
        result = run_simulation(case["config"])
        measured = case["measured"]
        for percentile in (95, 99):
            ttft_key = f"p{percentile}_ttft_ms"
            e2e_key = f"p{percentile}_e2e_ms"
            if ttft_key in measured:
                predicted_ttft = _predict_metric(result, ttft_key)
                if predicted_ttft > 1e-9:
                    prefill_ratios.append(float(measured[ttft_key]) / predicted_ttft)
            if ttft_key in measured and e2e_key in measured:
                predicted_ttft = _predict_metric(result, ttft_key)
                predicted_e2e = _predict_metric(result, e2e_key)
                measured_decode = max(float(measured[e2e_key]) - float(measured[ttft_key]), 1e-9)
                predicted_decode = max(predicted_e2e - predicted_ttft, 1e-9)
                decode_ratios.append(measured_decode / predicted_decode)
    return {
        "prefill_time_scale": _ratio_median(prefill_ratios),
        "decode_time_scale": _ratio_median(decode_ratios),
        "transfer_time_scale": 1.0,
    }


def _apply_scales(cases: list[dict[str, Any]], scales: dict[str, float]) -> list[dict[str, Any]]:
    output = deepcopy(cases)
    for case in output:
        case["config"].update(scales)
    return output


def calibrate_measurements(
    cases: list[dict[str, Any]],
    holdout_fraction: float = 0.33,
    seed: int = 7,
) -> dict[str, Any]:
    if not cases:
        raise ValueError("Calibration requires at least one measurement case")
    holdout_fraction = min(max(float(holdout_fraction), 0.0), 0.80)
    indices = list(range(len(cases)))
    random.Random(seed).shuffle(indices)
    if len(cases) >= 3 and holdout_fraction > 0:
        holdout_count = max(1, min(len(cases) - 1, int(round(len(cases) * holdout_fraction))))
        holdout_idx = set(indices[:holdout_count])
        train = [case for idx, case in enumerate(cases) if idx not in holdout_idx]
        holdout = [case for idx, case in enumerate(cases) if idx in holdout_idx]
        validation_mode = "held-out"
    else:
        train = list(cases)
        holdout = list(cases)
        validation_mode = "resubstitution-insufficient-cases-for-holdout"

    scales = _fit_scales(train)
    baseline = validate_cases(holdout)
    calibrated = validate_cases(_apply_scales(holdout, scales))
    return {
        "case_count": len(cases),
        "train_count": len(train),
        "holdout_count": len(holdout),
        "validation_mode": validation_mode,
        "holdout_fraction": holdout_fraction,
        "fitted_scales": scales,
        "baseline": baseline,
        "calibrated": calibrated,
        "improvement": {
            "mape_points": baseline["mape_pct"] - calibrated["mape_pct"],
            "relative_mape_reduction": (
                (baseline["mape_pct"] - calibrated["mape_pct"]) / baseline["mape_pct"]
                if baseline["mape_pct"] > 1e-12
                else 0.0
            ),
        },
        "provenance": "robust-median-scale-calibration-with-heldout-validation",
        "note": (
            "Calibration fits robust global prefill/decode multipliers from training cases only. It is intentionally "
            "simple: calibration can correct global timing bias but cannot validate scheduler semantics or unseen hardware."
        ),
    }