feat: wandb logging for microbenchmarks and lm-eval accuracy

veeksha/evaluator/accuracy/base.py

@@ -405,6 +405,78 @@ def save(self, output_dir: str) -> None:
         path = os.path.join(output_dir, "lmeval_results.json")
         with open(path, "w", encoding="utf-8") as f:
             json.dump(self._results_dict, f, indent=2, default=str)
+        self._log_wandb_metrics()
+
+    def _log_wandb_metrics(self) -> None:
+        """Log lm-eval accuracy metrics to Weights & Biases (if a run is active)."""
+        try:
+            from typing import Any as _Any
+
+            import wandb  # type: ignore[import-not-found]
+
+            wandb = cast(_Any, wandb)
+            if not getattr(wandb, "run", None):
+                return
+
+            results = self._results_dict
+            if not results:
+                return
+
+            # Per-task scalar metrics
+            flat_metrics: Dict[str, float] = {}
+            for task_name, task_results in results.get("results", {}).items():
+                if not isinstance(task_results, dict):
+                    continue
+                for metric_name, value in task_results.items():
+                    if isinstance(value, (int, float)) and not isinstance(value, bool):
+                        clean_metric = metric_name.split(",")[0]
+                        flat_metrics[f"lmeval/{task_name}/{clean_metric}"] = float(
+                            value
+                        )
+
+            if flat_metrics:
+                wandb.log(flat_metrics)
+
+            # Results table
+            table_rows = []
+            for task_name, task_results in results.get("results", {}).items():
+                if not isinstance(task_results, dict):
+                    continue
+                for metric_name, value in task_results.items():
+                    if metric_name.startswith("alias"):
+                        continue
+                    table_rows.append(
+                        [
+                            task_name,
+                            metric_name.split(",")[0],
+                            metric_name.split(",")[1] if "," in metric_name else "none",
+                            value,
+                        ]
+                    )
+
+            if table_rows:
+                wandb.log(
+                    {
+                        "lmeval_results_table": wandb.Table(
+                            columns=["task", "metric", "filter", "value"],
+                            data=table_rows,
+                        )
+                    }
+                )
+
+            # Summary counts
+            wandb.log(
+                {
+                    "lmeval/num_completed_requests": self.num_completed_requests,
+                    "lmeval/num_errored_requests": self.num_errored_requests,
+                    "lmeval/num_tasks": len(results.get("results", {})),
+                }
+            )
+
+        except Exception as e:
+            from veeksha.logger import init_logger
+
+            init_logger(__name__).warning("Failed to log lm-eval wandb metrics: %s", e)
 
     # ---------------------------------------------------------------------
     # Progress tracking

veeksha/microbench/config.py

@@ -5,6 +5,7 @@
 from vidhi import BasePolyConfig, field, frozen_dataclass
 
 from veeksha.cli.base import VeekshaCommand
+from veeksha.config.wandb import WandbConfig  # noqa: F401
 
 
 class StressTrafficMode(StrEnum):
@@ -46,6 +47,10 @@ class BaseMicrobenchmarkConfig:
         False, help="Skip benchmark, only validate existing output"
     )
     skip_validation: bool = field(False, help="Skip post-run validation")
+    wandb: WandbConfig = field(
+        default_factory=WandbConfig,
+        help="Weights & Biases logging configuration.",
+    )
 
 
 @frozen_dataclass
