Spaces:
Sleeping
Sleeping
| """ | |
| observability/wandb_tracker.py | |
| ------------------------------- | |
| Weights & Biases experiment tracker for AutoDevAgent. | |
| Logs every benchmark run to W&B so results are visible as interactive | |
| charts, comparable across runs, and exportable for the README and | |
| LinkedIn posts. | |
| What gets logged: | |
| - Per-task metrics: success, iterations, tokens, exec time | |
| - Aggregate metrics: success rate, averages, per-language breakdown | |
| - A W&B Table with one row per benchmark task (visible as a chart) | |
| - Session history stats from the Gradio UI session | |
| Design: | |
| - All W&B calls are wrapped in try/except — a W&B failure must never | |
| crash the pipeline or the evaluation harness. | |
| - When W&B is disabled (enable_wandb=False or API key missing), all | |
| functions are transparent no-ops. | |
| - Uses wandb.init() with reinit=True so multiple benchmark runs in | |
| the same process don't conflict. | |
| Usage: | |
| from observability.wandb_tracker import WandbTracker | |
| tracker = WandbTracker() | |
| tracker.log_benchmark_run(results, summary) | |
| tracker.finish() | |
| """ | |
| import logging | |
| from typing import Any | |
| from config import settings | |
| logger = logging.getLogger(__name__) | |
| # ------------------------------------------------------------------ # | |
| # Tracker # | |
| # ------------------------------------------------------------------ # | |
| class WandbTracker: | |
| """ | |
| Handles all Weights & Biases logging for AutoDevAgent. | |
| Wraps every W&B call in error handling so tracking failures are | |
| never fatal. When W&B is disabled or unavailable, all methods | |
| are silent no-ops. | |
| Attributes: | |
| _active: True if W&B was successfully initialised. | |
| _run: The active wandb.Run object, or None. | |
| """ | |
| def __init__(self) -> None: | |
| """Attempt to initialise W&B. Sets _active=False on failure.""" | |
| self._active = False | |
| self._run = None | |
| if not settings.enable_wandb: | |
| logger.debug("WandbTracker: disabled in config") | |
| return | |
| if not settings.wandb_api_key: | |
| logger.debug("WandbTracker: WANDB_API_KEY not set — tracking disabled") | |
| return | |
| try: | |
| import wandb | |
| import os | |
| os.environ.setdefault("WANDB_API_KEY", settings.wandb_api_key) | |
| self._run = wandb.init( | |
| project = settings.wandb_project, | |
| name = "benchmark-run", | |
| reinit = True, | |
| tags = ["autodevagent", "benchmark"], | |
| config = { | |
| "groq_model_primary": settings.groq_model_primary, | |
| "groq_model_fast": settings.groq_model_fast, | |
| "max_debug_retries": settings.max_debug_retries, | |
| "flowchart_min_lines": settings.flowchart_min_lines, | |
| }, | |
| ) | |
| self._active = True | |
| logger.info( | |
| "WandbTracker: initialised — project: %s run: %s", | |
| settings.wandb_project, | |
| self._run.name, | |
| ) | |
| except ImportError: | |
| logger.warning( | |
| "WandbTracker: wandb not installed — " | |
| "run `pip install wandb` to enable tracking" | |
| ) | |
| except Exception as e: | |
| logger.warning("WandbTracker: init failed (non-fatal): %s", e) | |
| # ---------------------------------------------------------------- # | |
| # Public API # | |
| # ---------------------------------------------------------------- # | |
| def log_benchmark_run( | |
| self, | |
| results: list, | |
| summary: Any, | |
| ) -> None: | |
| """ | |
| Log the full benchmark run results to W&B. | |
| Logs: | |
| 1. Aggregate scalar metrics (success rate, averages) | |
| 2. Per-task metrics as individual W&B log calls | |
| 3. A W&B Table with all task results for visual comparison | |
| Args: | |
| results: List of BenchmarkResult objects from the runner. | |
| summary: MetricsSummary from compute_metrics(). | |
| """ | |
| if not self._active: | |
| return | |
| try: | |
| import wandb | |
| # ── 1. Aggregate metrics ───────────────────────────────── # | |
| self._run.log(summary.to_dict()) | |
| logger.info( | |
| "WandbTracker: logged aggregate metrics — " | |
| "success_rate=%.2f avg_iterations=%.1f avg_tokens=%.0f", | |
| summary.success_rate, | |
| summary.avg_iterations, | |
| summary.avg_tokens, | |
| ) | |
| # ── 2. Per-task metrics ────────────────────────────────── # | |
| for result in results: | |
| self._run.log({ | |
| f"task/{result.task_name}/success": int(result.success), | |
| f"task/{result.task_name}/iterations": result.iterations, | |
| f"task/{result.task_name}/tokens": result.total_tokens, | |
| f"task/{result.task_name}/exec_time": result.exec_time, | |
| }) | |
| # ── 3. Results table ──────────────────────────────────── # | |
| table = wandb.Table( | |
| columns=[ | |
| "task_name", "language", "category", | |
| "success", "iterations", "total_tokens", | |
| "exec_time", "final_status", "tests_passed", | |
| ] | |
| ) | |
| for r in results: | |
| table.add_data( | |
| r.task_name, | |
| r.language, | |
| r.category, | |
| r.success, | |
| r.iterations, | |
| r.total_tokens, | |
| r.exec_time, | |
| r.final_status, | |
| r.tests_passed, | |
| ) | |
| self._run.log({"benchmark_results": table}) | |
| logger.info("WandbTracker: logged results table with %d rows", len(results)) | |
| except Exception as e: | |
| logger.warning("WandbTracker.log_benchmark_run failed (non-fatal): %s", e) | |
| def log_pipeline_run(self, state: Any) -> None: | |
| """ | |
| Log a single pipeline run's metrics to W&B. | |
| Called from app.py after each user-submitted task completes. | |
| Logs lightweight per-run metrics without starting a new W&B run. | |
| Args: | |
| state: The final PipelineState after the pipeline run. | |
| """ | |
| if not self._active: | |
| return | |
| try: | |
| self._run.log({ | |
| "pipeline/language": state.language.value, | |
| "pipeline/status": state.status.value, | |
| "pipeline/debug_iterations": state.debug_iterations, | |
| "pipeline/total_tokens": state.token_usage.total_tokens, | |
| "pipeline/exec_time": state.elapsed_time(), | |
| "pipeline/success": int(state.status.value == "success"), | |
| }) | |
| except Exception as e: | |
| logger.warning("WandbTracker.log_pipeline_run failed (non-fatal): %s", e) | |
| def log_session_summary(self, summary_stats: dict[str, Any]) -> None: | |
| """ | |
| Log end-of-session aggregate stats to W&B. | |
| Called from app.py when the user ends or resets their session. | |
| Args: | |
| summary_stats: Dict from SessionHistory.summary_stats(). | |
| """ | |
| if not self._active: | |
| return | |
| try: | |
| self._run.log({ | |
| f"session/{k}": v for k, v in summary_stats.items() | |
| }) | |
| logger.info("WandbTracker: logged session summary stats") | |
| except Exception as e: | |
| logger.warning("WandbTracker.log_session_summary failed (non-fatal): %s", e) | |
| def finish(self) -> None: | |
| """ | |
| Close the W&B run cleanly. | |
| Call this after all logging is complete to ensure W&B uploads | |
| all data before the process exits. | |
| """ | |
| if not self._active or not self._run: | |
| return | |
| try: | |
| self._run.finish() | |
| self._active = False | |
| logger.info("WandbTracker: run finished and closed") | |
| except Exception as e: | |
| logger.warning("WandbTracker.finish failed (non-fatal): %s", e) | |
| def is_active(self) -> bool: | |
| """Return True if W&B tracking is active.""" | |
| return self._active | |
| def run_url(self) -> str: | |
| """ | |
| Return the W&B dashboard URL for the current run. | |
| Returns: | |
| URL string, or empty string if tracking is inactive. | |
| """ | |
| if not self._active or not self._run: | |
| return "" | |
| try: | |
| return self._run.get_url() or "" | |
| except Exception: | |
| return "" | |