Tracking API Reference¶
The tracking module implements experiment management. The DualTracker allows for simultaneous logging to local JSON/TensorBoard files and a remote or local MLflow server, ensuring both real-time visualization and long-term experiment comparison.
base ¶
Base class and registry for experiment trackers.
Classes¶
BaseTracker ¶
Bases: ABC
Abstract base for experiment tracking.
All trackers must implement: - start_run: Begin a new experiment run - log_params: Log hyperparameters - log_metrics: Log metrics (optionally at specific step) - log_artifact: Log file artifacts - end_run: Finalize the run
Functions¶
end_run
abstractmethod
¶
log_artifact
abstractmethod
¶
Log a file artifact.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path |
Path
|
Path to the artifact file |
required |
name |
Optional[str]
|
Optional name for the artifact |
None
|
log_metrics
abstractmethod
¶
Log metrics (optionally at a specific step).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics |
Dict[str, float]
|
Dictionary of metric name to value |
required |
step |
Optional[int]
|
Optional step/epoch number |
None
|
Source code in src/tracking/base.py
log_params
abstractmethod
¶
Log hyperparameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
params |
Dict[str, Any]
|
Dictionary of parameters |
required |
start_run
abstractmethod
¶
Start a new run, return run_id.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run_name |
str
|
Name for this run |
required |
config |
Dict[str, Any]
|
Configuration dictionary |
required |
Returns:
| Type | Description |
|---|---|
str
|
Unique run identifier |
TrackerRegistry ¶
Registry for tracker classes.
Functions¶
get
classmethod
¶
Get a tracker instance by name.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
name |
str
|
Registry name |
required |
**kwargs |
Arguments for tracker constructor |
{}
|
Returns:
| Type | Description |
|---|---|
BaseTracker
|
Tracker instance |
Source code in src/tracking/base.py
list_available
classmethod
¶
register
classmethod
¶
Decorator to register a tracker class.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
name |
str
|
Registry name |
required |
Returns:
| Type | Description |
|---|---|
|
Decorator function |
Source code in src/tracking/base.py
local ¶
Local file-based tracking (JSON + TensorBoard).
Classes¶
LocalTracker ¶
Bases: BaseTracker
Local file-based tracking (JSON + TensorBoard).
Outputs to: artifacts/runs/
Example usage
tracker = LocalTracker() run_id = tracker.start_run("experiment_1", {"lr": 0.001}) tracker.log_metrics({"loss": 0.5}, step=1) tracker.end_run()
Initialize the tracker.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
base_dir |
str
|
Base directory for run outputs |
'artifacts/runs'
|
use_tensorboard |
bool
|
Whether to log to TensorBoard |
True
|
Source code in src/tracking/local.py
Functions¶
end_run ¶
End the current run.
Source code in src/tracking/local.py
get_run_dir ¶
log_artifact ¶
Log a file artifact.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path |
Path
|
Path to artifact |
required |
name |
Optional[str]
|
Optional name |
None
|
Source code in src/tracking/local.py
log_figure ¶
Log a matplotlib figure.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
figure |
Matplotlib figure |
required | |
name |
str
|
Figure name |
required |
Source code in src/tracking/local.py
log_metrics ¶
Log metrics.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics |
Dict[str, float]
|
Metric name to value mapping |
required |
step |
Optional[int]
|
Optional step number |
None
|
Source code in src/tracking/local.py
log_params ¶
Log additional parameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
params |
Dict[str, Any]
|
Parameters to log |
required |
Source code in src/tracking/local.py
start_run ¶
Start a new run.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run_name |
str
|
Name for this run |
required |
config |
Dict[str, Any]
|
Configuration dictionary |
required |
Returns:
| Type | Description |
|---|---|
str
|
Run ID |
Source code in src/tracking/local.py
mlflow_tracker ¶
MLflow-based experiment tracking.
Classes¶
MLflowTracker ¶
MLflowTracker(tracking_uri: str = 'file:./artifacts/mlruns', experiment_name: str = 'battery_degradation')
Bases: BaseTracker
MLflow-based tracking.
Provides: - Centralized experiment tracking - Model versioning - Artifact storage - UI for experiment comparison
Example usage
tracker = MLflowTracker(tracking_uri="file:./mlruns") run_id = tracker.start_run("experiment_1", {"lr": 0.001}) tracker.log_metrics({"loss": 0.5}, step=1) tracker.end_run()
Initialize the tracker.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tracking_uri |
str
|
MLflow tracking URI |
'file:./artifacts/mlruns'
|
experiment_name |
str
|
Experiment name |
'battery_degradation'
|
Source code in src/tracking/mlflow_tracker.py
Functions¶
end_run ¶
End the current run.
Source code in src/tracking/mlflow_tracker.py
get_run_id ¶
log_artifact ¶
Log a file artifact.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path |
Path
|
Path to artifact |
required |
name |
Optional[str]
|
Optional subfolder name |
None
|
Source code in src/tracking/mlflow_tracker.py
log_metrics ¶
Log metrics.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics |
Dict[str, float]
|
Metric name to value mapping |
required |
step |
Optional[int]
|
Optional step number |
None
|
Source code in src/tracking/mlflow_tracker.py
log_model ¶
Log a PyTorch model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model |
PyTorch model |
required | |
artifact_path |
str
|
Path in artifact store |
'model'
|
Source code in src/tracking/mlflow_tracker.py
log_params ¶
Log additional parameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
params |
Dict[str, Any]
|
Parameters to log |
required |
Source code in src/tracking/mlflow_tracker.py
start_run ¶
Start a new run.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run_name |
str
|
Name for this run |
required |
config |
Dict[str, Any]
|
Configuration dictionary |
required |
Returns:
| Type | Description |
|---|---|
str
|
Run ID |
Source code in src/tracking/mlflow_tracker.py
dual_tracker ¶
Dual tracker that logs to both local and MLflow simultaneously.
Classes¶
DualTracker ¶
DualTracker(local_base_dir: str = 'artifacts/runs', use_tensorboard: bool = True, mlflow_tracking_uri: str = 'file:./artifacts/mlruns', mlflow_experiment_name: str = 'battery_degradation')
Bases: BaseTracker
Tracker that logs to both local files and MLflow.
Provides the best of both worlds: - Local: Fast access, TensorBoard, works offline - MLflow: Centralized UI, model registry, comparison
Example usage
tracker = DualTracker() run_id = tracker.start_run("experiment_1", {"lr": 0.001}) tracker.log_metrics({"loss": 0.5}, step=1) tracker.end_run()
Initialize dual tracker.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
local_base_dir |
str
|
Base directory for local runs |
'artifacts/runs'
|
use_tensorboard |
bool
|
Whether to enable TensorBoard |
True
|
mlflow_tracking_uri |
str
|
MLflow tracking URI |
'file:./artifacts/mlruns'
|
mlflow_experiment_name |
str
|
MLflow experiment name |
'battery_degradation'
|
Source code in src/tracking/dual_tracker.py
Functions¶
end_run ¶
End run on both backends.
get_run_dir ¶
log_artifact ¶
Log artifact to both backends.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
path |
Path
|
Path to artifact |
required |
name |
Optional[str]
|
Optional name |
None
|
Source code in src/tracking/dual_tracker.py
log_figure ¶
Log matplotlib figure.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
figure |
Matplotlib figure |
required | |
name |
str
|
Figure name |
required |
log_metrics ¶
Log metrics to both backends.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics |
Dict[str, float]
|
Metric name to value mapping |
required |
step |
Optional[int]
|
Optional step number |
None
|
Source code in src/tracking/dual_tracker.py
log_model ¶
Log PyTorch model to MLflow.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model |
PyTorch model |
required | |
artifact_path |
str
|
Artifact path |
'model'
|
Source code in src/tracking/dual_tracker.py
log_params ¶
Log parameters to both backends.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
params |
Dict[str, Any]
|
Parameters to log |
required |
Source code in src/tracking/dual_tracker.py
start_run ¶
Start a run on both backends.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run_name |
str
|
Name for this run |
required |
config |
Dict[str, Any]
|
Configuration dictionary |
required |
Returns:
| Type | Description |
|---|---|
str
|
Local run ID |