@@ -175,6 +180,9 @@ class StressMicrobenchmarkConfig(
     max_tokens_per_second_estimate: float = field(
         500.0, help="Estimated max output tok/s (for session budget)"
     )
+    num_gpus: int = field(
+        0, help="Number of GPUs serving the model (0 = unknown, TPS/GPU not computed)"
+    )
     mode: BaseStressModeConfig = field(
         default_factory=ManualStressModeConfig,
         help="Stress mode configuration",

veeksha/microbench/diff.py

@@ -471,6 +471,65 @@ def _plot_stress(data1: dict, data2: dict, label1: str, label2: str, out: str) -
     fig.savefig(os.path.join(out, "output_throughput_vs_interactivity.png"), dpi=150)
     plt.close(fig)
 
+    # 8. TPS/GPU vs Load (if either run has per-GPU data)
+    has_tps1 = any(r.get("output_tps_per_gpu", 0) > 0 for r in results1)
+    has_tps2 = any(r.get("output_tps_per_gpu", 0) > 0 for r in results2)
+    if has_tps1 or has_tps2:
+        fig, ax = plt.subplots(figsize=(8, 5))
+        if has_tps1:
+            ax.plot(
+                levels1,
+                [r.get("output_tps_per_gpu", 0) for r in results1],
+                "o-",
+                label=f"{label1}",
+                linewidth=2,
+            )
+        if has_tps2:
+            ax.plot(
+                levels2,
+                [r.get("output_tps_per_gpu", 0) for r in results2],
+                "s--",
+                label=f"{label2}",
+                linewidth=2,
+            )
+        ax.set_xlabel(level_label)
+        ax.set_ylabel("TPS / GPU (tok/s/gpu)")
+        ax.set_title("Output Throughput per GPU vs Load")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+        fig.tight_layout()
+        fig.savefig(os.path.join(out, "tps_per_gpu_vs_load.png"), dpi=150)
+        plt.close(fig)
+
+        # 9. TPS/GPU vs TPS/User comparison
+        fig, ax = plt.subplots(figsize=(8, 5))
+        if has_tps1:
+            ax.plot(
+                [r["interactivity_p50"] for r in results1],
+                [r.get("output_tps_per_gpu", 0) for r in results1],
+                "o-",
+                label=f"{label1}",
+                linewidth=2,
+                markersize=8,
+            )
+        if has_tps2:
+            ax.plot(
+                [r["interactivity_p50"] for r in results2],
+                [r.get("output_tps_per_gpu", 0) for r in results2],
+                "s--",
+                label=f"{label2}",
+                linewidth=2,
+                markersize=8,
+            )
+        ax.set_xlabel("TPS / User (tok/s/user)")
+        ax.set_ylabel("TPS / GPU (tok/s/gpu)")
+        ax.set_title("Throughput Curve: TPS/GPU vs TPS/User")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+        fig.tight_layout()
+        fig.savefig(os.path.join(out, "tps_per_gpu_vs_tps_per_user.png"), dpi=150)
+        plt.close(fig)
+
 
 # ---------------------------------------------------------------------------
 # Dispatcher

veeksha/microbench/runner.py

@@ -8,12 +8,18 @@
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
+from vidhi import dataclass_to_dict
 
 from veeksha.cli.benchmarks import BenchmarkCliRunner
 from veeksha.config.benchmark import BenchmarkConfig
 from veeksha.logger import init_logger
 from veeksha.microbench.common import ValidationResult
 from veeksha.microbench.config import BaseMicrobenchmarkConfig
+from veeksha.wandb_integration import (
+    maybe_finish_microbench_wandb_run,
+    maybe_init_microbench_wandb_run,
+    maybe_log_microbench_results,
+)
 
 logger = init_logger(__name__)
 console = Console()
@@ -31,11 +37,28 @@ def run(
     cfg = _make_run_dir(cfg, type_name)
     _print_banner(cfg, type_name, banner_rows)
 
-    if not cfg.validate_only:
-        benchmark_configs: list[BenchmarkConfig] = build_benchmark_configs(cfg)
-        BenchmarkCliRunner(benchmark_configs).run_all()
+    maybe_init_microbench_wandb_run(
+        cfg.wandb,
+        microbench_type=type_name,
+        output_dir=cfg.output_dir,
+        config_dict=dataclass_to_dict(cfg),
+    )
+
+    try:
+        if not cfg.validate_only:
+            benchmark_configs: list[BenchmarkConfig] = build_benchmark_configs(cfg)
+            BenchmarkCliRunner(benchmark_configs).run_all()
 
-    print_results_table(cfg)
+        print_results_table(cfg)
+
+        maybe_log_microbench_results(
+            cfg.wandb,
+            microbench_type=type_name,
+            results_json_path=os.path.join(cfg.output_dir, f"{type_name}_results.json"),
+            output_dir=cfg.output_dir,
+        )
+    finally:
+        maybe_finish_microbench_wandb_run(cfg.wandb, cfg.output_dir)
 
     if not cfg.skip_validation:
         result: ValidationResult = validate(cfg, cfg.output_dir)