Skip to content

MLflow integration guide#

Open in Colab

The Neptune-MLflow integration lets send your metadata to Neptune while using MLflow logging code. You can also use the migration tool to export existing run metadata from MLflow to Neptune.

See example in Neptune  Code examples 

Before you start#

Installing the integration#

To use your preinstalled version of Neptune together with the integration:

pip install -U neptune-mlflow

To install both Neptune and the integration:

pip install -U neptune neptune-mlflow

Setting up the Neptune URI (plugin)#

To enable Neptune logging in your MLflow code, you set up the tracking URI to point to Neptune. This way, you can send your metadata to Neptune without changing your MLflow code.

  1. Use the create_neptune_tracking_uri() function to set up a Neptune tracking URI.

    You can pass extra arguments to customize the Neptune run that will be created (see More Neptune URI options).

    from neptune_mlflow_plugin import create_neptune_tracking_uri

neptune_uri = create_neptune_tracking_uri()  # (1)!

    1. If you haven't registered or set up your Neptune credentials, you can use the following arguments for testing:

      neptune_uri = create_neptune_tracking_uri(
    api_token=neptune.ANONYMOUS_API_TOKEN,
    project="common/mlflow-integration",
)

  2. Then pass the created neptune_uri to the set_tracking_uri() function:

    mlflow.set_tracking_uri(neptune_uri)

    MLflow will now start logging the run metadata to Neptune.

  3. Run your MLflow logging script as you normally would.

    with mlflow.start_run():
    ...

To open the run, click the Neptune link that appears in the console output.

Sample output

[neptune] [info ] Neptune initialized. Open in the app: https://app.neptune.ai/workspace/project/e/RUN-1

In the above example, the run ID is RUN-1.

If Neptune can't find your project name or API token

As a best practice, you should save your Neptune API token and project name as environment variables:

export NEPTUNE_API_TOKEN="h0dHBzOi8aHR0cHM6Lkc78ghs74kl0jv...Yh3Kb8"

export NEPTUNE_PROJECT="ml-team/classification"

Alternatively, you can pass the information when using a function that takes api_token and project as arguments:

run = neptune.init_run( # (1)!
    api_token="h0dHBzOi8aHR0cHM6Lkc78ghs74kl0jv...Yh3Kb8",  # your token here
    project="ml-team/classification",  # your full project name here
)

  1. Also works for init_model(), init_model_version(), init_project(), and integrations that create Neptune runs underneath the hood, such as NeptuneLogger or NeptuneCallback.

  2. API token: In the bottom-left corner, expand the user menu and select Get my API token.

  3. Project name: You can copy the path from the project details ( → Edit project details).

If you haven't registered, you can log anonymously to a public project:

api_token=neptune.ANONYMOUS_API_TOKEN
project="common/quickstarts"

Make sure not to publish sensitive data through your code!

More Neptune URI options#

To customize the Neptune run, you can provide additional neptune.init_run() options.

neptune_uri = create_neptune_tracking_uri(
    name="MLflow run metadata",
    description="Some longer description",
    tags=["mlflow", "plugin"],
    dependencies="infer",
    source_files="**/.py",
    ...,
)

The exceptions are with_id and custom_run_id, which will be ignored if provided (as calling the create_neptune_tracking_uri() function always creates a new Neptune run).

Exporting MLflow runs to Neptune#

neptune-notebooks incompatibility

Currently, the CLI component of the integration does not work together with the Neptune-Jupyter extension (neptune-notebooks).

Until a fix is released, if you have neptune-notebooks installed, you must uninstall it to be able to use the neptune mlflow command.

You can export your logged metadata from MLflow to Neptune with the neptune mlflow CLI command:

neptune mlflow

If you haven't set your credentials as environment variables, you can also pass them as arguments:

neptune mlflow --api-token YourNeptuneApiToken --project YourProjectName
How do I save my credentials as environment variables?

Set your Neptune API token and full project name to the NEPTUNE_API_TOKEN and NEPTUNE_PROJECT environment variables, respectively.

export NEPTUNE_API_TOKEN="h0dHBzOi8aHR0cHM.4kl0jvYh3Kb8...ifQ=="

export NEPTUNE_PROJECT="ml-team/classification"

export NEPTUNE_API_TOKEN="h0dHBzOi8aHR0cHM.4kl0jvYh3Kb8...ifQ=="

export NEPTUNE_PROJECT="ml-team/classification"

setx NEPTUNE_API_TOKEN "h0dHBzOi8aHR0cHM.4kl0jvYh3Kb8...ifQ=="

setx NEPTUNE_PROJECT "ml-team/classification"

You can also navigate to Settings → Edit the system environment variables and add the variables there.

%env NEPTUNE_API_TOKEN="h0dHBzOi8aHR0cHM.4kl0jvYh3Kb8...ifQ=="

%env NEPTUNE_PROJECT="ml-team/classification"

To find your credentials:

  • API token: In the bottom-left corner of the Neptune app, expand your user menu and select Get your API token. If you need the token of a service account, go to the workspace or project settings and enter the Service accounts settings.
  • Project name: Your full project name has the form workspace-name/project-name. You can copy it from the project menu ( → Edit project details).

If you're working in Google Colab, you can set your credentials with the os and getpass libraries:

import os
from getpass import getpass
os.environ["NEPTUNE_API_TOKEN"] = getpass("Enter your Neptune API token: ")
os.environ["NEPTUNE_PROJECT"] = "workspace-name/project-name"

You'll find the imported metadata in All metadata, in the following structure:

artifacts
    |-- <logged artifacts>
experiment
    |-- creation_time
    |-- experiment_id
    |-- last_update_time
    |-- name
run_data
    |-- metrics
    |-- params
    |-- tags
run_info
    |-- artifact_uri
    |-- end_time
    |-- experiment_id
    |-- lifecycle_stage
    |-- run_id
    |-- run_name
    |-- run_uuid
    |-- start_time
    |-- status
    |-- user_id

Artifact options#

You can set the maximum MB size of artifacts to be exported to Neptune. The default is 50.

neptune mlflow --max-artifact-size=100

To exclude artifacts from the export entirely, use the --exclude-artifacts option:

neptune mlflow --exclude-artifacts

Related