Upload agent/tools/compare_runs.py with huggingface_hub

agent/tools/compare_runs.py

@@ -0,0 +1,234 @@
+"""compare_runs tool — build the side-by-side Report from two RunMetrics + Patch.
+
+Mostly a pure transform. Resilient to one specific LLM failure mode:
+sometimes the model collapses the Patch dict down to its ``new_config`` fields
+when forwarding it between turns (passing
+``patch={"precision": "bf16", "attention_impl": "flash_rocm", ...}`` instead of
+``patch={"new_config": {...}, "diff": ..., "rationale": [...], ...}``).
+``_normalize_patch`` detects that shape and either:
+
+  1. Substitutes the cached ``propose_patch`` result if available
+     (preferred — preserves rationale, diff, predicted speedup, confidence).
+  2. Wraps the flat config as a minimal Patch (fallback — Report has empty
+     rationale and no predicted speedup, but still ships).
+
+Either way the audit produces a Report instead of bailing on a Pydantic
+ValidationError.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent.schemas import (
+    MetricDelta,
+    Patch,
+    Report,
+    RunMetrics,
+    ToolResult,
+    WorkloadConfig,
+)
+from agent.tools import Tool
+
+
+# Patch-shape sentinel keys: a real Patch dict has at minimum new_config + diff.
+_PATCH_KEYS = {"new_config", "diff", "rationale", "expected_speedup_low"}
+
+# WorkloadConfig sentinel keys: presence of any of these (without _PATCH_KEYS)
+# means the LLM passed a flat WorkloadConfig instead of a Patch envelope.
+# Broader than just the "always-set" fields — Qwen has been observed to send
+# only the *changed* fields after propose_patch (e.g. just the dataloader
+# group), so we accept any WorkloadConfig field name as a signal.
+_FLAT_CONFIG_KEYS = frozenset(
+    {
+        "model_name",
+        "precision",
+        "attention_impl",
+        "batch_size",
+        "grad_accum_steps",
+        "seq_len",
+        "optimizer",
+        "gradient_checkpointing",
+        "lora_rank",
+        "dataloader_workers",
+        "dataloader_pin_memory",
+        "dataloader_prefetch_factor",
+        "dataloader_persistent_workers",
+        "torch_compile",
+        "lr",
+        "warmup_steps",
+        "env_vars",
+    }
+)
+
+
+def _looks_like_flat_config(d: dict[str, Any]) -> bool:
+    """Return True iff `d` looks like a flat WorkloadConfig (or partial diff)
+    rather than a Patch envelope.
+
+    A real Patch always carries at least one of `_PATCH_KEYS` (`new_config`,
+    `diff`, etc.). If none of those are present and at least one
+    WorkloadConfig field is, the LLM almost certainly forwarded a flat config
+    or a partial diff. ``_normalize_patch`` then recovers via the cached
+    propose_patch result.
+    """
+    if not isinstance(d, dict):
+        return False
+    if any(k in d for k in _PATCH_KEYS):
+        return False
+    return any(k in d for k in _FLAT_CONFIG_KEYS)
+
+
+def _normalize_patch(patch: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
+    """Return ``(patch_dict, notes)`` — never raises on a malformed input.
+
+    If the LLM passed a flat WorkloadConfig dict instead of a Patch envelope,
+    we recover by checking ``propose_patch.latest_patch()`` first (full
+    fidelity) and falling back to wrapping the flat config (low fidelity).
+    """
+    notes: list[str] = []
+    if isinstance(patch, dict) and any(k in patch for k in _PATCH_KEYS):
+        return patch, notes
+
+    if _looks_like_flat_config(patch):
+        # Try the cached patch first — preserves rationale + diff + uplift.
+        from agent.tools.propose_patch import latest_patch as _latest_patch
+
+        cached = _latest_patch()
+        if cached is not None:
+            notes.append(
+                "patch arg looked like a flat WorkloadConfig; substituted the "
+                "cached propose_patch result so rationale and diff survive."
+            )
+            return cached, notes
+
+        # No cache — wrap the flat config minimally so Report still renders.
+        try:
+            wrapped_cfg = WorkloadConfig.model_validate(
+                {"model_name": "unknown", **patch}
+            ).model_dump()
+        except Exception:
+            wrapped_cfg = patch
+        notes.append(
+            "patch arg looked like a flat WorkloadConfig and no cached "
+            "propose_patch result was available; synthesized a minimal "
+            "Patch (rationale empty, no predicted speedup)."
+        )
+        return (
+            {
+                "new_config": wrapped_cfg,
+                "diff": "(diff unavailable — patch was passed as flat config)",
+                "rationale": [],
+                "expected_speedup_low": 1.0,
+                "expected_speedup_high": 1.0,
+                "confidence": 0.0,
+            },
+            notes,
+        )
+
+    # Some other malformed shape — let pydantic produce a clear error.
+    return patch, notes
+
+
+def _compare_runs(
+    workload_name: str, before: dict, after: dict, patch: dict
+) -> ToolResult:
+    patch_dict, patch_notes = _normalize_patch(patch)
+
+    try:
+        before_m = RunMetrics.model_validate(before)
+        after_m = RunMetrics.model_validate(after)
+        patch_m = Patch.model_validate(patch_dict)
+    except Exception as exc:
+        return ToolResult(ok=False, error=f"{type(exc).__name__}: {exc}")
+
+    speedup = (
+        after_m.tokens_per_sec / before_m.tokens_per_sec
+        if before_m.tokens_per_sec
+        else 0.0
+    )
+
+    deltas = [
+        MetricDelta(
+            name="tokens_per_sec",
+            before=before_m.tokens_per_sec,
+            after=after_m.tokens_per_sec,
+            unit="tok/s",
+        ),
+        MetricDelta(
+            name="mfu_pct", before=before_m.mfu_pct, after=after_m.mfu_pct, unit="%"
+        ),
+        MetricDelta(
+            name="hbm_peak_gb",
+            before=before_m.hbm_peak_gb,
+            after=after_m.hbm_peak_gb,
+            unit="GB",
+        ),
+        MetricDelta(
+            name="gpu_util_pct",
+            before=before_m.gpu_util_pct,
+            after=after_m.gpu_util_pct,
+            unit="%",
+        ),
+    ]
+
+    summary = (
+        f"Tokens/sec: {before_m.tokens_per_sec:.0f} → {after_m.tokens_per_sec:.0f} "
+        f"({speedup:.2f}×). MFU: {before_m.mfu_pct:.0f}% → {after_m.mfu_pct:.0f}%."
+    )
+
+    report = Report(
+        workload_name=workload_name,
+        before=before_m,
+        after=after_m,
+        patch=patch_m,
+        metric_deltas=deltas,
+        waste_budget_before=before_m.waste_budget,
+        waste_budget_after=after_m.waste_budget,
+        speedup_actual=round(speedup, 2),
+        speedup_predicted_low=patch_m.expected_speedup_low,
+        speedup_predicted_high=patch_m.expected_speedup_high,
+        confidence=patch_m.confidence,
+        summary_line=summary,
+    )
+    payload = report.model_dump()
+    if patch_notes:
+        payload["notes"] = patch_notes
+    return ToolResult(ok=True, result=payload)
+
+
+COMPARE_RUNS = Tool(
+    name="compare_runs",
+    description=(
+        "Build the final side-by-side Report from a baseline RunMetrics, an "
+        "optimized RunMetrics, and the Patch that connects them. Pure function — "
+        "always call this last.\n"
+        "\n"
+        "`patch` should be the FULL Patch dict you got back from propose_patch "
+        "(with `new_config`, `diff`, `rationale`, `expected_speedup_low`, etc.) "
+        "— NOT just the optimized config fields. If you forward a flat config "
+        "by mistake, compare_runs will recover by looking up the cached "
+        "propose_patch result, but you'll lose detail."
+    ),
+    input_schema={
+        "type": "object",
+        "properties": {
+            "workload_name": {
+                "type": "string",
+                "description": "Human-readable workload label for the report header.",
+            },
+            "before": {"type": "object", "description": "Baseline RunMetrics dict."},
+            "after": {"type": "object", "description": "Optimized RunMetrics dict."},
+            "patch": {
+                "type": "object",
+                "description": (
+                    "FULL Patch dict from propose_patch — must include "
+                    "`new_config`, `diff`, `rationale`, `expected_speedup_low/high`, "
+                    "`confidence`. Do not pass just the optimized config fields."
+                ),
+            },
+        },
+        "required": ["workload_name", "before", "after", "patch"],
+    },
+    fn=_compare_runs,
+)