diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md index 6ec220b..5a01d26 100644 --- a/docs/DEVELOPMENT.md +++ b/docs/DEVELOPMENT.md @@ -59,3 +59,25 @@ Verify your setup by running the JAX initialization test: uv run pytest tests/test_jax_init.py ``` In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is expected if a GPU is detected and the `cuda` extra was installed. + +## Logging & Monitoring + +This project uses a unified logging system through the `experiment_logger` package. For a full API reference, see the [package README](../src/experiment_logger/README.md). + +### Quick Setup + +1. **Authorization**: Export your API key in your terminal to enable WandB synchronization: + ```bash + export WANDB_API_KEY=your_copied_api_key_here + ``` +2. **Toggle Tracking**: Use the `--track` flag in `scripts/train.py` to enable online sync. +3. **Local Monitoring**: All runs are recorded in the `runs/` directory. View scalars with TensorBoard: + ```bash + tensorboard --logdir runs/ + ``` + +### Environment Awareness + +The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job). It will automatically disable progress bars and switch to robust fallback modes (offline logging) to ensure your experiments never hang. + + diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md index ab6394d..ce31d69 100644 --- a/src/experiment_logger/README.md +++ b/src/experiment_logger/README.md @@ -1,17 +1,68 @@ # Experiment Logger -A lightweight logging framework supporting Weights & Biases, local JSON, and stdout. +A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk). -## Usage +This library is designed to be a standalone package that decouples the logging logic from the core training routines in the `brittle_star_project`. + +## Quick Start + +The recommended way to use the logger is through the `get_logger()` singleton: ```python -from experiment_logger import UnifiedLogger +from experiment_logger import get_logger -logger = UnifiedLogger(run_name="my_experiment", config={"lr": 0.001}) +logger = get_logger() +# Initialize at the start of your script (e.g., in train.py) +logger.init( + project_name="MyProject", + run_name="my_experiment_run", + base_dir="runs", + use_wandb=True +) -logger.log({"loss": 0.5}, step=1) -logger.save_checkpoint(params=model_params, step=1) -logger.finish() +# Log metrics (Scalar values, numpy scalars, or JAX types) +logger.log({"loss": 0.5, "accuracy": 0.98}, step=100) + +# Standard logging (Mirrored to disk and stdout) +logger.info("Training started") +logger.warning("Learning rate is very high") + +# Save checkpoints (Automatically synced to WandB as artifacts) +logger.save_checkpoint(params, step=5000) ``` -Logs and checkoints are saved in the `runs/` directory. If `track=True` (or `use_wandb=True`), everything is additionally synced to Weights & Biases. +## Logger Classes + +### `UnifiedLogger` + +The full suite for production training. It manages: +- **WandB**: Syncs metrics and uploads model checkpoints as artifacts. +- **TensorBoard**: Writes events for local visualization. +- **Local Disk**: Stores metrics in `metrics.yaml` and textual logs in `run.log`. + +### `SimpleLogger` + +A zero-dependency fallback that uses standard Python `print()` statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring. + +```python +from experiment_logger import SimpleLogger +logger = SimpleLogger(run_name="test_run") +``` + +## API Features + +### `logger.progress_bar(iterable, **kwargs)` + +A smart wrapper around `tqdm` that automatically detects its environment. +- **Interactive Terminal**: Displays a normal progress bar. +- **Non-Interactive (HPC)**: Automatically disables the bar to prevent log file bloat in `slurm.out`. + +### `logger.log_non_interactive(msg: str)` + +Prints a message *only* when running in non-interactive environments. Useful for high-level progress tracking (e.g., "Epoch 5 Complete") without interactive noise. + +### `logger.save_checkpoint(params, step, prefix="checkpoint")` + +Saves model parameters using Flax serialization. +- **Local Location**: `runs//checkpoints/` +- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking.