Tracking model versions using MLflow Model Registry

While the MLflow Tracking server lets you track all the attributes of your ML experiments, MLflow Model Registry provides a central model repository that lets you track all the aspects of your model life cycle. MLflow Model Registry consists of a user interface and APIs to track the model's version, lineage, stage transitions, annotations, and any developer comments. MLflow Model Registry also contains webhooks for CI/CD integrations and a model server for online model serving.

MLflow Model Registry provides us with a way to track and organize the many ML models that are produced and used by businesses during development, testing, and production. Model Registry provides a secure way to share models by leveraging access control lists and provides a way to integrate with model governance and approval workflows. Model Registry also allows us to monitor ML deployments and their performance via its API.

Tip

Model Registry's access controls and the ability to set permissions on registered models are only available in the full version of Databricks. They are not available in the Community Edition of Databricks or the open source version of MLflow.

Once a model has been logged to the Model Registry, you can add, modify, update, transition, or delete the model through the UI or the API, as shown in the following code sample:

import mlflow
from mlflow.tracking.client import MlflowClient
client = MlflowClient()
model_name = "linear-regression-model"
artifact_path = "best_model"
model_uri = "runs:/{run_id}/{artifact_path}".format (run_id=run_id, artifact_path=artifact_path)
registered_model = mlflow.register_model(model_uri=model_uri, name=model_name, )
client.update_model_version(
  name=registered_model.name,
  version=registered_model.version,
  description="This predicts the age of a customer using transaction history."
)
client.transition_model_version_stage(
  name=registered_model.name,
  version=registered_model.version,
  stage='Staging',
)
model_version = client.get_model_version(
  name=registered_model.name,
  version=registered_model.version,
)
model_uri = "models:/{model_name}/staging".format(model_name=model_name)
spark_model = mlflow.spark.load_model(model_uri)

In the previous code snippet, we did the following:

1. First, we imported the relevant MLflow libraries and MlflowClient, an MLflow interface for accessing the MLflow Tracking server and Model Registry artifacts via Python. The client interface is invoked using the MlflowClient() method.
2. Then, we constructed model_uri by providing the model, name, the model artifact location, and the run_id property from a tracked experiment. This information can be accessed via the MLflow Tracking server, either via the API or the UI.
3. Since we had reconstructed the model URI from the tracking server, we registered the model to the Model Registry using the register_model() function. If the model doesn't already exist in the Model Registry, a new model with a version of 1 is registered instead.
4. Once the model was registered, we had the opportunity to add or update the model description using the update_model_version() method.
5. During the life cycle of a model, it often evolves and needs to transition its stage, say from staging to production or the archival stage. This can be accomplished using the transition_model_version_stage() method.
6. A registered model's version and stage can be listed using the get_model_version() method.
7. Finally, we loaded a specific version/stage of the model from the Model Registry using the load_model() method after constructing the model URI using the models keyword, the model's name, and its version or stage name.

Model Registry also provides other handy functions for listing and searching through various versions of an individual registered model, as well as archiving and deleting registered models. References for those methods can be found at https://www.mlflow.org/docs/latest/model-registry.html#model-registry-workflows.

Important Note

Databricks Community Edition, which we will be using throughout this book, does not include Model Registry, so you will have to use the full version of Databricks to execute the preceding code sample. Alternatively, you can deploy your own MLflow Remote Tracking Server, along with a database-backed backend store outside of the Databricks environment, as described here: https://mlflow.org/docs/latest/tracking.html#mlflow-tracking-servers.

Now that you have learned how to store, version, and retrieve models using MLflow Model Registry, the next step of the ML life cycle is putting the trained, evaluated, and tuned models to practical use by leveraging them in business applications to draw inferences. We will discuss this in the next section.