Best ML Model Registry Tools

A model registry is a central repository that is used to version control Machine Learning (ML) models. It simply tracks the models while they move between training, production, monitoring, and deployment. It stores all the predominant information such as:

• metadata,
• lineage,
• model versions,
• annotations,
• and training jobs.

As the model registry is shared by multiple team members working on the same machine learning project, model governance is a major advantage that these teams have. This governance data tells them:

• which dataset was used for training,
• who trained and published a model,
• what’s the predictive performance of the model,
• and finally, when the model was deployed to production.

Read also

Best Tools for ML Model Governance, Provenance, and Lineage

Usually, while working in a team, different team members tend to try out different things, and only a few of them are finalized and pushed to the version control tool they use. The model registry helps them solve this issue as each team member can try their own versions of models, and they will have a record of all the things they have experimented with throughout the project journey.

This article will discuss the model registry tools and evaluation criteria for such tools. You will also see a comparison of different model registry and model management tools, such as:

• MLflow,
• Verta.ai,
• Comet,
• and neptune.ai,

So let’s get started!

Evaluation criteria for choosing model registry tools 

The model registry is an important part of MLOps platforms/tools. There are plenty of tools available in the market that can fulfill your ML workflow needs. Here is an illustration that classifies these tools on the basis of their specialization.

The products on the bottom right are focused on deployment and monitoring; those on the bottom-left focus on training and tracking. Those at the very top aim to cover every aspect of the ML lifecycle, while those in the middle-top do most or all of the spectrum with leaning one way or another.

To visualize it even more precisely, let’s have a look at another image:

From the above image, it can be inferred that tools like Kubeflow and other cloud providers are the most balanced and cover every stage of an ML pipeline development equally. Specialized tools like Neptune and Polyaxon are closest to their axis, i.e., majorly focused on model training. 

NOTE: The aforementioned evaluation criteria for these tools are subjective to the features these tools had at that point in time (November 2021). Many of these tools have moved much beyond their area of specialization in the past year, so take this discussion with a pinch of salt.

However, there are some evergreen factors that are integral to determining a registry tool’s effectiveness. From my own experience, some of them are:

Ease of automation

One of the requirements of a model registry tool is how easily the development team can make use of that tool.

• Some tools require you to code all the things needed to store the model versions,
• While some tools require very less coding, and you just need to drag and drop different components to use them.
• There are also some tools fully based on the concept of AutoML and do not require you to write any code for storing your model versions. 

Auto-ML tools have less flexibility for customizations while Low-Code tools provide both custom and automation options finally, Code-First tools only provide a writing code facility. You can choose a tool based on your requirement.

Updated model overview and model stages tracking 

The entire purpose of a model registry tool is to provide an easy overview of all the versions of models that the development team has tried. While selecting the tool, you must remember that the tool must provide the model overview of each version at every stage. Tracking models extend beyond development; it is done for maintenance and enhancement in staging and production as well. The machine learning model lifetime including:

• training,
• staging,
• and production,

must be tracked by the model registration tool.

Competence in managing the model dependencies

The model registry tool must have compatibility with all the dependencies the ML model needs. You should check the dependencies competence for the Machine Learning libraries, Python version, and data. If you are working on some use case that requires a special ML library and the registry tool does not support it, that tool would not make much sense for you.

Providing the flexibility of team collaboration

You may evaluate whether you and your team can collaborate on the registered model or not. If the model registry enables you to work with your team on the same ML model, then you can choose that tool.

Thus, you can follow the evaluation criteria to select the best model registry tool according to your requirements.

Comparison of model registry tools

Every model registry tool has different features and performs various unique operations. Here’s how they compare:

Model registry tools

Here are a number of model registry tools that are used across the industry:

MLflow

An open-source platform that you can use for managing the ML model lifecycle. MLFlow enables you to track the MLOps life cycle with the help of its APIs. It provides model versioning, model lineage, annotations, and transitions from development to deployment functionalities.

Some features of MLflow model registry are as follows:

