How to Migrate From MLflow to Neptune

TL;DR

MLflow proved to have many limitations that neptune.ai can address, providing better security, more robust collaboration tools, and a user-friendly interface.

The migration is not as complex as you might think. neptune.ai developed solutions to ease this process.

Your MLflow run logs can easily be exported to the neptune.ai app using a dedicated plugin.

Use our MLflow vs neptune.ai API comparison table to migrate your training scripts faster.

As an MLflow user, it is straightforward to adapt to neptune.ai’s UI.

MLflow is a framework widely used for its experiment-tracking capabilities, but many organizations are searching for alternatives. It is often used as an early tracking solution that teams can easily adopt as it is well-known and open-source. However, they often end up needing more advanced features and better support than MLflow can provide.

MLflow lacks robust security, collaboration, scaling, and interface modularity features. neptune.ai is a strong alternative to MLflow as it addresses these limitations. It’s built to manage larger ML projects and offers a fully managed service to abstract away the need for infrastructure maintenance.

This hands-on guide will walk you through all the necessary steps to perform the migration. By the end of this article, you will have transferred your data, adapted your training scripts, and have all the information you need to accustom your team to the new UI.

Seven reasons for migrating from MLflow to neptune.ai

Before diving into the migration process, let’s see how MLflow and Neptune differ and why migrating from MLflow to Neptune is worth the effort.

1. Security and compliance: MLflow has limited permissions management and lacks built-in authentication features. Neptune offers advanced access controls and meets strict compliance standards, including SOC 2 certification, which is crucial for organizations working with sensitive data.

2. User access management: MLflow does not support the management of user access and resource ownership, while Neptune does.

3. Team collaboration: MLflow does not provide collaboration features such as reviewing projects, sharing data, or creating detailed reports. Neptune provides utilities that make this process easier with customizable workspaces and persistent shareable links.

4. User interface: Teams may find MLflow’s interface less adaptable for analyzing experiments. Neptune’s interface is richer, customizable, and more intuitive, making it easier to manage and review experiments.

5. Infrastructure maintenance: Self-hosting open-source MLflow involves configuring and managing the server supporting the application. On the other hand, Neptune is available as a fully managed tool, saving your team members precious time.

6. Scalability and performance: MLflow is not designed to handle increasingly large data volumes produced by your expanding training pipelines. Neptune can process large streams of logs running 1000s of experiments at once.

7. Support: While MLflow’s community is supportive, there are no guarantees that your question or issue will be addressed. Neptune offers dedicated user support, helping to solve issues quickly.

How Veo Eliminated Work Loss With neptune.ai

Computer-vision models are an integral part of Veo’s products. Their sports cameras produce 4K video data that is analyzed to provide insights for coaches and players.

Initially, the team started with MLflow as the experiment tracker but quickly found it unreliable, especially under heavy computational loads. This led to connectivity issues and crashes, disrupting their product development and sometimes causing them to lose a week’s worth of work.

Switching to Neptune not only provided a reliable platform for their ML experiments but also allowed them to track and analyze tens of metrics effortlessly.

See in the app Full screen preview

• 

  Read the full case study

• 

  Get in touch if you’d like to go through a custom demo with your team

Challenges when migrating between machine learning experiment trackers

Migrating from one experiment tracker to another might seem overwhelming at first. However, breaking it down into clear steps makes migrating from MLflow to Neptune a lot less daunting.

These are the essential tasks you will need to do:

1. Exporting training metadata: The first step is to export your existing training logs from MLflow to Neptune. Neptune created the neptune-mlflow plugin to handle this process for you.

2. Copy dashboards: To maintain your experiment-tracking workflow and data visibility, you will need to replicate your MLflow dashboards in the new tool. Neptune offers no-code dashboard features. You can add pre-built widgets with drag-and-drop and resize to organize your data visualization as you wish.

3. Adapting training code: You have to modify your existing MLflow training scripts to redirect run logs to the new tool. You can either redirect your MLflow logs to Neptune using the neptune-mlflow plugin or perform a full migration to Neptune’s native client.

4. Setting up user accounts: To manage collaboration on the new tool, you will need to create accounts and set up user permissions. Neptune’s documentation offers guidance on how to add collaborators to your workspace and manage user roles.

