From 7e92c0e8b52e9a953e41e98e4dcf716fbd693189 Mon Sep 17 00:00:00 2001 From: Amey Agrawal Date: Mon, 30 Mar 2026 23:36:55 -0400 Subject: [PATCH] feat: add wandb logging for microbenchmarks and lm-eval accuracy Add Weights & Biases integration to all three microbenchmarks (prefill, decode, stress) and the lm-eval accuracy evaluator, bringing them to parity with the full benchmark system's wandb support. Changes: - Add `wandb: WandbConfig` to `BaseMicrobenchmarkConfig` so all microbench CLIs expose `--wandb.enabled`, `--wandb.project`, etc. - Add `num_gpus` to `StressMicrobenchmarkConfig` for per-GPU throughput metrics (TPS/GPU vs TPS/User plots, the Sarvam/InferenceX style). - 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) for prefill/decode/stress. - Wire wandb lifecycle into `runner.run()` (prefill, decode, stress manual/range) and `_run_auto_main()` (stress auto mode). - Add `output_tps_per_gpu` / `input_tps_per_gpu` to StressPointResult, two new stress plots (TPS/GPU vs Load, TPS/GPU vs TPS/User), and TPS/GPU columns in the results 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. Co-Authored-By: Claude Opus 4.6 (1M context) --- veeksha/evaluator/accuracy/base.py | 72 ++++++ veeksha/microbench/config.py | 8 + veeksha/microbench/diff.py | 59 +++++ veeksha/microbench/runner.py | 31 ++- veeksha/microbench/stress.py | 105 ++++++++- veeksha/wandb_integration.py | 358 +++++++++++++++++++++++++++++ 6 files changed, 619 insertions(+), 14 deletions(-) diff --git a/veeksha/evaluator/accuracy/base.py b/veeksha/evaluator/accuracy/base.py index a9960240..18dcb3d2 100644 --- a/veeksha/evaluator/accuracy/base.py +++ b/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 diff --git a/veeksha/microbench/config.py b/veeksha/microbench/config.py index 1d3220ac..48c4714b 100644 --- a/veeksha/microbench/config.py +++ b/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", diff --git a/veeksha/microbench/diff.py b/veeksha/microbench/diff.py index 1b09e4aa..d4a5be6e 100644 --- a/veeksha/microbench/diff.py +++ b/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 diff --git a/veeksha/microbench/runner.py b/veeksha/microbench/runner.py index a59ab06a..3c5d2d8e 100644 --- a/veeksha/microbench/runner.py +++ b/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) diff --git a/veeksha/microbench/stress.py b/veeksha/microbench/stress.py index a6aa91f5..70cc7bb9 100644 --- a/veeksha/microbench/stress.py +++ b/veeksha/microbench/stress.py @@ -57,6 +57,7 @@ ("Traffic mode", "traffic_mode"), ("Point duration", "point_duration"), ("Warmup duration", "warmup_duration"), + ("Num GPUs", "num_gpus"), ] @@ -203,6 +204,8 @@ class StressPointResult: interactivity_p50: float # 1 / tpot_p50 (tok/s/user) interactivity_p99: float # 1 / tpot_p99 (tok/s/user) num_requests: int + output_tps_per_gpu: float = 0.0 # output_throughput / num_gpus (0 = unknown) + input_tps_per_gpu: float = 0.0 # input_throughput / num_gpus (0 = unknown) def _extract_stress_point( @@ -289,6 +292,12 @@ def _collect_results( results.append(point) results.sort(key=lambda r: r.level) + + if cfg.num_gpus > 0: + for r in results: + r.output_tps_per_gpu = r.output_throughput / cfg.num_gpus + r.input_tps_per_gpu = r.input_throughput / cfg.num_gpus + return results @@ -468,6 +477,45 @@ def _save_plots( ) plt.close(fig) + # 8. TPS/GPU vs Load Level (only when num_gpus is known) + if any(r.output_tps_per_gpu > 0 for r in results): + tps_per_gpu = [r.output_tps_per_gpu for r in results] + + fig, ax = plt.subplots(figsize=(8, 5)) + ax.plot(levels, tps_per_gpu, "o-", label="Output TPS/GPU", linewidth=2) + ax.plot( + levels, + [r.input_tps_per_gpu for r in results], + "s--", + label="Input TPS/GPU", + linewidth=2, + ) + ax.set_xlabel(level_label) + ax.set_ylabel("TPS / GPU (tok/s/gpu)") + ax.set_title("Throughput per GPU vs Load") + ax.legend() + ax.grid(True, alpha=0.3) + fig.tight_layout() + fig.savefig(os.path.join(plots_dir, "tps_per_gpu_vs_load.png"), dpi=150) + plt.close(fig) + + # 9. TPS/GPU vs TPS/User (the Sarvam-style throughput curve) + fig, ax = plt.subplots(figsize=(8, 5)) + ax.plot( + [r.interactivity_p50 for r in results], + tps_per_gpu, + "o-", + 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.grid(True, alpha=0.3) + fig.tight_layout() + fig.savefig(os.path.join(plots_dir, "tps_per_gpu_vs_tps_per_user.png"), dpi=150) + plt.close(fig) + def print_results_table(cfg: StressMicrobenchmarkConfig) -> None: """Print stress results table and save JSON, CSV, and plots.""" @@ -479,10 +527,14 @@ def print_results_table(cfg: StressMicrobenchmarkConfig) -> None: "QPS" if cfg.traffic_mode == StressTrafficMode.FIXED_RATE else "Concurrency" ) + show_tps_per_gpu = cfg.num_gpus > 0 + table = Table(title="Stress Results (Throughput vs Latency)") table.add_column(level_label, justify="right", style="cyan") table.add_column("In Tput\n(tok/s)", justify="right") table.add_column("Out Tput\n(tok/s)", justify="right") + if show_tps_per_gpu: + table.add_column("Out TPS\n/GPU", justify="right") table.add_column("E2E P50\n(ms)", justify="right") table.add_column("E2E P99\n(ms)", justify="right") table.add_column("TTFC P50\n(ms)", justify="right") @@ -492,18 +544,25 @@ def print_results_table(cfg: StressMicrobenchmarkConfig) -> None: table.add_column("Reqs", justify="right", style="dim") for r in results: - table.add_row( + row = [ str(r.level), f"{r.input_throughput:.1f}", f"{r.output_throughput:.1f}", - fmt_ms(r.e2e_latency_p50), - fmt_ms(r.e2e_latency_p99), - fmt_ms(r.ttfc_p50), - fmt_ms(r.ttfc_p99), - f"{r.interactivity_p50:.1f}", - f"{r.interactivity_p99:.1f}", - str(r.num_requests), + ] + if show_tps_per_gpu: + row.append(f"{r.output_tps_per_gpu:.1f}") + row.extend( + [ + fmt_ms(r.e2e_latency_p50), + fmt_ms(r.e2e_latency_p99), + fmt_ms(r.ttfc_p50), + fmt_ms(r.ttfc_p99), + f"{r.interactivity_p50:.1f}", + f"{r.interactivity_p99:.1f}", + str(r.num_requests), + ] ) + table.add_row(*row) console.print() console.print(table) @@ -516,6 +575,7 @@ def print_results_table(cfg: StressMicrobenchmarkConfig) -> None: "traffic_mode": str(cfg.traffic_mode), "input_length": cfg.input_length, "output_length": cfg.output_length, + "num_gpus": cfg.num_gpus, "results": [asdict(r) for r in results], }, os.path.join(cfg.output_dir, "stress_results.json"), @@ -764,7 +824,14 @@ def _probe(c: int) -> StressPointResult | None: def _run_auto_main(cfg: StressMicrobenchmarkConfig) -> None: """Full auto mode entrypoint.""" + from vidhi import dataclass_to_dict + from veeksha.microbench.runner import _make_run_dir, _print_banner + from veeksha.wandb_integration import ( + maybe_finish_microbench_wandb_run, + maybe_init_microbench_wandb_run, + maybe_log_microbench_results, + ) assert isinstance(cfg.mode, AutoStressModeConfig) if cfg.mode.resume_dir: @@ -774,9 +841,27 @@ def _run_auto_main(cfg: StressMicrobenchmarkConfig) -> None: cfg = _make_run_dir(cfg, "stress") # type: ignore[assignment] _print_banner(cfg, "stress (auto)", BANNER_ROWS) - _run_auto_sweep(cfg) + maybe_init_microbench_wandb_run( + cfg.wandb, + microbench_type="stress", + output_dir=cfg.output_dir, + config_dict=dataclass_to_dict(cfg), + extra_tags=["auto-mode"], + ) + + try: + _run_auto_sweep(cfg) + + print_results_table(cfg) - print_results_table(cfg) + maybe_log_microbench_results( + cfg.wandb, + microbench_type="stress", + results_json_path=os.path.join(cfg.output_dir, "stress_results.json"), + output_dir=cfg.output_dir, + ) + finally: + maybe_finish_microbench_wandb_run(cfg.wandb, cfg.output_dir) if not cfg.skip_validation: result = validate(cfg, cfg.output_dir) diff --git a/veeksha/wandb_integration.py b/veeksha/wandb_integration.py index 4dbc899c..9ebd905a 100644 --- a/veeksha/wandb_integration.py +++ b/veeksha/wandb_integration.py @@ -610,3 +610,361 @@ def update_run_tags(output_dir: str, tags: Iterable[str]) -> None: logger.info("Updated tags for run %s: %s", run_path, tags) except Exception as exc: logger.warning("Failed to update tags for run %s: %s", run_path, exc) + + +# --------------------------------------------------------------------------- +# Microbenchmark wandb helpers +# --------------------------------------------------------------------------- + + +def maybe_init_microbench_wandb_run( + wandb_cfg: "WandbConfig", + *, + microbench_type: str, + output_dir: str, + config_dict: Dict[str, Any], + extra_tags: Optional[Iterable[str]] = None, +) -> None: + """Initialize a wandb run for a microbenchmark (prefill/decode/stress). + + Safe no-op when wandb is disabled or unavailable. + """ + if not wandb_cfg.enabled: + return + + try: + from typing import Any as _Any + + import wandb # type: ignore[import-not-found] + + wandb = cast(_Any, wandb) + except Exception as exc: + raise ImportError( + "Weights & Biases logging is enabled, but `wandb` could not be imported. " + "Install wandb (e.g. `pip install wandb`) or disable `wandb.enabled`." + ) from exc + + if getattr(wandb, "run", None): + logger.info("wandb.run already initialized; skipping wandb.init()") + return + + config_dict = cast(Dict[str, Any], _scrub_secrets(config_dict)) + wandb_cfg_dict: Dict[str, Any] = { + "run_kind": f"microbench-{microbench_type}", + **config_dict, + } + + run_name = wandb_cfg.run_name or os.path.basename(output_dir.rstrip("/")) + tags = dedup_tags( + [*wandb_cfg.tags, "microbench", microbench_type, *(extra_tags or [])] + ) + + logger.info( + "Initializing wandb run for microbench-%s (project=%s name=%s)", + microbench_type, + wandb_cfg.project, + run_name, + ) + + wandb.init( + project=wandb_cfg.project, + entity=wandb_cfg.entity, + group=wandb_cfg.group, + name=run_name, + tags=tags or None, + notes=wandb_cfg.notes, + dir=output_dir, + config=wandb_cfg_dict, + reinit="finish_previous", + mode=wandb_cfg.mode, + ) + + +def _log_microbench_prefill(results: Dict[str, Any], wandb: Any) -> None: + """Log prefill microbenchmark results to wandb.""" + rows = results.get("results", []) + if not rows: + return + + # Scalars per input length + for row in rows: + il = row["input_length"] + ttfc = row.get("ttfc", {}) + wandb.log( + { + "prefill/input_length": il, + "prefill/ttfc_p50_ms": ttfc.get("median", 0) * 1000, + "prefill/ttfc_p99_ms": ttfc.get("p99", 0) * 1000, + "prefill/ttfc_mean_ms": ttfc.get("mean", 0) * 1000, + } + ) + + # Table + columns = [ + "input_length", + "mean_ms", + "p50_ms", + "p99_ms", + "min_ms", + "max_ms", + "count", + ] + data = [] + for row in rows: + s = row.get("ttfc", {}) + data.append( + [ + row["input_length"], + round(s.get("mean", 0) * 1000, 2), + round(s.get("median", 0) * 1000, 2), + round(s.get("p99", 0) * 1000, 2), + round(s.get("min", 0) * 1000, 2), + round(s.get("max", 0) * 1000, 2), + s.get("count", 0), + ] + ) + wandb.log({"prefill_results": wandb.Table(columns=columns, data=data)}) + + +def _log_microbench_decode(results: Dict[str, Any], wandb: Any) -> None: + """Log decode microbenchmark results to wandb.""" + rows = results.get("results", []) + if not rows: + return + + # Scalars per (batch_size, input_length) + for row in rows: + bs = row["batch_size"] + il = row["input_length"] + tbt = row.get("tbt", {}) + wandb.log( + { + "decode/batch_size": bs, + "decode/input_length": il, + "decode/tbt_p50_ms": tbt.get("median", 0) * 1000, + "decode/tbt_p99_ms": tbt.get("p99", 0) * 1000, + "decode/tbt_mean_ms": tbt.get("mean", 0) * 1000, + } + ) + + # Table + columns = [ + "batch_size", + "input_length", + "mean_ms", + "p50_ms", + "p99_ms", + "min_ms", + "max_ms", + "samples", + ] + data = [] + for row in rows: + s = row.get("tbt", {}) + data.append( + [ + row["batch_size"], + row["input_length"], + round(s.get("mean", 0) * 1000, 2), + round(s.get("median", 0) * 1000, 2), + round(s.get("p99", 0) * 1000, 2), + round(s.get("min", 0) * 1000, 2), + round(s.get("max", 0) * 1000, 2), + s.get("count", 0), + ] + ) + wandb.log({"decode_results": wandb.Table(columns=columns, data=data)}) + + +def _log_microbench_stress(results: Dict[str, Any], wandb: Any) -> None: + """Log stress microbenchmark results to wandb.""" + rows = results.get("results", []) + if not rows: + return + + num_gpus = results.get("num_gpus", 0) + + # Step-indexed scalars (one wandb.log per concurrency level) + for r in rows: + metrics: Dict[str, Any] = { + "stress/level": r["level"], + "stress/output_throughput": r["output_throughput"], + "stress/input_throughput": r["input_throughput"], + "stress/e2e_latency_p50_ms": r["e2e_latency_p50"] * 1000, + "stress/e2e_latency_p99_ms": r["e2e_latency_p99"] * 1000, + "stress/ttfc_p50_ms": r["ttfc_p50"] * 1000, + "stress/ttfc_p99_ms": r["ttfc_p99"] * 1000, + "stress/interactivity_p50": r["interactivity_p50"], + "stress/interactivity_p99": r["interactivity_p99"], + "stress/num_requests": r["num_requests"], + } + if num_gpus > 0: + metrics["stress/output_tps_per_gpu"] = r.get( + "output_tps_per_gpu", r["output_throughput"] / num_gpus + ) + metrics["stress/input_tps_per_gpu"] = r.get( + "input_tps_per_gpu", r["input_throughput"] / num_gpus + ) + wandb.log(metrics, step=r["level"]) + + # Table + columns = [ + "level", + "output_throughput", + "input_throughput", + "e2e_p50_ms", + "e2e_p99_ms", + "ttfc_p50_ms", + "ttfc_p99_ms", + "interactivity_p50", + "interactivity_p99", + "num_requests", + ] + if num_gpus > 0: + columns.extend(["output_tps_per_gpu", "input_tps_per_gpu"]) + + data = [] + for r in rows: + row = [ + r["level"], + round(r["output_throughput"], 1), + round(r["input_throughput"], 1), + round(r["e2e_latency_p50"] * 1000, 1), + round(r["e2e_latency_p99"] * 1000, 1), + round(r["ttfc_p50"] * 1000, 1), + round(r["ttfc_p99"] * 1000, 1), + round(r["interactivity_p50"], 1), + round(r["interactivity_p99"], 1), + r["num_requests"], + ] + if num_gpus > 0: + row.extend( + [ + round( + r.get("output_tps_per_gpu", r["output_throughput"] / num_gpus), + 1, + ), + round( + r.get("input_tps_per_gpu", r["input_throughput"] / num_gpus), 1 + ), + ] + ) + data.append(row) + wandb.log({"stress_results": wandb.Table(columns=columns, data=data)}) + + # Summary + if rows: + peak = max(rows, key=lambda r: r["output_throughput"]) + wandb.run.summary["peak_output_throughput"] = peak["output_throughput"] + wandb.run.summary["interactivity_at_peak"] = peak["interactivity_p50"] + wandb.run.summary["peak_concurrency"] = peak["level"] + if num_gpus > 0: + wandb.run.summary["peak_tps_per_gpu"] = peak["output_throughput"] / num_gpus + + +_MICROBENCH_LOGGERS = { + "prefill": _log_microbench_prefill, + "decode": _log_microbench_decode, + "stress": _log_microbench_stress, +} + + +def maybe_log_microbench_results( + wandb_cfg: "WandbConfig", + *, + microbench_type: str, + results_json_path: str, + output_dir: str, +) -> None: + """Log microbenchmark results (tables, scalars, plots, artifacts) to wandb. + + Safe no-op when wandb is disabled or no active run. + """ + if not wandb_cfg.enabled: + return + + 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 + except Exception: + return + + # Log type-specific metrics + if os.path.exists(results_json_path): + with open(results_json_path, "r", encoding="utf-8") as f: + results = json.load(f) + + log_fn = _MICROBENCH_LOGGERS.get(microbench_type) + if log_fn: + try: + log_fn(results, wandb) + except Exception as exc: + logger.warning( + "Failed to log %s results to wandb: %s", microbench_type, exc + ) + + # Log plot images + plots_dir = os.path.join(output_dir, "plots") + if os.path.isdir(plots_dir): + for fname in sorted(os.listdir(plots_dir)): + if fname.endswith(".png"): + try: + wandb.log({fname[:-4]: wandb.Image(os.path.join(plots_dir, fname))}) + except Exception: + pass + + # Upload artifacts + if wandb_cfg.log_artifacts: + artifact = wandb.Artifact( + name=f"microbench-{microbench_type}-{wandb.run.id}", + type=f"veeksha-microbench-{microbench_type}", + ) + has_entries = False + for fname in sorted(os.listdir(output_dir)): + fpath = os.path.join(output_dir, fname) + if os.path.isfile(fpath) and any( + fname.endswith(ext) for ext in (".json", ".csv") + ): + artifact.add_file(fpath, name=fname) + has_entries = True + if os.path.isdir(plots_dir): + for fname in sorted(os.listdir(plots_dir)): + fpath = os.path.join(plots_dir, fname) + if os.path.isfile(fpath): + artifact.add_file(fpath, name=f"plots/{fname}") + has_entries = True + if has_entries: + wandb.log_artifact(artifact) + + +def maybe_finish_microbench_wandb_run( + wandb_cfg: "WandbConfig", + output_dir: str, +) -> None: + """Finish a microbenchmark wandb run. + + Safe no-op when wandb is disabled or no active run. + """ + if not wandb_cfg.enabled: + return + + 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 + maybe_persist_wandb_run_info(output_dir) + try: + wandb.finish() + except Exception as exc: + logger.warning("wandb.finish() failed: %s", exc) + except Exception: + return