• Model lineage tracking, showing which experiment and run produced a given model version
• Predefined model stages as Archived, Staging, and Production but allocates one model stage at a time for different model versions.
• Annotations and versioning, allowing you to document and manage top-level models and individual versions using Markdown
• Webhooks, triggering actions based on registry events.
• Email notifications, to stay informed about model lifecycle changes.

MLflow can be self-hosted or used as part of a managed service. While Databricks offers a full-featured hosted version, Amazon SageMaker and Azure Machine Learning also support the MLflow client, letting you track and register models within their ecosystems. However, in these cloud integrations, model data is logged to proprietary backends, and not all MLflow features are supported. These integrations provide convenience for teams operating within AWS or Azure, while still benefiting from MLflow’s open interface.

Learn more

Check detailed comparison between neptune.ai and MLflow.

Comet

Developers can use the Comet platform to manage machine learning experiments. This system allows you to version, register, and deploy the model using its Python SDK Experiment.  

Comet keeps track of model versions and the experiment history of the model. You can check the detailed information of all model versions. Besides, you can maintain ML workflow more efficiently using model reproduction and optimization.  

The feature-rich Comet has various functionalities for running and tracking ML model experiments, including:

• Comet allows you to easily check the history of evaluation/testing runs.
• You can easily compare different experiments using the Comet model registry. 
• It allows you to access the code, dependencies, hyperparameters, and metrics within a single UI. 
• It has in-built reporting and visualization features to communicate with team members and stakeholders.
• It lets you configure webhooks and integrate the Comet model registry with your CI/CD pipeline.

May be useful

Check detailed comparison between neptune.ai and Comet.

Verta.ai

You can use the Verta AI tool for the management and operations of the model in one unified space. It provides an interactive UI where you can register the ML models and publish the metadata, artefacts, and documents. Then, to manage the end-to-end experiment, you may connect the model to the experiment tracker. Version control solutions for ML projects are also offered by Verta AI.

Additionally, it enables you to keep track of changes made to data, code, environments, and model configuration. With the audit log’s accessibility, you may also examine the model’s dependability and compatibility at any time. You can also create a unique approval sequence that is appropriate for your project and incorporate it with the selected ticketing system.

Some of the main features of Verta AI’s model registry are:

• It enables end-to-end information tracking such as Model ID, description, tags, documentation, model versions, release stage, artifacts, model metadata, and more, which helps in selecting the best model. 
• It works on container tools like Kubernetes and Docker and is integrable with GitOps and Jenkins, which helps in automatically tracking model versions.
• It provides access to detailed audit logs for compliance. 
• It has an environment like Git that makes it intuitive.
• You can set up granular access control for editors, reviewers, and collaborators.

neptune.ai

Neptune is primarily an experiment tracker, but it provides model registry functionality to a great extent. 

Neptune allows you to log, visualize, compare, and query all metadata related to ML experiments and models. It only takes a few lines of code to integrate Neptune with your code. The API is flexible, and the UI is user-friendly but also prepared for the high volume of logged metadata.

 Some of the features of Neptune:

• It lets you track models and model versions, along with the associated metadata. You can version model code, images, datasets, Git info, and notebooks.
• It allows you to filter and sort the versioned data easily.
• It lets you manage model stages using tags.
• You can query and download any stored model files and metadata.
• And it helps your team to collaborate on experiments by providing persistent links to the UI or building reports tailored to specific stakeholders or project.
• It supports different connection modes such as asynchronous (default), synchronous, offline, read-only, and debug modes for the versioned metadata tracking.   

Summary

After reading this article, I hope you now know what model registry tools are and the different criteria that one must look for while selecting a model registry tool. To offer a practical perspective, we also discussed some of the popular model registry tools and compared them with each other in several aspects. Now, let’s wrap the article with a few key takeaways:

• Model registry performs model versioning and publishes them into production.
• Before selecting a model registry tool, you must evaluate each model according to your requirement.
• Model registry evaluation criteria can range from the capability to monitor and manage the different ML model stages and versions to its ease of use and pricing.
• You may refer to the highlighted features of different model registry tools to get a better idea of that tool’s compatibility with your use case.

With these points in mind, I hope your model registry tool search will be much easier.