Experiment Tracking#

Source file

  • twiga/tracking/tracker.py - TwigaTracker

Overview#

TwigaTracker is a context manager that wraps an MLflow run lifecycle and exposes helpers tuned to Twiga’s data structures - Pydantic configs, metrics DataFrames, and checkpoint directories. It is installed as part of the [mlops] extra.

pip install twiga[mlops]

Quick start#

from twiga.tracking import TwigaTracker

with TwigaTracker(experiment="load-forecast", run_name="lgbm-v1") as tracker:
    forecaster.fit(train_df, val_df)
    _, metrics_df = forecaster.evaluate(test_df)
    tracker.log_metrics(metrics_df)
    tracker.log_forecaster(forecaster)

The block above:

  1. Creates (or resumes) the MLflow experiment "load-forecast".

  2. Starts a new run named "lgbm-v1".

  3. Logs evaluation metrics keyed as {model}/{metric}.

  4. Logs data pipeline params and checkpoint artefacts.

  5. Ends the run with status FINISHED on exit (or FAILED on exception).

API reference#

TwigaTracker#

class twiga.tracking.TwigaTracker(experiment='twiga', run_name=None, tracking_uri=None, tags=None)#

Bases: object

Context manager for tracking Twiga experiments in MLflow.

Provides high-level utilities for logging metadata, models, and evaluation results while maintaining strict MLOps standards for lineage and reproducibility.

__init__(experiment='twiga', run_name=None, tracking_uri=None, tags=None)#

Initializes the tracker and sets the tracking URI.

Parameters:

  • experiment (str) – Name of the MLflow experiment.

  • run_name (str | None) – Optional name for this specific run.

  • tracking_uri (str | None) – Optional URI for the MLflow tracking server.

  • tags (dict[str, str] | None) – Initial tags to attach to the run.

log_dataset(df, name, source='unknown', context='training')#

Logs dataset lineage for reproducibility.

Parameters:

  • df (DataFrame) – The pandas DataFrame to track.

  • name (str) – Human-readable name of the dataset.

  • source (str) – Origin of data (e.g., S3 URI, SQL Query).

  • context (str) – Context of the data usage (e.g., ‘training’, ‘test’).

Return type:

None

log_evaluation(metrics_df, results_df)#

Logs summary metrics and interactive evaluation tables.

Parameters:

  • metrics_df (DataFrame) – DataFrame containing computed metrics per split/sample.

  • results_df (DataFrame) – DataFrame containing actual values and predictions.

Return type:

None

log_forecaster_metadata(forecaster)#

Harvests all available configurations and hyperparameters.

Consolidates Pydantic configs and top-level primitive attributes into MLflow parameters.

Return type:

None

log_model(forecaster, sample_input)#

Logs the entire forecaster as a PyFunc model with environment dependencies.

Parameters:

  • forecaster (TwigaForecaster) – The fitted TwigaForecaster.

  • sample_input (DataFrame) – Sample data used to infer the model schema/signature.

Return type:

None

property run_id: str#

The active MLflow run ID, or 'unknown' if no run is active.

Logging helpers#

Parameters from Pydantic configs#

log_config uses model_dump() and flattens nested dicts with dot-separated keys so that all Pydantic config fields appear as searchable MLflow parameters:

tracker.log_config(forecaster.data_pipeline, prefix="data")
# → logs  data.target_feature, data.forecast_horizon, data.lookback_window_size, …

Metrics from a DataFrame#

log_metrics reads the standard Twiga evaluation DataFrame (columns include mae, rmse, corr, Model). Each model gets its own metric namespace:

tracker.log_metrics(metrics_df)
# → logs  lgbm/mae, lgbm/rmse, catboost/mae, catboost/rmse, …

Full forecaster log#

log_forecaster is a convenience wrapper that logs:

What

MLflow path

data_pipeline fields

params

project_name, model_type

params

Checkpoint directory

artefacts/checkpoints/

 
tracker.log_forecaster(forecaster)

Using a remote tracking server#

Pass tracking_uri to point at a remote MLflow instance:

with TwigaTracker(
    experiment="solar-forecast",
    tracking_uri="http://mlflow.internal:5000",
) as tracker:
    ...

Without tracking_uri, MLflow defaults to a local mlruns/ directory in the working directory.

Prefect integration#

Inside a Prefect flow, the log_to_mlflow task wraps TwigaTracker so you do not need to call it directly:

from twiga.pipeline import training_flow

training_flow(
    forecaster=forecaster,
    data_path="data/load.parquet",
    experiment="load-forecast",
    run_name="lgbm-v1",
)

See Pipeline for the complete flow reference.