5. Onboarding: You will have to train your team to use the new tool effectively.
   

In the remainder of this article, we’ll work through these steps in detail.

How to migrate from MLflow to Neptune: step-by-step guide

Here are the steps we will perform together:

1. Prepare for migration: We will ensure we have access to the necessary resources to perform the migration.

2. Create a neptune.ai account and project: We will sign up for Neptune, access our workspace, and create a project to organize our experiments.

3. Set up a local Python environment: We are going to create a Python environment with the necessary packages to perform the migration.
   
4. Configure the Neptune client: We will retrieve an API key and pass it to the client along with the project name.

5. Export MLflow logs to neptune.ai: We will migrate your existing MLflow run logs and models to the Neptune experiment tracker.
   
6. Inspect the imported data: After the export has finished, we will verify that all data has been transferred and can be accessed in Neptune.

7. Adapt your training scripts: We will adapt your training scripts to send experiment metadata to Neptune by adding just a few lines of Python code.

8. Optional: Migrate from the MLflow client to Neptune’s native client: We are going to walk through the refactoring of an MLflow-instrumented training script to use Neptune’s native client.

Prepare for migration

Before you start, make sure that:

• You have a computer with Python and pip installed. (We developed this tutorial using Python 3.10.12.) You will need to be able to create new Python environments and download software packages from the internet.

• Your MLflow tracking server is running and you know its URI. (This is the address you pass to your MLflow client via mlflow.set_tracking_uri() or the MLFLOW_TRACKING_URI environment variable.)
  
• You can access the MLflow API from the computer you’re working on. (If you can log data to MLflow from your computer, this is fulfilled.)

Create a Neptune account and project

If you do not have a Neptune account yet, create one on the neptune.ai registration page.

Once your account is created, access your workspace by logging into the Neptune app and create a project as follows:

Click on All projects at the top left. Then click on Create project at the top right.

The Create new project form will appear as shown below:

Choose a name for your project such as mlflow-to-neptune, and a project key with upper case letters like MIGRATION. This key will be used to identify runs and other Neptune objects.

Finally, click the Create button at the bottom right. You now have a new, empty Neptune project to transfer your MLflow data to.

Set up a local Python environment

To perform the migration, we will need to install the Neptune and MLflow clients in a Python environment. We will build the environment using virtualenv to isolate it from your other Python projects. (If you’re already using a different Python environment manager, feel free to use it instead.)

First, install virtualenv using pip:

pip install virtualenv

If you store your MLflow data locally, navigate to the root directory of your MLflow repository:

cd your_mlflow_project_directory

In all other cases, create a new directory (e.g., mlflow_migration).

Create a new virtual environment:

virtualenv venv

Activate the environment you just created:

source venv/bin/activate

You now have a functional Python environment.

Let’s proceed by installing the necessary packages to perform the migration using pip:

pip install -U mlflow neptune neptune-mlflow

Configure the Neptune client

Next, we need to configure the Neptune client. It needs the name and an access token of the project to which it should send data.

In your Neptune workspace, access your Neptune API token as shown in the screenshot below:

In the bottom-left corner of your screen, click on your username and then on Get your API token.

The following pop-up will show up:

Click the copy to clipboard icon on the left to copy your API token.

Then, set the NEPTUNE_API_TOKEN environment variable in your terminal:

export NEPTUNE_API_TOKEN=

Now, you need to find the precise identifier of your Neptune project. On the main page of your workspace, click on your project. Then, click on the Edit project details option in the menu in the top right:

The Edit project details page will appear.

Copy your project identifier by clicking on the copy to clipboard icon:

Then, set the NEPTUNE_PROJECT variable to this identifier:

export NEPTUNE_PROJECT=

These two environment variables will be picked up by the Neptune client and allow it to access your newly created project.

Export MLflow logs to Neptune

Depending on whether you store your MLflow data locally or work with a remote tracking server, the export works slightly differently.

Local MLflow

In your terminal, navigate to the directory containing your MLflow logs and run:

neptune mlflow

If you don’t run the MLflow tracking server at its default address, you’ll have to pass it via the –mlflow-tracking-uri command-line option:

