Skip to content

ZenML integration guide#

Custom dashboard displaying metadata logged with ZenML

ZenML is a technology-agnostic, open source pipelines framework. With the Neptune-ZenML integration, you can log and visualize information from your ZenML pipeline steps (such as models, parameters, metrics).

See in Neptune  Code examples 

Setup#

Installing the Neptune-ZenML integration#

The Neptune Experiment Tracker flavor is provided by the Neptune-ZenML integration. To register the Neptune Experiment Tracker and add it to your stack, start by installing the integration:

zenml integration install neptune -y

Storing your Neptune credentials#

As a best practice, you should save your Neptune credentials as environment variables.

Setting your API token#

  1. In the bottom-left corner of the app, expand your user menu.
  2. Select Get Your API token.

    How to find your Neptune API token

  3. Depending on your system:

    From the API token dialog in Neptune, copy the export command and append the line to your .profile or other shell initialization file.

    Example line
    export NEPTUNE_API_TOKEN="uyVrZXkiOiIzNTd..."
    

    From the API token dialog in Neptune, copy the export command and append the line to your .profile or other shell initialization file.

    Example line
    export NEPTUNE_API_TOKEN="uyVrZXkiOiIzNTd..."
    
    1. From the API token dialog in Neptune, copy the setx command.

      Example line
      setx NEPTUNE_API_TOKEN "uyVrZXkiOiIzNTd..."
      
    2. Open a terminal app, such as PowerShell or Command Prompt.

    3. Paste the line you copied and press enter.
    4. To activate the change, restart the terminal app.

    You can also navigate to SettingsEdit the system environment variables and add the variable there.

    You can use the os library to set the token as an environment variable.

    Add the following to a notebook cell:

    import os
    from getpass import getpass
    os.environ["NEPTUNE_API_TOKEN"] = getpass("Enter your Neptune API token: ")
    

    From the API token dialog, copy your token, paste it in the input box, and hit Enter.

    Note that any environment variables declared this way won't persist after the notebook kernel shuts down. If you start a new kernel, they need to be set again.

Setting your project name#

  1. Open the project settings ( ).
  2. Select Details & privacy.

    How to access project details

  3. Find the copy button () next to the project name.

  4. Assign the project name to an environment variable named NEPTUNE_PROJECT:

    Append the following line to your .profile (or other shell configuration file):

    export NEPTUNE_PROJECT="workspace-name/project-name"
    

    where workspace-name/project-name is the full project name you just copied.

    Append the following line to your .profile (or other shell configuration file):

    export NEPTUNE_PROJECT="workspace-name/project-name"
    

    where workspace-name/project-name is the full project name you just copied.

    To set the project name permanently:

    1. Open a terminal app, such as PowerShell or Command Prompt.
    2. Paste in the following command and press enter:

      setx NEPTUNE_PROJECT "workspace-name/project-name"
      

      where workspace-name/project-name is the full project name you just copied.

    3. To activate the change, restart the terminal.

    You can also navigate to SettingsEdit the system environment variables and add the variable there.

    You can use the os library to set the project name as an environment variable:

    import os
    os.environ["NEPTUNE_PROJECT"] = "workspace-name/project-name"
    

    where workspace-name/project-name is the full project name you just copied.

    Note that any environment variables declared this way won't persist after the notebook kernel shuts down. If you start a new kernel, they need to be set again.

Registering the Neptune Experiment Tracker#

Register neptune_tracker with the neptune flavor:

zenml experiment-tracker register neptune_tracker --flavor=neptune # (1)!
  1. Neptune fetches the values of project and api_token from the NEPTUNE_PROJECT and NEPTUNE_API_TOKEN environment variables. If these have not been set, you can pass them directly to zenml experiment-tracker register (not recommended in production settings):

    zenml experiment-tracker register neptune_tracker \
        --flavor=neptune \
        --project="YOUR_NEPTUNE_PROJECT" \
        --api_token="YOUR_NEPTUNE_API_TOKEN"
    

Creating a stack with neptune_tracker as the experiment tracker#

Create a new stack with the Neptune Experiment Tracker.

The --set flag sets the new stack as the active one.

zenml stack register neptune_stack \
    -e neptune_tracker \
    ...
    --set

Adding Neptune to an existing stack

You can add the Neptune Experiment Tracker to an existing stack by using zenml stack update:

zenml stack update EXISTING_STACK_NAME -e neptune_tracker

Using the Neptune-ZenML integration#

Manually logging metadata to the run object#

You can access the Neptune Run object with the get_neptune_run() function. Use the object to track metadata from your training pipeline in a structure of your choosing.

from zenml.integrations.neptune.experiment_trackers.run_state import (
    get_neptune_run
)

@step(experiment_tracker="neptune_tracker", ...)
def my_step():
    neptune_run = get_neptune_run()
    neptune_run["sys/name"] = "My custom run name" # (1)!
    neptune_run["params/lr"] = params.lr # (2)!
    ...
  1. Sets a custom name, which you can use as a human-friendly ID.

    To display it in the app, add sys/name as a column.

    You can also edit the name in the run information view ( Run information).

  2. Assign any value to a path of your choice. This would store the learning rate under a field called lr inside a namespace called params.

Adding tags to the Neptune run#

You can use NeptuneExperimentTrackerSettings to add tags to the Neptune run during initialization:

from zenml.integrations.neptune.flavors import NeptuneExperimentTrackerSettings

neptune_settings = NeptuneExperimentTrackerSettings(
    tags={"regression", "script", "sklearn"}
)

@step(
    experiment_tracker="neptune_tracker",
    settings={"experiment_tracker.neptune": neptune_settings},
    ...
)
def my_step():
    ...

Adding tags later

You can also add tags to the runs later. To learn more, see Adding and managing tags.

Logging pipeline and step metadata to the run#

from zenml import get_step_context
from zenml.integrations.neptune.experiment_trackers.run_state import (
    get_neptune_run
)
from neptune.utils import stringify_unsupported

@step(experiment_tracker="neptune_tracker", ...)
def my_step():
    neptune_run = get_neptune_run()
    context = get_step_context()

    neptune_run["pipeline_metadata"] = stringify_unsupported(
        context.pipeline_run.get_metadata().dict()
    )    
    neptune_run[f"step_metadata/{context.step_name}"] = stringify_unsupported(
        context.step_run.get_metadata().dict()
    )
    ...

Related

More Neptune docs: