feat(utils): add Trackio experiment tracking backend#1113
Conversation
There was a problem hiding this comment.
Code Review
This pull request integrates Trackio for experiment tracking, adding a TrackioConfig dataclass and incorporating it into the StatsLogger and logging utilities. The changes include dependency updates, documentation generation, and a comprehensive test suite. A review comment suggests implementing validation for the mode field in TrackioConfig to ensure it only accepts 'disabled', 'online', or 'local'.
| class TrackioConfig: | ||
| """Configuration for Trackio experiment tracking (Hugging Face). | ||
|
|
||
| Trackio is a lightweight, local-first experiment tracking library | ||
| with a wandb-compatible API. Dashboards can be viewed locally or | ||
| deployed to Hugging Face Spaces. | ||
|
|
||
| See: https://github.com/gradio-app/trackio | ||
| """ | ||
|
|
||
| mode: str = "disabled" | ||
| """Tracking mode. One of "disabled", "online", or "local".""" | ||
| project: str | None = None | ||
| """Project name. Defaults to experiment_name if not set.""" | ||
| name: str | None = None | ||
| """Run name. Defaults to trial_name if not set.""" | ||
| space_id: str | None = None | ||
| """HF Space ID for remote dashboard deployment (e.g. "user/my-space"). | ||
| When set, metrics are also pushed to the specified Hugging Face Space.""" | ||
|
|
There was a problem hiding this comment.
To improve robustness and prevent unexpected behavior with invalid modes, it's a good practice to validate the mode field. Many other configuration dataclasses in this file (e.g., NormConfig, RecoverConfig) already use a __post_init__ method for validation. I suggest adding one here to ensure mode is one of the expected values ('disabled', 'online', or 'local').
@dataclass
class TrackioConfig:
"""Configuration for Trackio experiment tracking (Hugging Face).
Trackio is a lightweight, local-first experiment tracking library
with a wandb-compatible API. Dashboards can be viewed locally or
deployed to Hugging Face Spaces.
See: https://github.com/gradio-app/trackio
"""
mode: str = "disabled"
"""Tracking mode. One of "disabled", "online", or "local"."""
project: str | None = None
"""Project name. Defaults to experiment_name if not set."""
name: str | None = None
"""Run name. Defaults to trial_name if not set."""
space_id: str | None = None
"""HF Space ID for remote dashboard deployment (e.g. "user/my-space").
When set, metrics are also pushed to the specified Hugging Face Space."""
def __post_init__(self):
"""Validate Trackio configuration."""
valid_modes = {"disabled", "online", "local"}
if self.mode not in valid_modes:
raise ValueError(
f"Invalid trackio mode: '{self.mode}'. Must be one of {valid_modes}."
)Integrate Trackio (Hugging Face) as a new experiment tracking option alongside existing WandB, SwanLab, and TensorBoard backends. Key changes: - Add TrackioConfig dataclass with mode/project/name/space_id fields - Integrate trackio init/log/finish lifecycle in StatsLogger - Add trackio.log() fallback in logging.py helper function - Add trackio to pyproject.toml dependencies - Update CLI docs generator to include TrackioConfig - Add unit tests for config and StatsLogger integration
939633b to
ced8391
Compare
) Integrate Trackio (Hugging Face) as a new experiment tracking option alongside existing WandB, SwanLab, and TensorBoard backends. Key changes: - Add TrackioConfig dataclass with mode/project/name/space_id fields - Integrate trackio init/log/finish lifecycle in StatsLogger - Add trackio.log() fallback in logging.py helper function - Add trackio to pyproject.toml dependencies - Update CLI docs generator to include TrackioConfig - Add unit tests for config and StatsLogger integration
Description
Add Trackio (Hugging Face) as a new experiment tracking backend for AReaL, alongside existing WandB, SwanLab, and TensorBoard support. Trackio is a lightweight, local-first tracking library with a wandb-compatible API that supports local dashboards and Hugging Face Spaces deployment.
Related Issue
#1089
#907
Type of Change
Checklist
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prBreaking Change Details (if applicable):
N/A
Additional Context
Key changes:
areal/api/cli_args.py: NewTrackioConfigdataclass withmode/project/name/space_idfields; addedtrackiofield toStatsLoggerConfigareal/utils/stats_logger.py: Trackioinit/log/finishlifecycle integration inStatsLoggerareal/utils/logging.py: Addtrackio.log()with graceful fallback in helper functionpyproject.toml: Addtrackiodependencydocs/generate_cli_docs.py: IncludeTrackioConfigin CLI docs generationtests/test_trackio_backend.py: Unit tests for config defaults andStatsLoggerintegration (mocked)Trackio defaults to
disabledmode, so this change is fully backward-compatible. Enable withstats_logger.trackio.mode=onlineorstats_logger.trackio.mode=local.