neptune mlflow --mlflow-tracking-uri https://127.0.0.1:8080

If you encounter a NeptuneSynchronizationAlreadyStopped exception while running neptune mlflow, don’t panic! You will simply need to run the synchronization manually:

neptune sync

Remote MLflow

If you store the MLflow data remotely, e.g., using a centralized MLflow tracking server, you need to pass the tracking URI to the export command:

neptune mlflow --mlflow-tracking-uri https://my-mlflow-tracking-server.example.com

The export plugin uses the MLflow tracking client internally, so you can pass any kind of URI supported by MLflow.

Inspect the imported data

Full screen preview

In your Neptune project, you should see all the experiments under the Run tab. On the left, you will find a list of the experiments migrated from MLflow. 

You might notice that the ID of the exported runs does not match the run names in MLflow.

Full screen preview

To match the run names, perform these steps as shown in the screenshot above:

1. Click on Runs table at the top of the web page.
2. Click on Add column in the top-right corner.
3. Write run_data/tags/mlflow.runName in the search bar.
4. Click on the item run_data/tags/mlflow.runName below the search bar.

You now have a column in your custom view that displays the original run names from MLflow:

Full screen preview

Adapt your training scripts

We can use the same neptune-mlflow plugin that we used for data migration in our training script to instruct the MLflow client to log all data to Neptune instead of MLflow.

This allows you to start using Neptune for all your experiment tracking quickly, without having to replace the MLflow-specific functions such as mlflow.log_metric() or mlflow.log_params().

Here’s an example of how to create a Neptune tracking URI and pass it to MLflow:

import os
import mlflow
from neptune_mlflow_plugin import create_neptune_tracking_uri

neptune_uri = create_neptune_tracking_uri(
    api_token=os.environ['NEPTUNE_API_TOKEN'],
    project=os.environ['NEPTUNE_PROJECT'],
    tags=["mlflow", "plugin", "script"],
)

mlflow.set_tracking_uri(uri=neptune_uri)

