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.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 # nosec
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)
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.