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-client>=0.16.10",
    ]

    @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.token = self.config.api_token
        self.run_state.project = self.config.project
        self.run_state.run_name = info.run_name
        self.run_state.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_active_run()

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_active_run()

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.token = self.config.api_token
    self.run_state.project = self.config.project
    self.run_state.run_name = info.run_name
    self.run_state.tags = list(settings.tags)

run_state

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

InvalidExperimentTrackerSelected (Exception)

Raised if a Neptune run is fetched while using a different experiment tracker.

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

class InvalidExperimentTrackerSelected(Exception):
    """Raised if a Neptune run is fetched while using a different experiment tracker."""

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 = None
        self._project: Optional[str]
        self._run_name: Optional[str]
        self._token: Optional[str]
        self._tags: Optional[List[str]]

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

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

    @project.setter
    def project(self, project: str) -> None:
        """Setter for project name.

        Args:
            project: Neptune project name
        """
        self._project = project

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

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

    @token.setter
    def token(self, token: str) -> None:
        """Setter for API token.

        Args:
            token: Neptune API token
        """
        self._token = token

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

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

    @run_name.setter
    def run_name(self, run_name: str) -> None:
        """Setter for run name.

        Args:
            run_name: name of the pipeline run
        """
        self._run_name = run_name

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

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

    @tags.setter
    def tags(self, tags: List[str]) -> None:
        """Setter for run tags.

        Args:
            tags: list of tags associated with a Neptune run
        """
        self._tags = tags

    @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
                tags=self.tags,
            )
            run[_INTEGRATION_VERSION_KEY] = zenml.__version__
            self._active_run = run
        return self._active_run

    def reset_active_run(self) -> None:
        """Resets the active run state to None."""
        self._active_run = None

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

project: Optional[Any] property writable

Getter for project name.

Returns:

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

run_name: Optional[Any] property writable

Getter for run name.

Returns:

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

tags: Optional[Any] property writable

Getter for run tags.

Returns:

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

token: Optional[Any] property writable

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 = None
    self._project: Optional[str]
    self._run_name: Optional[str]
    self._token: Optional[str]
    self._tags: Optional[List[str]]

reset_active_run(self)

Resets the active run state to None.

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

def reset_active_run(self) -> None:
    """Resets the active run state to None."""
    self._active_run = None

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
InvalidExperimentTrackerSelected	when called while using an experiment tracker other than Neptune

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:
        InvalidExperimentTrackerSelected: when called while using an experiment tracker other than Neptune
    """
    client = Client()
    experiment_tracker = client.active_stack.experiment_tracker
    if experiment_tracker.flavor == NEPTUNE:  # type: ignore
        return experiment_tracker.run_state.active_run  # type: ignore
    raise InvalidExperimentTrackerSelected(
        "Fetching a Neptune run works only with the 'neptune' flavor of "
        "the experiment tracker. The flavor currently selected is %s"
        % experiment_tracker.flavor  # type: ignore
    )

flavors special

Neptune integration flavors.

neptune_experiment_tracker_flavor

Neptune experiment tracker flavor.

NeptuneExperimentTrackerConfig (BaseExperimentTrackerConfig) pydantic-model

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()

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) pydantic-model

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.