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.
This commit is contained in:
parent
c27f5fcdf2
commit
3ce107a560
4 changed files with 495 additions and 0 deletions
145
src/experiment_logger/README.md
Normal file
145
src/experiment_logger/README.md
Normal file
|
|
@ -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.
|
||||
11
src/experiment_logger/__init__.py
Normal file
11
src/experiment_logger/__init__.py
Normal file
|
|
@ -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"
|
||||
270
src/experiment_logger/unified_logger.py
Normal file
270
src/experiment_logger/unified_logger.py
Normal file
|
|
@ -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()
|
||||
69
src/experiment_logger/wandb_utils.py
Normal file
69
src/experiment_logger/wandb_utils.py
Normal file
|
|
@ -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}")
|
||||
Reference in a new issue