wandb

wandb is the official Python client for Weights & Biases, a hosted (and self-hostable) platform for tracking machine learning experiments, versioning datasets and models, and visualizing training runs. A few lines of wandb.init() / wandb.log() calls capture metrics, hyperparameters, system stats, and artifacts from any training loop and stream them to a web dashboard for comparison across runs.

Beyond logging, the SDK ships integrations for the major ML frameworks (PyTorch, Keras, TensorFlow, JAX, Hugging Face, scikit-learn), an artifact system for versioning datasets and model checkpoints, a hyperparameter sweep engine, and a CLI (wandb/wb) for authentication, syncing offline runs, and managing local configuration. It’s built as a thin client around the hosted W&B backend rather than a standalone local tool.

What You Get

• wandb.init()/wandb.log() API for capturing metrics, hyperparameters, and system stats from any training loop
• Framework integrations for PyTorch, Keras, TensorFlow, JAX, Hugging Face Transformers, and scikit-learn
• Artifact versioning for datasets, models, and other large binary outputs, with lineage tracking between runs
• Built-in hyperparameter sweep engine (wandb sweep) for grid/random/Bayesian search across configs
• CLI (wandb/wb) for authentication, offline-run syncing, and local configuration management

Common Use Cases

• Tracking loss curves, accuracy, and other metrics across many training runs for side-by-side comparison
• Versioning datasets and model checkpoints as artifacts with full lineage from raw data to trained model
• Running automated hyperparameter sweeps to find optimal training configurations
• Sharing training results and dashboards with a team for collaborative model development and review

Under The Hood

Architecture — The SDK is organized under wandb/sdk/ with a wandb_run.py (top-level Run object), wandb_controller.py, and dedicated subpackages for filesync, agents (sweep agents), automations, and apis (public/reporting API clients). A significant portion of the performance-critical path — the local sync/networking layer that streams logged data to the backend — is implemented as a separate core component (visible as its own top-level directory alongside the Python package) rather than pure Python, reflecting the project’s evolution from a pure-Python client to a hybrid Python+Go/Rust architecture for throughput. integration/ contains per-framework adapters (PyTorch, Keras, etc.) that hook into each framework’s callback/logging system. Tech Stack — Python 3.10+, packaged with Hatchling. Core runtime dependencies are deliberately minimal (Click for the CLI, requests, sentry-sdk for error reporting, protobuf, PyYAML, platformdirs, pydantic, packaging) with heavy framework integrations (PyTorch, TensorFlow, boto3, kubernetes, etc.) gated behind optional extras (kubeflow, gcp, aws, media, etc.) so a bare install stays lightweight. The core directory and parquet-rust-wrapper indicate Go and Rust components in the sync/data path. Code Quality — Extensive tests/ tree split into unit_tests, system_tests, and fixtures, with a dedicated tests/ruff.toml for test-specific lint rules — indicating a mature, actively maintained test suite distinct from the main codebase’s linting config. BREAKING.md and CHANGELOG.unreleased.md show disciplined change management for a project with a large, versioned public API. API Design — The core API is deliberately minimal (wandb.init(), wandb.log(), wandb.finish()) so the common case requires almost no boilerplate, while power users can reach into wandb.Artifact, wandb.sweep, and framework-specific callback classes for deeper control. package_readme.md (a separate, PyPI-facing README from the repo’s own README.md) shows deliberate attention to first-impression documentation for library consumers versus contributors.