src/inferscale/diagnostics.py

from __future__ import annotations

from typing import Any

from .models import SimulationConfig


def diagnose_run(summary: dict[str, Any], latency: dict[str, Any], resource: dict[str, Any], cfg: SimulationConfig) -> dict[str, Any]:
    """Explain the simulator's limiting regime from its own state.

    Diagnostics are heuristic labels over simulated telemetry; they are not
    hardware-counter measurements and must be interpreted with the profile
    provenance attached to every run.
    """

    target = cfg.slo_attainment_target
    attainment = float(summary.get("slo_attainment", 0.0))
    ttft_attainment = float(summary.get("ttft_slo_attainment", 0.0))
    e2e_attainment = float(summary.get("e2e_slo_attainment", 0.0))
    unfinished = int(summary.get("requests_unfinished", 0))
    busy = float(summary.get("busy_fraction", 0.0))
    kv_util = float(resource.get("peak_kv_utilization", 0.0))
    queue_p95 = float(latency.get("queue_ms", {}).get("p95", 0.0))
    prefill_busy = float(resource.get("prefill_busy_fraction", 0.0))
    decode_busy = float(resource.get("decode_busy_fraction", 0.0))
    transfer_busy = float(resource.get("transfer_busy_fraction", 0.0))
    topology = str(resource.get("topology", cfg.topology))

    code = "mixed_pressure"
    label = "Mixed SLO pressure"
    explanation = "Several constraints are active; inspect queueing, KV pressure, and tail latency together."
    recommendation = "Compare schedulers and run the capacity planner to isolate the limiting regime."

    if topology == "disaggregated_pd" and transfer_busy >= 0.85 and attainment < target:
        code, label = "kv_transfer_pressure", "KV transfer pressure"
        explanation = "The simulated P/D interconnect is highly utilized and transfer time is contributing to first-token latency."
        recommendation = "Increase link bandwidth, reduce prompt pressure, or compare against colocated serving."
    elif topology == "disaggregated_pd" and prefill_busy >= 0.92 and prefill_busy > decode_busy + 0.08:
        code, label = "prefill_pool_pressure", "Prefill pool pressure"
        explanation = "Prefill workers are the hottest role in the disaggregated pipeline."
        recommendation = "Add prefill capacity, use prefix reuse, or reduce prompt-side pressure."
    elif topology == "disaggregated_pd" and decode_busy >= 0.92 and decode_busy > prefill_busy + 0.08:
        code, label = "decode_pool_pressure", "Decode pool pressure"
        explanation = "Decode workers are the hottest role in the disaggregated pipeline."
        recommendation = "Add decode capacity or reduce output-length / offered-load pressure."
    elif unfinished > 0 and kv_util >= 0.85:
        code, label = "kv_pressure", "KV memory pressure"
        explanation = "Requests remain unfinished while the simulated KV cache approaches its configured capacity."
        recommendation = "Reduce sequence pressure, lower offered load, use a smaller KV block, or choose a larger-memory profile."
    elif unfinished > 0 and busy >= 0.92:
        code, label = "over_capacity", "Offered load above capacity"
        explanation = "The virtual serving system remains busy while the workload fails to drain."
        recommendation = "Reduce offered load or use the capacity planner to find a sustainable operating point."
    elif unfinished > 0:
        code, label = "admission_stall", "Admission stall"
        explanation = "Some requests cannot be admitted or completed under the current serving constraints."
        recommendation = "Inspect batch-token and KV-cache limits, then compare a less restrictive configuration."
    elif attainment >= target:
        if busy >= 0.90:
            code, label = "healthy_near_saturation", "Healthy, little headroom"
            explanation = "The workload meets the configured SLO target, but simulated resource utilization is already high."
            recommendation = "Run the capacity planner before increasing traffic; the current point is close to saturation."
        else:
            code, label = "healthy", "Healthy"
            explanation = "The workload drains and meets the configured SLO-attainment target with visible headroom."
            recommendation = "Use the capacity planner or design explorer to quantify additional headroom."
    elif kv_util >= 0.90:
        code, label = "kv_pressure", "KV memory pressure"
        explanation = "Peak simulated KV utilization is high enough to constrain admission and batching flexibility."
        recommendation = "Reduce sequence pressure, use smaller KV blocks, enable reusable prefixes, or select a larger-memory profile."
    elif ttft_attainment < target and queue_p95 >= max(25.0, cfg.slo_ttft_ms * 0.20):
        code, label = "queue_prefill_pressure", "Queue / prefill pressure"
        explanation = "TTFT misses coincide with substantial queueing before prefill begins."
        recommendation = "Lower offered load, increase safe batching capacity, or try chunked prefill / SLO-aware scheduling."
    elif ttft_attainment < target:
        code, label = "ttft_pressure", "TTFT pressure"
        explanation = "First-token latency is the dominant SLO failure even though queueing is not the only contributor."
        recommendation = "Compare prefill-oriented scheduling, prefix reuse, and P/D disaggregation scenarios."
    elif e2e_attainment < target:
        code, label = "decode_pressure", "Decode / E2E pressure"
        explanation = "Most first tokens arrive within target, but end-to-end latency still misses the configured SLO."
        recommendation = "Reduce output-length pressure or offered load; larger decode capacity may help when memory allows."
    elif busy >= 0.95:
        code, label = "compute_saturation", "Compute saturation"
        explanation = "The simulated serving resources are effectively saturated even though no single SLO component dominates."
        recommendation = "Reduce load or switch to a higher-throughput accelerator profile."

    return {
        "code": code,
        "label": label,
        "explanation": explanation,
        "recommendation": recommendation,
        "evidence": {
            "slo_attainment": attainment,
            "ttft_slo_attainment": ttft_attainment,
            "e2e_slo_attainment": e2e_attainment,
            "busy_fraction": busy,
            "peak_kv_utilization": kv_util,
            "queue_p95_ms": queue_p95,
            "unfinished_requests": unfinished,
            "prefill_busy_fraction": prefill_busy,
            "decode_busy_fraction": decode_busy,
            "transfer_busy_fraction": transfer_busy,
        },
        "provenance": "heuristic-simulator-diagnosis",
    }