Introduction
MLFlow is an open-source platform that allows you, among other things, to track experiments to compare different hyperparameters and results, and to serve ML models on different platforms. In this tutorial, we provide basic scripts that you can use to track experiments made with TimeGPT, or to serve TimeGPT through MLFLow. Each script is customizable to your own needs. The goal is to provide an easy-to-use template that you can extend.Experiment Tracking with TimeGPT and MLFlow
The following scripts provide functions for logging experiment results when testing different parameter combinations in TimeGPT. Of course, to use these scripts make sure to have MLFLow installed and any other required dependencies. To install MLFlow, runpip install mlflow. Here, we use the following packages:
import os
import mlflow
import pandas as pd
from datetime import datetime
from dotenv import load_dotenv # Used when working with the API directly or a Docker deployment
from api.serverless import make_client # Used when working with the Python wheel
NixtlaClient and pass it to the functions. You can see how to setup your client with an API key here.
Logging experiments with forecasting
The following script can be used to log results when using theforecast method.
def log_timegpt_forecast(
client: NixtlaClient,
df: pd.DataFrame,
h: int = 12,
freq: str = "MS",
level: list = None,
model: str = "timegpt-2-mini",
experiment_name: str = "basic_forecast",
time_col: str = "ds",
target_col: str = "y",
id_col: str = "unique_id",
**kwargs
):
"""
Perform TimeGPT forecast and log to MLFlow.
Parameters:
-----------
client : NixtlaClient
Initialized Nixtla client
df : pd.DataFrame
Input dataframe with time series data
h : int
Forecast horizon
freq : str
Frequency of the time series
level : list
Confidence levels for prediction intervals
model : str
TimeGPT model variant to use
experiment_name : str
Name for this MLFlow run
time_col : str
Name of the time column in df
target_col : str
Name of the target column in df
id_col : str
Name of the series identifier column in df
**kwargs : dict
Additional parameters to pass to forecast()
"""
with mlflow.start_run(run_name=experiment_name):
# Log parameters
mlflow.log_param("horizon", h)
mlflow.log_param("frequency", freq)
mlflow.log_param("model", model)
mlflow.log_param("n_series", df[id_col].nunique() if id_col in df.columns else 1)
mlflow.log_param("n_observations", len(df))
if level:
mlflow.log_param("prediction_intervals", level)
# Log any additional parameters
for key, value in kwargs.items():
mlflow.log_param(key, value)
# Log dataset info
mlflow.log_param("start_date", df[time_col].min())
mlflow.log_param("end_date", df[time_col].max())
# Forecast
forecast_df = client.forecast(
df=df,
h=h,
freq=freq,
level=level,
model=model,
id_col=id_col,
time_col=time_col,
target_col=target_col,
**kwargs
)
# Log metrics
mlflow.log_metric("forecast_points", len(forecast_df))
# Log forecast as artifact
forecast_path = "forecast_output.csv"
forecast_df.to_csv(forecast_path, index=False)
mlflow.log_artifact(forecast_path)
# Log input data as artifact
input_path = "input_data.csv"
df.to_csv(input_path, index=False)
mlflow.log_artifact(input_path)
# Add tags
mlflow.set_tag("task_type", "forecasting")
mlflow.set_tag("timestamp", datetime.now().isoformat())
# Clean up temporary files
os.remove(forecast_path)
os.remove(input_path)
return forecast_df
Logging experiments with cross-validation
The following is a simple script to log metrics when performing experiments with thecross_validation method. This script requires utilsforecast to compute performance metrics, so make sure pip install utilsforecast.
The script allows you to log overall performance metrics across multiple windows, and also per-window metrics.
def log_timegpt_cross_validation(
client: NixtlaClient,
df: pd.DataFrame,
h: int = 12,
n_windows: int = 1,
step_size: int = None,
freq: str = "MS",
level: list = None,
model: str = "timegpt-2-mini",
experiment_name: str = "cross_validation",
time_col: str = "ds",
target_col: str = "y",
id_col: str = "unique_id",
**kwargs
):
"""
Perform TimeGPT cross-validation and log to MLFlow.
Parameters:
-----------
client : NixtlaClient
Initialized Nixtla client
df : pd.DataFrame
Input dataframe with time series data
h : int
Forecast horizon
n_windows : int
Number of cross-validation windows
step_size : int
Step size between windows (default: h)
freq : str
Frequency of the time series
level : list
Confidence levels for prediction intervals
model : str
TimeGPT model variant to use
experiment_name : str
Name for this MLFlow run
time_col : str
Name of the time column in df
target_col : str
Name of the target column in df
id_col : str
Name of the series identifier column in df
**kwargs : dict
Additional parameters to pass to cross_validation()
"""
with mlflow.start_run(run_name=experiment_name):
# Log parameters
mlflow.log_param("horizon", h)
mlflow.log_param("n_windows", n_windows)
mlflow.log_param("step_size", step_size or h)
mlflow.log_param("frequency", freq)
mlflow.log_param("model", model)
mlflow.log_param("n_series", df[id_col].nunique() if id_col in df.columns else 1)
if level:
mlflow.log_param("prediction_intervals", level)
for key, value in kwargs.items():
mlflow.log_param(key, value)
# Perform cross-validation
cv_df = client.cross_validation(
df=df,
h=h,
n_windows=n_windows,
step_size=step_size,
freq=freq,
level=level,
model=model,
time_col=time_col,
target_col=target_col,
id_col=id_col,
**kwargs
)
# Calculate and log metrics
from utilsforecast.losses import mae, mse, rmse
# MAE
mae_value = mae(
df=cv_df.drop(columns=['cutoff']),
models=['TimeGPT'],
target_col=target_col,
id_col=id_col,
)['TimeGPT'].values[0]
# MSE
mse_value = mse(
df=cv_df.drop(columns=['cutoff']),
models=['TimeGPT'],
target_col=target_col,
id_col=id_col,
)['TimeGPT'].values[0]
# RMSE
rmse_value = rmse(
df=cv_df.drop(columns=['cutoff']),
models=['TimeGPT'],
target_col=target_col,
id_col=id_col,
)['TimeGPT'].values[0]
mlflow.log_metric("mae", mae_value)
mlflow.log_metric("mse", mse_value)
mlflow.log_metric("rmse", rmse_value)
mlflow.log_metric("total_cv_predictions", len(cv_df))
# Log per-window metrics (MAE as an example)
cutoffs = cv_df['cutoff'].unique()
for i, cutoff in enumerate(cutoffs):
window_df = cv_df[cv_df['cutoff'] == cutoff]
window_mae = mae(
df=window_df.drop(columns=['cutoff']),
models=['TimeGPT'],
target_col=target_col,
id_col=id_col,
)['TimeGPT'].values[0]
mlflow.log_metric(f"mae_window_{i+1}", window_mae)
# Log artifacts
cv_path = "cross_validation_results.csv"
cv_df.to_csv(cv_path, index=False)
mlflow.log_artifact(cv_path)
# Log summary statistics
summary = {
'metric': ['MAE', 'MSE', 'RMSE'],
'value': [mae_value, mse_value, rmse_value]
}
summary_df = pd.DataFrame(summary)
summary_path = "metrics_summary.csv"
summary_df.to_csv(summary_path, index=False)
mlflow.log_artifact(summary_path)
mlflow.set_tag("task_type", "cross_validation")
mlflow.set_tag("timestamp", datetime.now().isoformat())
# Clean up
os.remove(cv_path)
os.remove(summary_path)
return cv_df
Logging experiments with anomaly detection
This script showcases how you can track experiments done with thedetect_anomalies_online method.
def log_timegpt_online_anomaly_detection(
client: NixtlaClient,
df: pd.DataFrame,
h: int,
detection_size: int,
threshold_method: str = "univariate",
freq: str = "D",
level: int | float = 99,
model: str = "timegpt-2-mini",
experiment_name: str = "anomaly_detection",
time_col: str = "ds",
target_col: str = "y",
id_col: str = "unique_id",
**kwargs
):
"""
Perform TimeGPT anomaly detection and log results to MLFlow.
Parameters:
-----------
client : NixtlaClient
Initialized Nixtla client
df : pd.DataFrame
Input dataframe with time series data
freq : str
Frequency of the time series
level : int
Confidence level for anomaly detection threshold
model : str
TimeGPT model variant to use
experiment_name : str
Name for this MLFlow run
time_col : str
Name of the time column in df
target_col : str
Name of the target column in df
id_col : str
Name of the series identifier column in df
**kwargs : dict
Additional parameters to pass to detect_anomalies()
"""
with mlflow.start_run(run_name=experiment_name):
# Log parameters
mlflow.log_param("horizon", h)
mlflow.log_param("detection_size", detection_size)
mlflow.log_param("threshold_method", threshold_method)
mlflow.log_param("frequency", freq)
mlflow.log_param("detection_level", level)
mlflow.log_param("model", model)
mlflow.log_param("n_observations", len(df))
for key, value in kwargs.items():
mlflow.log_param(key, value)
# Detect anomalies
anomalies_df = client.detect_anomalies_online(
df=df,
h=h,
detection_size=detection_size,
threshold_method=threshold_method,
freq=freq,
level=level,
model=model,
time_col=time_col,
target_col=target_col,
id_col=id_col,
**kwargs
)
# Calculate metrics
n_anomalies = anomalies_df['anomaly'].sum()
mlflow.log_metric("n_anomalies", n_anomalies)
# Log results
anomaly_path = "anomaly_detection_results.csv"
anomalies_df.to_csv(anomaly_path, index=False)
mlflow.log_artifact(anomaly_path)
# Log only the detected anomalies
if n_anomalies > 0:
detected_anomalies = anomalies_df[anomalies_df['anomaly'] == True]
detected_path = "detected_anomalies_only.csv"
detected_anomalies.to_csv(detected_path, index=False)
mlflow.log_artifact(detected_path)
os.remove(detected_path)
mlflow.set_tag("task_type", "anomaly_detection")
mlflow.set_tag("timestamp", datetime.now().isoformat())
# Clean up
os.remove(anomaly_path)
return anomalies_df
Sample usage
With the above functions, you can now run experiments and log it to MLFLow. First, you must instantiate your client using either:from nixtla import NixtlaClient
nixtla_client = NixtlaClient(api_key='your_api_key_here')
from api.serverless import make_client
client = make_client()
df = pd.read_csv(
'https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/air_passengers.csv'
)
df.columns = ['ds', 'y']
df["unique_id"] = 0
df = df[["unique_id", "ds", "y"]]
mlflow.set_tracking_uri("mlruns")
experiment = mlflow.set_experiment("timegpt_experiments")
forecast_df = log_timegpt_forecast(
client=client,
df=df,
h=h,
freq=freq,
level=[80, 90],
experiment_name="forecast_example",
time_col=time_col,
target_col=target_col,
id_col=id_col,
model="timegpt-2"
)
cv_df = log_timegpt_cross_validation(
client=client,
df=df,
h=h,
n_windows=1,
freq=freq,
level=[90],
experiment_name="cv_example",
time_col=time_col,
target_col=target_col,
id_col=id_col,
)
anomaly_df = log_timegpt_online_anomaly_detection(
client=client,
df=df,
h=h,
detection_size=h,
threshold_method="univariate",
freq=freq,
level=60,
model="timegpt-2-mini",
experiment_name="anomaly_detection_example",
time_col=time_col,
target_col=target_col,
id_col=id_col,
)
Serving TimeGPT with MLFlow
You can also use MLFlow for model serving. This can be useful if you need a unified interface and serve the model on many endpoints. Although it is not required to use MLFlow to use TimeGPT, it can be a practice your organization enforces. The following script shows how you can wrap the TimeGPT model in anmlflow.pyfunc.PythonModel to save the model and call it.
class MLFLowTimeGPTModel(mlflow.pyfunc.PythonModel):
"""
Unified MLflow pyfunc wrapper for TimeGPT
Can perform forecasting, cross-validation, and anomaly detection
based on the 'operation' parameter in the input.
"""
def __init__(self,
client: NixtlaClient,
model: str = "timegpt-2-mini",
default_h: int = 96,
default_freq: str = "15min",
default_level: Optional[list] = None,
default_n_windows: int = 1,
default_anomaly_level: int = 99,
):
"""
Initialize the unified TimeGPT model wrapper.
Parameters:
-----------
model : str
TimeGPT model variant (timegpt-2-mini, timegpt-2)
default_h : int
Default forecast horizon
default_freq : str
Default frequency
default_level : list
Default confidence levels for prediction intervals
default_n_windows : int
Default number of cross-validation windows
default_anomaly_level : int
Default confidence level for anomaly detection
"""
self.model = model
self.default_h = default_h
self.default_freq = default_freq
self.default_level = default_level or [80, 90]
self.default_n_windows = default_n_windows
self.default_anomaly_level = default_anomaly_level
self.client = client
def load_context(self, context):
"""
Load the model context. Called once when the model is loaded.
Parameters:
-----------
context : mlflow.pyfunc.PythonModelContext
Context containing artifacts and other model metadata
"""
# Load configuration from artifacts if present
if context.artifacts and "config" in context.artifacts:
with open(context.artifacts["config"], 'r') as f:
config = json.load(f)
self.model = config.get("model", self.model)
self.default_h = config.get("h", self.default_h)
self.default_freq = config.get("freq", self.default_freq)
self.default_level = config.get("level", self.default_level)
self.default_n_windows = config.get("n_windows", self.default_n_windows)
self.default_anomaly_level = config.get("anomaly_level", self.default_anomaly_level)
def predict(self, context, model_input):
"""Perform operation"""
# Parse input
if isinstance(model_input, dict):
df = model_input.get('data')
operation = model_input.get('operation', 'forecast')
time_col = model_input.get('time_col', 'ds')
target_col = model_input.get('target_col', 'y')
id_col = model_input.get('id_col', 'unique_id')
model_input.pop('time_col', None)
model_input.pop('target_col', None)
model_input.pop('id_col', None)
else:
# If just DataFrame is passed, default to forecast
df = model_input
operation = 'forecast'
id_col = 'unique_id'
time_col = 'ds'
target_col = 'y'
# Validate input
if df is None or not isinstance(df, pd.DataFrame):
raise ValueError("Input must contain a pandas DataFrame under 'data' key")
# Route to appropriate operation
if operation == 'forecast':
return self._forecast(df, model_input, id_col, time_col, target_col)
elif operation == 'cross_validation':
return self._cross_validation(df, model_input, id_col, time_col, target_col)
elif operation == 'anomaly_detection':
return self._anomaly_detection(df, model_input, id_col, time_col, target_col)
else:
raise ValueError(f"Unknown operation: {operation}. Must be 'forecast', 'cross_validation', or 'anomaly_detection'")
def _forecast(self, df, params, id_col, time_col, target_col):
"""Perform forecasting operation."""
h = params.get('h', self.default_h) if isinstance(params, dict) else self.default_h
freq = params.get('freq', self.default_freq) if isinstance(params, dict) else self.default_freq
level = params.get('level', self.default_level) if isinstance(params, dict) else self.default_level
# Extract additional parameters
additional_params = {}
if isinstance(params, dict):
exclude_keys = {'data', 'operation', 'h', 'freq', 'level', 'time_col', 'target_col'}
additional_params = {k: v for k, v in params.items() if k not in exclude_keys}
result = self.client.forecast(
df=df,
h=h,
freq=freq,
level=level,
id_col=id_col,
time_col=time_col,
target_col=target_col,
model=self.model,
**additional_params
)
return result
def _cross_validation(self, df, params, id_col, time_col, target_col):
"""Perform cross-validation operation."""
h = params.get('h', self.default_h) if isinstance(params, dict) else self.default_h
n_windows = params.get('n_windows', self.default_n_windows) if isinstance(params, dict) else self.default_n_windows
freq = params.get('freq', self.default_freq) if isinstance(params, dict) else self.default_freq
level = params.get('level', self.default_level) if isinstance(params, dict) else self.default_level
step_size = params.get('step_size') if isinstance(params, dict) else None
# Extract additional parameters
additional_params = {}
if isinstance(params, dict):
exclude_keys = {'data', 'operation', 'h', 'n_windows', 'freq', 'level', 'step_size', 'time_col', 'target_col'}
additional_params = {k: v for k, v in params.items() if k not in exclude_keys}
result = self.client.cross_validation(
df=df,
h=h,
n_windows=n_windows,
step_size=step_size,
freq=freq,
level=level,
id_col=id_col,
time_col=time_col,
target_col=target_col,
model=self.model,
**additional_params
)
return result
def _anomaly_detection(self, df, params, id_col, time_col, target_col):
"""Perform anomaly detection operation."""
h = params.get('h', self.default_h) if isinstance(params, dict) else self.default_h
detection_size = params.get("detection_size", self.default_h) if isinstance(params, dict) else self.default_h
threshold_method = params.get("threshold_method", "univariate") if isinstance(params, dict) else "univariate"
freq = params.get('freq', self.default_freq) if isinstance(params, dict) else self.default_freq
level = params.get('level', self.default_anomaly_level) if isinstance(params, dict) else self.default_anomaly_level
# Extract additional parameters
additional_params = {}
if isinstance(params, dict):
exclude_keys = {'data', 'operation', 'h', 'detection_size',
'threshold_method', 'freq', 'level', 'time_col', 'target_col'}
additional_params = {k: v for k, v in params.items() if k not in exclude_keys}
result = self.client.detect_anomalies_online(
df=df,
h=h,
detection_size=detection_size,
threshold_method=threshold_method,
freq=freq,
level=level,
time_col=time_col,
id_col=id_col,
target_col=target_col,
model=self.model,
**additional_params
)
return result
def save_unified_model(
client: NixtlaClient,
model_path: str,
model_variant: str = "timegpt-2-mini",
**default_params
):
"""
Save a unified TimeGPT model that can perform all operations.
Parameters:
-----------
client : NixtlaClient
Initialized Nixtla client
model_path : str
Path where the model will be saved
model_variant : str
TimeGPT model variant to use
**default_params : dict
Default parameters for operations
"""
python_model = MLFLowTimeGPTModel(
client=client,
model=model_variant,
**default_params
)
# Save the model
mlflow.pyfunc.save_model(
path=model_path,
python_model=python_model,
)
Sample usage
First, you must instantiate your client using either:from nixtla import NixtlaClient
nixtla_client = NixtlaClient(api_key='your_api_key_here')
from api.serverless import make_client
client = make_client()
df = pd.read_csv(
'https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/air_passengers.csv'
)
df.columns = ['ds', 'y']
df["unique_id"] = 0
df = df[["unique_id", "ds", "y"]]
save_unified_model(client=client, model_path="test_model", model_variant="timegpt-2-mini", default_h=h, default_freq=freq)
timegpt-2, then you must save another instance and specify that model variant.
Now, you can perform forecasting, cross-validation and anomaly detection with the saved model in MLFlow.
# Load model
model = mlflow.pyfunc.load_model("test_model")
# Forecast
forecast = model.predict({'data': df, 'operation': 'forecast', 'h': h, 'id_col': id_col, 'time_col': time_col, 'target_col': target_col, 'freq': freq})
# Cross-validation
cv = model.predict({'data': df, 'operation': 'cross_validation', 'h': h, 'n_windows': 1, 'id_col': id_col, 'time_col': time_col, 'target_col': target_col, 'freq': freq})
# Anomaly detection
anomalies = model.predict({'data': df, 'operation': 'anomaly_detection', 'detection_size': h, 'id_col': id_col, 'time_col': time_col, 'target_col': target_col, 'freq': freq})