feat: wandb logging for microbenchmarks and lm-eval accuracy

[Anirudhaagrawal]

Summary

• Add wandb: WandbConfig to BaseMicrobenchmarkConfig so all microbench CLIs expose --wandb.* flags (prefill, decode, stress)
• Add num_gpus to StressMicrobenchmarkConfig for TPS/GPU metrics and plots (TPS/GPU vs Load, TPS/GPU vs TPS/User — the Sarvam/InferenceX-style throughput curves)
• Add three lifecycle functions in wandb_integration.py: maybe_init_microbench_wandb_run, maybe_log_microbench_results, maybe_finish_microbench_wandb_run with type-specific logging (tables, scalars, plots, artifacts)
• Wire wandb into runner.run() (all manual/range microbench modes) and _run_auto_main() (stress auto mode)
• Add output_tps_per_gpu / input_tps_per_gpu to StressPointResult, TPS/GPU columns in rich table and JSON output
• Add TPS/GPU comparison plots to veeksha diff for stress results
• Add _log_wandb_metrics() to LMEvalAccuracyEvaluator.save() to log per-task accuracy scalars and results table when a wandb run is active

Test plan

• make lint passes (excluding pre-existing black bug in slo/metrics.py)
• python -m pytest tests/unit/ -x -q passes (213/214 pass; 1 failure is pre-existing env issue: port 8000 in use)
• uvx --python 3.14t veeksha stress --help shows --wandb.* and --num_gpus flags
• uvx --python 3.14t veeksha prefill --help shows --wandb.* flags
• uvx --python 3.14t veeksha decode --help shows --wandb.* flags

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added Weights & Biases integration for automatic logging of accuracy metrics and microbenchmark results
  • Introduced per-GPU throughput metrics for stress testing scenarios
  • Added comparative plots analyzing throughput-per-GPU vs load and vs user throughput
  • Added GPU count configuration to microbenchmark settings

[coderabbitai]

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Free

Run ID: cf6ef186-a60a-4b03-8f93-0a4517d79207

📥 Commits

Reviewing files that changed from the base of the PR and between f250f0c and 7e92c0e.

📒 Files selected for processing (6)

• veeksha/evaluator/accuracy/base.py
• veeksha/microbench/config.py
• veeksha/microbench/diff.py
• veeksha/microbench/runner.py
• veeksha/microbench/stress.py
• veeksha/wandb_integration.py

---

📝 Walkthrough

Walkthrough

This PR integrates Weights & Biases (W&B) logging throughout the evaluator and microbenchmark systems, adds per-GPU throughput metrics to stress benchmarks, and introduces new comparative performance plots. It adds comprehensive W&B helper functions for run initialization, metrics logging, and finalization.

Changes

Cohort / File(s)	Summary
W&B Integration Infrastructure veeksha/wandb_integration.py	Added three microbench-specific W&B helper functions: maybe_init_microbench_wandb_run() for W&B run initialization with config and tags, maybe_log_microbench_results() for logging metrics/tables/artifacts by benchmark type, and maybe_finish_microbench_wandb_run() for run finalization.
Configuration Updates veeksha/microbench/config.py	Added wandb: WandbConfig field to BaseMicrobenchmarkConfig for W&B logging configuration. Added num_gpus: int field to StressMicrobenchmarkConfig to track GPU count for per-GPU throughput calculations.
Evaluator Accuracy Logging veeksha/evaluator/accuracy/base.py	Added _log_wandb_metrics() helper method that logs flattened per-task accuracy metrics and results tables to W&B after writing lmeval results JSON.
Microbenchmark Runner W&B Integration veeksha/microbench/runner.py	Integrated W&B run lifecycle (initialization, result logging, finalization) into run() function with try/finally structure to ensure W&B cleanup regardless of execution outcome.
Stress Benchmark Enhancements veeksha/microbench/stress.py	Added per-GPU TPS fields (output_tps_per_gpu, input_tps_per_gpu) to StressPointResult. Updated results collection to compute per-GPU metrics when num_gpus > 0. Added W&B integration calls for run initialization, result logging, and finalization. Enhanced results table and banner to display GPU count and per-GPU throughput.
Comparative Plot Generation veeksha/microbench/diff.py	Added two new stress-only comparative plots: tps_per_gpu_vs_load.png and tps_per_gpu_vs_tps_per_user.png, conditionally generated when per-GPU output TPS data is available.

Sequence Diagram

sequenceDiagram
    participant Runner as Microbench Runner
    participant WandB as W&B System
    participant Benchmark as Benchmark Executor
    participant Results as Results Logger
    
    Runner->>WandB: maybe_init_microbench_wandb_run()
    activate WandB
    WandB->>WandB: Initialize run with config & tags
    deactivate WandB
    
    Runner->>Benchmark: Execute benchmark
    activate Benchmark
    Benchmark->>Benchmark: Run microbench (prefill/decode/stress)
    Benchmark->>Benchmark: Generate results JSON
    deactivate Benchmark
    
    Runner->>Results: Print results table
    activate Results
    Results-->>Runner: Results displayed
    deactivate Results
    
    Runner->>WandB: maybe_log_microbench_results()
    activate WandB
    WandB->>WandB: Log metrics & tables by type
    WandB->>WandB: Upload artifacts (JSON/CSV/PNG)
    deactivate WandB
    
    Runner->>WandB: maybe_finish_microbench_wandb_run()
    activate WandB
    WandB->>WandB: Persist run metadata & finalize
    deactivate WandB

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

🐰 Hopping through metrics with paws held high,
GPU throughput curves across the sky,
W&B logs each benchmark's dance,
Per-GPU secrets spring from each glance,
Better benchmarks, oh what a prance! ✨

---

Note

🎁 Summarized by CodeRabbit Free

Your organization is on the Free plan. CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please upgrade your subscription to CodeRabbit Pro by visiting https://app.coderabbit.ai/login.

_{Comment @coderabbitai help to get the list of available commands and usage tips.}