Track and compare models using MLflow Logged Models (Beta)

Beta

This feature is in Beta.

MLflow Logged Models help you track a model's progress throughout its lifecycle. When you train a model, use mlflow.<model-flavor>.log_model() to create a LoggedModel that ties together all of its critical information using a unique ID.

The LoggedModel object persists throughout the model's lifecycle, across different environments and runs, and contains links to artifacts such as metadata, metrics, parameters, and the code used to generate the model.

Logged Model tracking lets you compare models against each other, find the most performant model, and track down information during debugging.

Improved tracking for gen AI and deep learning models

Generative AI and deep learning workflows especially benefit from the granular tracking that Logged Models provides.

Gen AI - unified evaluation and trace data:

• Gen AI models generate additional metrics during evaluation and deployment, such as reviewer feedback data and traces.
• The LoggedModel entity allows you to query all the information generated by a model using a single interface.

Deep learning - efficient checkpoint management:

• Deep learning training creates multiple checkpoints, which are snapshots of the model's state at a particular point during training.
• MLflow creates a separate LoggedModel for each checkpoint, containing the model's metrics and performance data. This allows you to compare and evaluate checkpoints to identify the best-performing models efficiently.

Create a Logged Model

To create a Logged Model, use the same log_model() API as existing MLflow workloads. The following code snippets show how to create a Logged Model for gen AI, deep learning, and traditional ML workflows.

For complete, runnable notebook examples, see Example notebooks.

Gen AI

The following code snippet shows how to log a LangChain agent. Use the log_model() method for your flavor of agent.

Python

# Log the chain with MLflow, specifying its parameters
# As a new feature, the LoggedModel entity is linked to its name and params
model_info = mlflow.langchain.log_model(
  lc_model=chain,
  artifact_path="basic_chain",
  params={
    "temperature": 0.1,
    "max_tokens": 2000,
    "prompt_template": str(prompt)
  },
  model_type="agent",
  input_example={"messages": "What is MLflow?"},
)

# Inspect the LoggedModel and its properties
logged_model = mlflow.get_logged_model(model_info.model_id)
print(logged_model.model_id, logged_model.params)

Start an evaluation job and link the metrics to a Logged Model by providing the unique model_id for the LoggedModel:

Python

# Start a run to represent the evaluation job
with mlflow.start_run() as evaluation_run:
  eval_dataset: mlflow.entities.Dataset = mlflow.data.from_pandas(
    df=eval_df,
    name="eval_dataset",
  )
  # Run the agent evaluation
  result = mlflow.evaluate(
    model=f"models:/{logged_model.model_id}",
    data=eval_dataset,
    model_type="databricks-agent"
  )
  # Log evaluation metrics and associate with agent
  mlflow.log_metrics(
    metrics=result.metrics,
    dataset=eval_dataset,
    # Specify the ID of the agent logged above
    model_id=logged_model.model_id
  )

Deep learning

The following code snippet shows how to create Logged Models during deep learning training. Use the log_model() method for your flavor of MLflow model.

Python

# Start a run to represent the training job
with mlflow.start_run():
  # Load the training dataset with MLflow. We will link training metrics to this dataset.
  train_dataset: Dataset = mlflow.data.from_pandas(train_df, name="train")
  X_train, y_train = prepare_data(train_dataset.df)

  criterion = nn.CrossEntropyLoss()
  optimizer = torch.optim.Adam(scripted_model.parameters(), lr=0.01)

  for epoch in range(101):
    X_train, y_train = X_train.to(device), y_train.to(device)
    out = scripted_model(X_train)
    loss = criterion(out, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Log a checkpoint with metrics every 10 epochs
    if epoch % 10 == 0:
      # Each newly created LoggedModel checkpoint is linked with its
      # name, params, and step
      model_info = mlflow.pytorch.log_model(
        pytorch_model=scripted_model,
        name=f"torch-iris-{epoch}",
        params={
          "n_layers": 3,
          "activation": "ReLU",
          "criterion": "CrossEntropyLoss",
          "optimizer": "Adam"
        },
        step=epoch,
        input_example=X_train.numpy(),
      )
      # Log metric on training dataset at step and link to LoggedModel
      mlflow.log_metric(
        key="accuracy",
        value=compute_accuracy(scripted_model, X_train, y_train),
        step=epoch,
        model_id=model_info.model_id,
        dataset=train_dataset
      )

Traditional ML

The following code snippet shows how to log a sklearn model and link metrics to the Logged Model. Use the log_model() method for your flavor of MLflow model.

Python

## Log the model
model_info = mlflow.sklearn.log_model(
  sk_model=lr,
  name="elasticnet",
  params={
    "alpha": 0.5,
    "l1_ratio": 0.5,
  },
  input_example = train_x
)

# Inspect the LoggedModel and its properties
logged_model = mlflow.get_logged_model(model_info.model_id)
print(logged_model.model_id, logged_model.params)

# Evaluate the model on the training dataset and log metrics
# These metrics are now linked to the LoggedModel entity
predictions = lr.predict(train_x)
(rmse, mae, r2) = compute_metrics(train_y, predictions)
mlflow.log_metrics(
  metrics={
    "rmse": rmse,
    "r2": r2,
    "mae": mae,
  },
  model_id=logged_model.model_id,
  dataset=train_dataset
)

Example notebooks

For example notebooks that illustrate the use of LoggedModels, see the following pages:

• MLflow 3.0 generative AI workflow (Beta)
• MLflow 3.0 deep learning workflow (Beta)
• MLflow 3.0 traditional ML workflow (Beta)

View models and track progress

You can view your Logged Models in the workspace UI:

1. Go to the Experiments tab in your workspace.
2. Select an experiment. Then, select the Models tab.

This page contains all the Logged Models associated with the experiment, along with their metrics, parameters, and artifacts.

You can search and filter on Logged Models and generate charts to track metrics across runs.

Search Logged Models programmatically

In addition to the UI, search for Logged Models using the MLflow API:

Python

## Get a Logged Model using a model_id
mlflow.get_logged_model(model_id = <my-model-id>)

## Get all Logged Models that you have access to
mlflow.search_logged_models()

For more information and additional search parameters see MLflow 3.0 API documentation.