Neptune

zenml.integrations.neptune special

Module containing Neptune integration.

NeptuneIntegration (Integration)

Definition of the neptune.ai integration with ZenML.

Source code in zenml/integrations/neptune/__init__.py

class NeptuneIntegration(Integration):
    """Definition of the neptune.ai integration with ZenML."""

    NAME = NEPTUNE
    REQUIREMENTS = ["neptune"]

    @classmethod
    def flavors(cls) -> List[Type[Flavor]]:
        """Declare the stack component flavors for the Neptune integration.

        Returns:
            List of stack component flavors for this integration.
        """
        from zenml.integrations.neptune.flavors import (
            NeptuneExperimentTrackerFlavor,
        )

        return [
            NeptuneExperimentTrackerFlavor,
        ]

flavors() classmethod

Declare the stack component flavors for the Neptune integration.

Returns:

Type	Description
List[Type[zenml.stack.flavor.Flavor]]	List of stack component flavors for this integration.

Source code in zenml/integrations/neptune/__init__.py

@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Declare the stack component flavors for the Neptune integration.

    Returns:
        List of stack component flavors for this integration.
    """
    from zenml.integrations.neptune.flavors import (
        NeptuneExperimentTrackerFlavor,
    )

    return [
        NeptuneExperimentTrackerFlavor,
    ]

experiment_trackers special

Initialization of Neptune experiment tracker.

neptune_experiment_tracker

Implementation of Neptune Experiment Tracker.

NeptuneExperimentTracker (BaseExperimentTracker)

Track experiments using neptune.ai.

Source code in zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

class NeptuneExperimentTracker(BaseExperimentTracker):
    """Track experiments using neptune.ai."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the experiment tracker.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.
        """
        super().__init__(*args, **kwargs)
        self.run_state: RunProvider = RunProvider()

    @property
    def config(self) -> NeptuneExperimentTrackerConfig:
        """Returns the `NeptuneExperimentTrackerConfig` config.

        Returns:
            The configuration.
        """
        return cast(NeptuneExperimentTrackerConfig, self._config)

    @property
    def settings_class(self) -> Optional[Type["BaseSettings"]]:
        """Settings class for the Neptune experiment tracker.

        Returns:
            The settings class.
        """
        return NeptuneExperimentTrackerSettings

    def prepare_step_run(self, info: "StepRunInfo") -> None:
        """Initializes a Neptune run and stores it in the run_state object.

        The run object can then be accessed later from other places, such as a step.

        Args:
            info: Info about the step that was executed.
        """
        settings = cast(
            NeptuneExperimentTrackerSettings, self.get_settings(info)
        )

        self.run_state.initialize(
            project=self.config.project,
            token=self.config.api_token,
            run_name=info.run_name,
            tags=list(settings.tags),
        )

    def get_step_run_metadata(
        self, info: "StepRunInfo"
    ) -> Dict[str, "MetadataType"]:
        """Get component- and step-specific metadata after a step ran.

        Args:
            info: Info about the step that was executed.

        Returns:
            A dictionary of metadata.
        """
        run_url = self.run_state.active_run.get_url()
        return {
            METADATA_EXPERIMENT_TRACKER_URL: Uri(run_url),
        }

    def cleanup_step_run(self, info: "StepRunInfo", step_failed: bool) -> None:
        """Stop the Neptune run.

        Args:
            info: Info about the step that was executed.
            step_failed: Whether the step failed or not.
        """
        self.run_state.active_run.sync()
        self.run_state.active_run.stop()
        self.run_state.reset()

config: NeptuneExperimentTrackerConfig property readonly

Returns the NeptuneExperimentTrackerConfig config.

Returns:

Type	Description
NeptuneExperimentTrackerConfig	The configuration.

settings_class: Optional[Type[BaseSettings]] property readonly

Settings class for the Neptune experiment tracker.

Returns:

Type	Description
Optional[Type[BaseSettings]]	The settings class.

__init__(self, *args, **kwargs) special

Initialize the experiment tracker.

Parameters:

Name	Type	Description	Default
*args	Any	Variable length argument list.	()
**kwargs	Any	Arbitrary keyword arguments.	{}

Source code in zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    """Initialize the experiment tracker.

    Args:
        *args: Variable length argument list.
        **kwargs: Arbitrary keyword arguments.
    """
    super().__init__(*args, **kwargs)
    self.run_state: RunProvider = RunProvider()

cleanup_step_run(self, info, step_failed)

Stop the Neptune run.

Parameters:

Name	Type	Description	Default
info	StepRunInfo	Info about the step that was executed.	required
step_failed	bool	Whether the step failed or not.	required

Source code in zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

def cleanup_step_run(self, info: "StepRunInfo", step_failed: bool) -> None:
    """Stop the Neptune run.

    Args:
        info: Info about the step that was executed.
        step_failed: Whether the step failed or not.
    """
    self.run_state.active_run.sync()
    self.run_state.active_run.stop()
    self.run_state.reset()

get_step_run_metadata(self, info)

Get component- and step-specific metadata after a step ran.

Parameters:

Name	Type	Description	Default
info	StepRunInfo	Info about the step that was executed.	required

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary of metadata.

Source code in zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

def get_step_run_metadata(
    self, info: "StepRunInfo"
) -> Dict[str, "MetadataType"]:
    """Get component- and step-specific metadata after a step ran.

    Args:
        info: Info about the step that was executed.

    Returns:
        A dictionary of metadata.
    """
    run_url = self.run_state.active_run.get_url()
    return {
        METADATA_EXPERIMENT_TRACKER_URL: Uri(run_url),
    }

prepare_step_run(self, info)

Initializes a Neptune run and stores it in the run_state object.

The run object can then be accessed later from other places, such as a step.

Parameters:

Name	Type	Description	Default
info	StepRunInfo	Info about the step that was executed.	required

Source code in zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

def prepare_step_run(self, info: "StepRunInfo") -> None:
    """Initializes a Neptune run and stores it in the run_state object.

    The run object can then be accessed later from other places, such as a step.

    Args:
        info: Info about the step that was executed.
    """
    settings = cast(
        NeptuneExperimentTrackerSettings, self.get_settings(info)
    )

    self.run_state.initialize(
        project=self.config.project,
        token=self.config.api_token,
        run_name=info.run_name,
        tags=list(settings.tags),
    )

run_state

Contains objects that create a Neptune run and store its state throughout the pipeline.

RunProvider

Singleton object used to store and persist a Neptune run state across the pipeline.

Source code in zenml/integrations/neptune/experiment_trackers/run_state.py

class RunProvider(metaclass=SingletonMetaClass):
    """Singleton object used to store and persist a Neptune run state across the pipeline."""

    def __init__(self) -> None:
        """Initialize RunProvider. Called with no arguments."""
        self._active_run: Optional["Run"] = None
        self._project: Optional[str] = None
        self._run_name: Optional[str] = None
        self._token: Optional[str] = None
        self._tags: Optional[List[str]] = None
        self._initialized = False

    def initialize(
        self,
        project: Optional[str] = None,
        token: Optional[str] = None,
        run_name: Optional[str] = None,
        tags: Optional[List[str]] = None,
    ) -> None:
        """Initialize the run state.

        Args:
            project: The neptune project.
            token: The neptune token.
            run_name: The neptune run name.
            tags: Tags for the neptune run.
        """
        self._project = project
        self._token = token
        self._run_name = run_name
        self._tags = tags
        self._initialized = True

    @property
    def project(self) -> Optional[Any]:
        """Getter for project name.

        Returns:
            Name of the project passed to the RunProvider.
        """
        return self._project

    @property
    def token(self) -> Optional[Any]:
        """Getter for API token.

        Returns:
            Neptune API token passed to the RunProvider.
        """
        return self._token

    @property
    def run_name(self) -> Optional[Any]:
        """Getter for run name.

        Returns:
            Name of the pipeline run.
        """
        return self._run_name

    @property
    def tags(self) -> Optional[Any]:
        """Getter for run tags.

        Returns:
            Tags associated with a Neptune run.
        """
        return self._tags

    @property
    def initialized(self) -> bool:
        """If the run state is initialized.

        Returns:
            If the run state is initialized.
        """
        return self._initialized

    @property
    def active_run(self) -> "Run":
        """Initializes a new neptune run every time it is called.

        The run is closed and the active run state is set to stopped
        after each step is completed.

        Returns:
            Neptune run object
        """
        if self._active_run is None:
            run = neptune.init_run(
                project=self.project,
                api_token=self.token,
                custom_run_id=md5(self.run_name.encode()).hexdigest(),  # type: ignore  # nosec
                tags=self.tags,
            )
            run[_INTEGRATION_VERSION_KEY] = zenml.__version__
            self._active_run = run
        return self._active_run

    def reset(self) -> None:
        """Reset the run state."""
        self._active_run = None
        self._project = None
        self._run_name = None
        self._token = None
        self._tags = None
        self._initialized = False

active_run: Run property readonly

Initializes a new neptune run every time it is called.

The run is closed and the active run state is set to stopped after each step is completed.

Returns:

Type	Description
Run	Neptune run object

initialized: bool property readonly

If the run state is initialized.

Returns:

Type	Description
bool	If the run state is initialized.

project: Optional[Any] property readonly

Getter for project name.

Returns:

Type	Description
Optional[Any]	Name of the project passed to the RunProvider.

run_name: Optional[Any] property readonly

Getter for run name.

Returns:

Type	Description
Optional[Any]	Name of the pipeline run.

tags: Optional[Any] property readonly

Getter for run tags.

Returns:

Type	Description
Optional[Any]	Tags associated with a Neptune run.

token: Optional[Any] property readonly

Getter for API token.

Returns:

Type	Description
Optional[Any]	Neptune API token passed to the RunProvider.

__init__(self) special

Initialize RunProvider. Called with no arguments.

Source code in zenml/integrations/neptune/experiment_trackers/run_state.py

def __init__(self) -> None:
    """Initialize RunProvider. Called with no arguments."""
    self._active_run: Optional["Run"] = None
    self._project: Optional[str] = None
    self._run_name: Optional[str] = None
    self._token: Optional[str] = None
    self._tags: Optional[List[str]] = None
    self._initialized = False

initialize(self, project=None, token=None, run_name=None, tags=None)

Initialize the run state.

Parameters:

Name	Type	Description	Default
project	Optional[str]	The neptune project.	None
token	Optional[str]	The neptune token.	None
run_name	Optional[str]	The neptune run name.	None
tags	Optional[List[str]]	Tags for the neptune run.	None

Source code in zenml/integrations/neptune/experiment_trackers/run_state.py

def initialize(
    self,
    project: Optional[str] = None,
    token: Optional[str] = None,
    run_name: Optional[str] = None,
    tags: Optional[List[str]] = None,
) -> None:
    """Initialize the run state.

    Args:
        project: The neptune project.
        token: The neptune token.
        run_name: The neptune run name.
        tags: Tags for the neptune run.
    """
    self._project = project
    self._token = token
    self._run_name = run_name
    self._tags = tags
    self._initialized = True

reset(self)

Reset the run state.

Source code in zenml/integrations/neptune/experiment_trackers/run_state.py

def reset(self) -> None:
    """Reset the run state."""
    self._active_run = None
    self._project = None
    self._run_name = None
    self._token = None
    self._tags = None
    self._initialized = False

get_neptune_run()

Helper function to fetch an existing Neptune run or create a new one.

Returns:

Type	Description
Run	Neptune run object

Exceptions:

Type	Description
RuntimeError	When unable to fetch the active neptune run.

Source code in zenml/integrations/neptune/experiment_trackers/run_state.py

def get_neptune_run() -> "Run":
    """Helper function to fetch an existing Neptune run or create a new one.

    Returns:
        Neptune run object

    Raises:
        RuntimeError: When unable to fetch the active neptune run.
    """
    from zenml.integrations.neptune.experiment_trackers import (
        NeptuneExperimentTracker,
    )

    experiment_tracker = Client().active_stack.experiment_tracker

    if not experiment_tracker:
        raise RuntimeError(
            "Unable to get neptune run: Missing experiment tracker in the "
            "active stack."
        )

    if not isinstance(experiment_tracker, NeptuneExperimentTracker):
        raise RuntimeError(
            "Unable to get neptune run: Experiment tracker in the active "
            f"stack ({experiment_tracker.flavor}) is not a neptune experiment "
            "tracker."
        )

    run_state = experiment_tracker.run_state
    if not run_state.initialized:
        raise RuntimeError(
            "Unable to get neptune run: The experiment tracker has not been "
            "initialized. To solve this, make sure you use the experiment "
            "tracker in your step. See "
            "https://docs.zenml.io/stack-components/experiment-trackers/neptune#how-do-you-use-it "
            "for more information."
        )

    return experiment_tracker.run_state.active_run

flavors special

Neptune integration flavors.

neptune_experiment_tracker_flavor

Neptune experiment tracker flavor.

NeptuneExperimentTrackerConfig (BaseExperimentTrackerConfig)

Config for the Neptune experiment tracker.

If attributes are left as None, the neptune.init_run() method will try to find the relevant values in the environment

Attributes:

Name	Type	Description
project	Optional[str]	name of the Neptune project you want to log the metadata to
api_token	Optional[str]	your Neptune API token

Source code in zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py

class NeptuneExperimentTrackerConfig(BaseExperimentTrackerConfig):
    """Config for the Neptune experiment tracker.

    If attributes are left as None, the neptune.init_run() method
    will try to find the relevant values in the environment

    Attributes:
        project: name of the Neptune project you want to log the metadata to
        api_token: your Neptune API token
    """

    project: Optional[str] = None
    api_token: Optional[str] = SecretField(default=None)

NeptuneExperimentTrackerFlavor (BaseExperimentTrackerFlavor)

Class for the NeptuneExperimentTrackerFlavor.

Source code in zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py

class NeptuneExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
    """Class for the `NeptuneExperimentTrackerFlavor`."""

    @property
    def name(self) -> str:
        """Name of the flavor.

        Returns:
            The name of the flavor.
        """
        return NEPTUNE_MODEL_EXPERIMENT_TRACKER_FLAVOR

    @property
    def docs_url(self) -> Optional[str]:
        """A url to point at docs explaining this flavor.

        Returns:
            A flavor docs url.
        """
        return self.generate_default_docs_url()

    @property
    def sdk_docs_url(self) -> Optional[str]:
        """A url to point at SDK docs explaining this flavor.

        Returns:
            A flavor SDK docs url.
        """
        return self.generate_default_sdk_docs_url()

    @property
    def logo_url(self) -> str:
        """A url to represent the flavor in the dashboard.

        Returns:
            The flavor logo.
        """
        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/experiment_tracker/neptune.png"

    @property
    def config_class(self) -> Type[NeptuneExperimentTrackerConfig]:
        """Returns `NeptuneExperimentTrackerConfig` config class.

        Returns:
                The config class.
        """
        return NeptuneExperimentTrackerConfig

    @property
    def implementation_class(self) -> Type["NeptuneExperimentTracker"]:
        """Implementation class for this flavor.

        Returns:
            The implementation class.
        """
        from zenml.integrations.neptune.experiment_trackers import (
            NeptuneExperimentTracker,
        )

        return NeptuneExperimentTracker

config_class: Type[zenml.integrations.neptune.flavors.neptune_experiment_tracker_flavor.NeptuneExperimentTrackerConfig] property readonly

Returns NeptuneExperimentTrackerConfig config class.

Returns:

Type	Description
Type[zenml.integrations.neptune.flavors.neptune_experiment_tracker_flavor.NeptuneExperimentTrackerConfig]	The config class.

docs_url: Optional[str] property readonly

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[NeptuneExperimentTracker] property readonly

Implementation class for this flavor.

Returns:

Type	Description
Type[NeptuneExperimentTracker]	The implementation class.

logo_url: str property readonly

A url to represent the flavor in the dashboard.

Returns:

Type	Description
str	The flavor logo.

name: str property readonly

Name of the flavor.

Returns:

Type	Description
str	The name of the flavor.

sdk_docs_url: Optional[str] property readonly

A url to point at SDK docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor SDK docs url.

NeptuneExperimentTrackerSettings (BaseSettings)

Settings for the Neptune experiment tracker.

Attributes:

Name	Type	Description
tags	Set[str]	Tags for the Neptune run.

Source code in zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py

class NeptuneExperimentTrackerSettings(BaseSettings):
    """Settings for the Neptune experiment tracker.

    Attributes:
        tags: Tags for the Neptune run.
    """

    tags: Set[str] = set()

neptune_constants

Some constants for reading environment variables.