From 3ce107a5606caef50f1345ca57a8c4dedc4281f9 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Tue, 31 Mar 2026 19:49:04 +0000 Subject: [PATCH] feat(experiment-logger): add standalone logging framework Create reusable experiment logging package with: - UnifiedLogger for multi-backend logging (WandB, disk, stdout) - Automatic checkpoint and model saving with metadata - WandB artifact upload support - Graceful degradation when WandB unavailable - Comprehensive API documentation This is a standalone, project-agnostic library that can be reused across different ML projects. --- src/experiment_logger/README.md | 145 +++++++++++++ src/experiment_logger/__init__.py | 11 + src/experiment_logger/unified_logger.py | 270 ++++++++++++++++++++++++ src/experiment_logger/wandb_utils.py | 69 ++++++ 4 files changed, 495 insertions(+) create mode 100644 src/experiment_logger/README.md create mode 100644 src/experiment_logger/__init__.py create mode 100644 src/experiment_logger/unified_logger.py create mode 100644 src/experiment_logger/wandb_utils.py diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md new file mode 100644 index 0000000..2beac1a --- /dev/null +++ b/src/experiment_logger/README.md @@ -0,0 +1,145 @@ +# Experiment Logger + +A lightweight, standalone logging framework for machine learning experiments with multi-backend support. + +## Features + +- **Multi-backend logging**: Simultaneously log to WandB, local disk (JSON), and stdout +- **Data preservation**: All metrics saved locally, even if WandB is unavailable +- **Checkpoint management**: Save model checkpoints with metadata +- **WandB integration**: Optional artifact upload for model versioning +- **Graceful degradation**: Works without WandB installed +- **Simple API**: Minimal configuration required + +## Installation + +This package is included in the project. To use it in your code: + +```python +from experiment_logger import UnifiedLogger +``` + +## Quick Start + +```python +from experiment_logger import UnifiedLogger + +# Initialize logger +logger = UnifiedLogger( + run_name="my_experiment", + config={"learning_rate": 0.001, "batch_size": 32}, + project_name="MyProject", + entity="my-wandb-username", # Optional + use_wandb=True, # Set to False to disable WandB +) + +# Log metrics +for step in range(100): + logger.log({ + "loss": 1.0 / (step + 1), + "accuracy": step * 0.01, + }, step=step) + +# Save checkpoint +logger.save_checkpoint( + params=model_params, + step=100, + metadata={"epoch": 1, "val_acc": 0.95}, +) + +# Save final model +logger.save_final_model( + params=final_params, + metadata={"final_accuracy": 0.98}, +) + +# Finalize (flushes remaining metrics) +logger.finish() +``` + +## Context Manager + +Use as a context manager for automatic cleanup: + +```python +with UnifiedLogger(run_name="my_exp", config={}) as logger: + logger.log({"metric": 1.0}) + # Automatically calls finish() on exit +``` + +## Configuration + +### Constructor Parameters + +- `run_name` (str): Unique name for this run +- `config` (dict): Configuration dictionary with hyperparameters +- `project_name` (str): WandB project name (default: "PPO-Modularity") +- `entity` (str, optional): WandB entity (team/user name) +- `base_dir` (str): Base directory for local storage (default: "runs") +- `use_wandb` (bool): Enable WandB logging (default: True) +- `save_code` (bool): Save code to WandB (default: True) + +### Directory Structure + +``` +runs/ +└── my_experiment/ + ├── config.json # Saved configuration + ├── metrics/ + │ └── metrics.jsonl # Line-delimited JSON metrics + ├── checkpoints/ + │ ├── checkpoint_step_100.flax + │ └── checkpoint_step_100_metadata.json + └── final_model.flax +``` + +## API Reference + +### `log(metrics, step=None, commit=True)` + +Log metrics to all backends. + +**Parameters:** +- `metrics` (dict): Dictionary of metric name -> value +- `step` (int, optional): Global step counter (auto-incremented if None) +- `commit` (bool): Whether to commit to WandB immediately + +### `save_checkpoint(params, step, prefix="checkpoint", metadata=None)` + +Save model checkpoint to disk and optionally to WandB. + +**Parameters:** +- `params`: Model parameters (Flax params or any serializable object) +- `step` (int): Current training step +- `prefix` (str): Prefix for checkpoint filename +- `metadata` (dict, optional): Additional metadata to save + +### `save_final_model(params, metadata=None)` + +Save the final trained model. + +**Parameters:** +- `params`: Model parameters +- `metadata` (dict, optional): Metadata about the final model + +### `finish()` + +Finalize logging and cleanup. Flushes remaining metrics to disk. + +## Usage in Projects + +This logger is designed to be: +- **Project-agnostic**: Use in any ML project, not just this one +- **Framework-agnostic**: Works with JAX, PyTorch, TensorFlow, etc. +- **Minimal dependencies**: Only requires `wandb` (optional), `flax` (for serialization), and `numpy` + +## Design Philosophy + +1. **Never lose data**: All metrics saved locally, regardless of WandB availability +2. **Simple API**: Minimal boilerplate, easy to integrate +3. **Fail gracefully**: Missing WandB shouldn't break experiments +4. **Reproducibility**: Save full configuration with every run + +## License + +Part of the 2026SEL3-project-BrittleStar repository. diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py new file mode 100644 index 0000000..7ac168d --- /dev/null +++ b/src/experiment_logger/__init__.py @@ -0,0 +1,11 @@ +"""Unified logging framework for machine learning experiments. + +This package provides a unified interface for logging to multiple backends +(WandB, disk, stdout) simultaneously, ensuring no data loss. +""" + +from experiment_logger.unified_logger import UnifiedLogger +from experiment_logger.wandb_utils import finish_wandb, init_wandb + +__all__ = ["UnifiedLogger", "init_wandb", "finish_wandb"] +__version__ = "0.1.0" diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py new file mode 100644 index 0000000..87e3006 --- /dev/null +++ b/src/experiment_logger/unified_logger.py @@ -0,0 +1,270 @@ +"""Unified logger that writes to multiple backends simultaneously. + +This logger ensures all experimental data is preserved by writing to: +1. Weights & Biases (when available) +2. Local disk (JSON files, model checkpoints) +3. stdout (for real-time monitoring) +""" + +import json +import logging +import time +from pathlib import Path +from typing import Any, Dict, Optional + +import flax +import numpy as np + +from experiment_logger.wandb_utils import finish_wandb, init_wandb + +logger = logging.getLogger(__name__) + + +class UnifiedLogger: + """Unified logger for scientific experiments with redundant backup.""" + + def __init__( + self, + run_name: str, + config: Dict[str, Any], + project_name: str = "PPO-Modularity", + entity: Optional[str] = None, + base_dir: str = "runs", + use_wandb: bool = True, + save_code: bool = True, + ): + """Initialize the unified logger. + + Args: + run_name: Unique name for this run + config: Configuration dictionary with hyperparameters + project_name: WandB project name + entity: WandB entity (team/user name) + base_dir: Base directory for local storage + use_wandb: Whether to use WandB logging + save_code: Whether to save code to WandB + """ + self.run_name = run_name + self.config = config + self.use_wandb = use_wandb + self.wandb_available = False + self.wandb_run = None + + # Setup local storage + self.run_dir = Path(base_dir) / run_name + self.run_dir.mkdir(parents=True, exist_ok=True) + + self.checkpoints_dir = self.run_dir / "checkpoints" + self.checkpoints_dir.mkdir(exist_ok=True) + + self.metrics_dir = self.run_dir / "metrics" + self.metrics_dir.mkdir(exist_ok=True) + + self.config_file = self.run_dir / "config.json" + + # Save config to disk + self._save_config() + + # Initialize WandB if requested + if self.use_wandb: + self._init_wandb(project_name, entity, save_code) + + # Initialize metrics storage + self.metrics_buffer: list[Dict[str, Any]] = [] + self.step_counter = 0 + + logger.info(f"Initialized for run: {run_name}") + logger.info(f"Local storage: {self.run_dir.absolute()}") + logger.info(f"WandB logging: {self.wandb_available}") + + def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool): + """Initialize Weights & Biases logging.""" + self.wandb_run = init_wandb( + project=project_name, + entity=entity, + name=self.run_name, + config=self.config, + save_code=save_code, + resume="allow", + ) + self.wandb_available = self.wandb_run is not None + + def _save_config(self): + """Save configuration to disk.""" + try: + with open(self.config_file, "w") as f: + json.dump(self.config, f, indent=2) + logger.info(f"Config saved to {self.config_file}") + except Exception as e: + logger.error(f"Error saving config: {e}") + + def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True): + """Log metrics to all backends. + + Args: + metrics: Dictionary of metric name -> value + step: Global step counter (auto-incremented if None) + commit: Whether to commit to WandB immediately + """ + if step is None: + step = self.step_counter + self.step_counter += 1 + + # Add timestamp + metrics_with_metadata = { + "step": step, + "timestamp": time.time(), + **metrics, + } + + # Log to stdout + self._log_to_stdout(metrics_with_metadata) + + # Log to WandB + if self.wandb_available: + try: + self.wandb_run.log(metrics, step=step, commit=commit) + except Exception as e: + logger.warning(f"WandB logging failed: {e}") + + # Buffer for disk storage + self.metrics_buffer.append(metrics_with_metadata) + + # Periodically flush to disk + if len(self.metrics_buffer) >= 100: + self._flush_metrics() + + def _log_to_stdout(self, metrics: Dict[str, Any]): + """Log metrics to stdout for real-time monitoring.""" + step = metrics.get("step", "?") + metric_str = ", ".join( + f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}" + for k, v in metrics.items() + if k not in ["step", "timestamp"] + ) + logger.info(f"[Step {step}] {metric_str}") + + def _flush_metrics(self): + """Flush buffered metrics to disk.""" + if not self.metrics_buffer: + return + + try: + metrics_file = self.metrics_dir / "metrics.jsonl" + with open(metrics_file, "a") as f: + for metric in self.metrics_buffer: + f.write(json.dumps(metric) + "\n") + self.metrics_buffer.clear() + except Exception as e: + logger.error(f"Error flushing metrics: {e}") + + def save_checkpoint( + self, + params: Any, + step: int, + prefix: str = "checkpoint", + metadata: Optional[Dict[str, Any]] = None, + ): + """Save model checkpoint to disk and optionally to WandB. + + Args: + params: Model parameters (Flax params or any serializable object) + step: Current training step + prefix: Prefix for checkpoint filename + metadata: Additional metadata to save with checkpoint + """ + checkpoint_name = f"{prefix}_step_{step}.flax" + checkpoint_path = self.checkpoints_dir / checkpoint_name + + try: + # Save to disk using Flax serialization + with open(checkpoint_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + # Save metadata if provided + if metadata: + metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.json" + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + + logger.info(f"Checkpoint saved: {checkpoint_path}") + + # Log to WandB as artifact + if self.wandb_available: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_{prefix}", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(checkpoint_path)) + if metadata: + artifact.add_file(str(metadata_path)) + self.wandb_run.log_artifact(artifact) + logger.info("Checkpoint uploaded to WandB") + except Exception as e: + logger.warning(f"Could not upload checkpoint to WandB: {e}") + + except Exception as e: + logger.error(f"Error saving checkpoint: {e}") + + def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): + """Save the final trained model. + + Args: + params: Model parameters + metadata: Additional metadata about the final model + """ + final_model_path = self.run_dir / "final_model.flax" + + try: + with open(final_model_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + if metadata: + metadata_path = self.run_dir / "final_model_metadata.json" + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + + logger.info(f"Final model saved: {final_model_path}") + + # Log to WandB + if self.wandb_available: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_final_model", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(final_model_path)) + if metadata: + artifact.add_file(str(metadata_path)) + self.wandb_run.log_artifact(artifact) + except Exception as e: + logger.warning(f"Could not upload final model to WandB: {e}") + + except Exception as e: + logger.error(f"Error saving final model: {e}") + + def finish(self): + """Finalize logging and cleanup.""" + # Flush remaining metrics + self._flush_metrics() + + logger.info(f"Run complete. Results saved to: {self.run_dir.absolute()}") + + # Finish WandB run + if self.wandb_available: + finish_wandb() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.finish() diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py new file mode 100644 index 0000000..5fd8837 --- /dev/null +++ b/src/experiment_logger/wandb_utils.py @@ -0,0 +1,69 @@ +"""Centralized WandB initialization utilities.""" + +import logging +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + + +def init_wandb( + project: str, + config: Dict[str, Any], + name: Optional[str] = None, + entity: Optional[str] = None, + sync_tensorboard: bool = False, + save_code: bool = True, + resume: str = "allow", + **kwargs, +): + """Initialize WandB with standardized settings. + + This function provides a centralized way to initialize WandB across different + scripts, ensuring consistent configuration and error handling. + + Args: + project: WandB project name + config: Configuration dictionary to log + name: Run name (auto-generated if None) + entity: WandB entity (team/user name) + sync_tensorboard: Whether to sync tensorboard logs + save_code: Whether to save code snapshots + resume: Resume strategy ("allow", "must", "never", "auto") + **kwargs: Additional arguments to pass to wandb.init() + + Returns: + wandb.Run object if successful, None otherwise + """ + try: + import wandb + + run = wandb.init( + project=project, + entity=entity, + name=name, + config=config, + sync_tensorboard=sync_tensorboard, + save_code=save_code, + resume=resume, + **kwargs, + ) + logger.info(f"WandB initialized successfully for project '{project}', run '{run.name}'") + return run + except ImportError: + logger.warning("WandB not installed. Skipping WandB initialization.") + return None + except Exception as e: + logger.error(f"Failed to initialize WandB: {e}") + return None + + +def finish_wandb(): + """Safely finish the current WandB run.""" + try: + import wandb + + if wandb.run is not None: + wandb.finish() + logger.info("WandB run finished successfully") + except Exception as e: + logger.warning(f"Error finishing WandB run: {e}")