# The rest of your MLflow training code
```

Optional: Migrate from the MLflow client to Neptune’s native client

Adapting your MLflow scripts using the Neptune MLflow plugin is convenient, as it redirects your logs to Neptune with minimal effort. However, in the long run, you will most likely want to switch to the Neptune client to get the most out of Neptune and its versatile logging capabilities centered around the Run object.

MLflow uses multiple distinct functions to set up experiment tracking. In contrast, Neptune’s init_run() function returns a Run object. This object is central to Neptune’s tracking system and provides all the methods needed to log information.

Step-by-step example: replacing MLflow’s client with Neptune’s client

To guide you through the process of migrating to the Neptune client, we will refactor the MLflow quickstart script using the neptune.ai API.

import pandas as pd
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# Load the Iris dataset
X, y = datasets.load_iris(return_X_y=True)

# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Define the model hyperparameters
params = {
    "solver": "lbfgs",
    "max_iter": 1000,
    "multi_class": "auto",
    "random_state": 8888,
}

# Train the model
lr = LogisticRegression(**params)
lr.fit(X_train, y_train)

# Predict on the test set
y_pred = lr.predict(X_test)

# Calculate metrics
accuracy = accuracy_score(y_test, y_pred)

import mlflow
from mlflow.models import infer_signature

# Set our tracking server uri for logging
mlflow.set_tracking_uri(uri="http://127.0.0.1:8080")

# Create a new MLflow Experiment
mlflow.set_experiment("MLflow Quickstart")

# Start an MLflow run
with mlflow.start_run():
                                                  

# Infer the model signature
signature = infer_signature(X_train, lr_model.predict(X_train))

# Log the model
model_info = mlflow.sklearn.log_model(
        sk_model=lr_model,
        artifact_path="iris_model",
        signature=signature,
        input_example=X_train,
        registered_model_name="tracking-quickstart",
    )

# Log the hyperparameters
mlflow.log_params(params)
                                   

# Log the loss metric
mlflow.log_metric("accuracy", accuracy)

# Set a tag that we can use to remind ourselves what this run was for
mlflow.set_tag("Training Info", "Basic LR model for iris data")

# MLflow uses model aliases instead of stages
client.set_registered_model_alias("tracking-quickstart", "staging", 1)

import pandas as pd
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# Load the Iris dataset
X, y = datasets.load_iris(return_X_y=True)

# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Define the model hyperparameters
params = {
    "solver": "lbfgs",
    "max_iter": 1000,
    "multi_class": "auto",
    "random_state": 8888,
}

# Train the model
lr = LogisticRegression(**params)
lr.fit(X_train, y_train)

# Predict on the test set
y_pred = lr.predict(X_test)

# Calculate metrics
accuracy = accuracy_score(y_test, y_pred)

import neptune
import joblib
                                 

# Initialize a Neptune run
run = neptune.init_run()
                                           

# Initialize a new model
model = neptune.init_model(
    name="Logistic Regression Model",
    key="LRMODEL",
)

# Create a new model version
model_id = model["sys/id"].fetch()
model_version = neptune.init_model_version(model=model_id)

# Save the model on disk
save_model_filename = 'lr_model.joblib'
dump(lr_model, save_model_filename)

# Upload the model weights to the new model version
model_version["model"].upload(save_model_filename)

# Add the training hyper-parameters to the model logs
model_version.assign({"parameters": params})

# Add the accuracy to the metrics of the model_version
model_version["metrics/accuracy"] = accuracy

# Add tags to the run
run["sys/tags"].add(["Training Info", "Basic LR model for iris data"])

# Change the stage of the model version to staging
model_version.change_stage("staging")

After running the refactored script, we can access all the uploaded metrics, parameters, and model weights in the Neptune UI:

Cheat Sheet: Converting training scripts from the MLflow client to Neptune’s client

For more information, read the Neptune API and the MLflow API documentation.

Navigating Neptune as an MLflow user

Experiment trackers are pretty similar, and switching from MLflow to Neptune is easier than you might think.

Let’s walk through Neptune’s interface, focusing on the similarities and differences to MLflow’s core features.

Experiments and projects

MLflow aggregates model runs in Experiments. In Neptune, these are called Projects.

When you log into Neptune, your projects are in the center of the workspace. This is similar to how MLflow displays experiments in the left-hand sidebar:

Run list view

Neptune and MLflow display the run list view in the center of the UI. This is where you see all your experiment runs at a glance:

Neptune provides more customization regarding which information to display:

• You can create multiple custom views. (MLflow only provides one.)
• You can create columns from any information you store in your runs, while MLflow limits this to attributes, metrics, and parameters.
• Neptune displays the owner of the run, whereas MLflow does not support user access management.

Run details view

For both tools, clicking on a row in the run list takes you to the run details page. This is where you see all the information related to a specific experiment run:

Neptune offers a multi-pane view that displays the run list, the logs, and the visualizations on one single page. MLflow is less flexible here and only displays the key information on the run details page. To access the visualizations, you need to click on a metric, which will open a new page.

Model list view

In MLflow, you access the model view from the top-left corner. In Neptune, you reach this view from the left-hand sidebar:

MLflow combines all the models from various experiments. This setup can be confusing when many teams use the same experiment tracker because everyone’s records get mixed together. Neptune organizes models into different projects, making it easier to view only your own models.

Model details view

For both trackers, clicking on a row in the model list view brings you to the model details view:

This is where you can view the information related to a specific model version.

Neptune supports logging detailed information for each model version. You can log training metrics and even the inference speed on specific hardware. Instead, MLflow only shows basic information, such as the input and output schema, tags, and description.

Conclusion

Migrating from MLflow to neptune.ai is easier than it seems. Neptune offers all the features, documentation, and tools to simplify transferring your experiment logs, adapting your training scripts, managing user accounts, and onboarding your team.

Starting the migration is simple. You just need a computer with Python and ensure you can connect to your MLflow tracking server via the MLflow API.

Here’s a summary of the main steps to perform the migration:

Using Neptune as an MLflow user is straightforward, as both tools share many similarities. Adapting your experiment workflows requires little effort and yields great benefits due to Neptune’s extensive improvements.

We hope this guide has helped you execute or plan your migration. If you have any further questions or feedback, please don’t hesitate to reach out. We’re here to support you in every step of the process.