Skip to content

Wandb

zenml.integrations.wandb special

Initialization for the wandb integration.

The wandb integrations currently enables you to use wandb tracking as a convenient way to visualize your experiment runs within the wandb ui.

WandbIntegration (Integration)

Definition of Plotly integration for ZenML.

Source code in zenml/integrations/wandb/__init__.py 
class WandbIntegration(Integration):
    """Definition of Plotly integration for ZenML."""

    NAME = WANDB
    REQUIREMENTS = ["wandb>=0.12.12", "Pillow>=9.1.0"]

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

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

        return [WandbExperimentTrackerFlavor]

flavors() classmethod

Declare the stack component flavors for the Weights and Biases integration.

Returns:

Type Description
List[Type[zenml.stack.flavor.Flavor]]

List of stack component flavors for this integration.
Source code in zenml/integrations/wandb/__init__.py 
@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Declare the stack component flavors for the Weights and Biases integration.

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

    return [WandbExperimentTrackerFlavor]

experiment_trackers special

Initialization for the wandb experiment tracker.

wandb_experiment_tracker

Implementation for the wandb experiment tracker.

WandbExperimentTracker (BaseExperimentTracker)

Track experiment using Wandb.

Source code in zenml/integrations/wandb/experiment_trackers/wandb_experiment_tracker.py 
class WandbExperimentTracker(BaseExperimentTracker):
    """Track experiment using Wandb."""

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

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

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

        Returns:
            The settings class.
        """
        return WandbExperimentTrackerSettings

    def prepare_step_run(self, info: "StepRunInfo") -> None:
        """Configures a Wandb run.

        Args:
            info: Info about the step that will be executed.
        """
        os.environ[WANDB_API_KEY] = self.config.api_key
        settings = cast(WandbExperimentTrackerSettings, self.get_settings(info))

        tags = settings.tags + [info.run_name, info.pipeline.name]
        wandb_run_name = (
            settings.run_name or f"{info.run_name}_{info.config.name}"
        )
        self._initialize_wandb(
            run_name=wandb_run_name, tags=tags, settings=settings.settings
        )

    def cleanup_step_run(self, info: "StepRunInfo") -> None:
        """Stops the Wandb run.

        Args:
            info: Info about the step that was executed.
        """
        wandb.finish()
        os.environ.pop(WANDB_API_KEY, None)

    def _initialize_wandb(
        self,
        run_name: str,
        tags: List[str],
        settings: Union[wandb.Settings, Dict[str, Any], None] = None,
    ) -> None:
        """Initializes a wandb run.

        Args:
            run_name: Name of the wandb run to create.
            tags: Tags to attach to the wandb run.
            settings: Additional settings for the wandb run.
        """
        logger.info(
            f"Initializing wandb with entity {self.config.entity}, project "
            f"name: {self.config.project_name}, run_name: {run_name}."
        )
        wandb.init(
            entity=self.config.entity,
            project=self.config.project_name,
            name=run_name,
            tags=tags,
            settings=settings,
        )
config: WandbExperimentTrackerConfig property readonly

Returns the WandbExperimentTrackerConfig config.

Returns:

Type Description
WandbExperimentTrackerConfig

The configuration.
settings_class: Optional[Type[BaseSettings]] property readonly

settings class for the Wandb experiment tracker.

Returns:

Type Description
Optional[Type[BaseSettings]]

The settings class.
cleanup_step_run(self, info)

Stops the Wandb run.

Parameters:

Name Type Description Default
info StepRunInfo

Info about the step that was executed.

 required
Source code in zenml/integrations/wandb/experiment_trackers/wandb_experiment_tracker.py 
def cleanup_step_run(self, info: "StepRunInfo") -> None:
    """Stops the Wandb run.

    Args:
        info: Info about the step that was executed.
    """
    wandb.finish()
    os.environ.pop(WANDB_API_KEY, None)
prepare_step_run(self, info)

Configures a Wandb run.

Parameters:

Name Type Description Default
info StepRunInfo

Info about the step that will be executed.

 required
Source code in zenml/integrations/wandb/experiment_trackers/wandb_experiment_tracker.py 
def prepare_step_run(self, info: "StepRunInfo") -> None:
    """Configures a Wandb run.

    Args:
        info: Info about the step that will be executed.
    """
    os.environ[WANDB_API_KEY] = self.config.api_key
    settings = cast(WandbExperimentTrackerSettings, self.get_settings(info))

    tags = settings.tags + [info.run_name, info.pipeline.name]
    wandb_run_name = (
        settings.run_name or f"{info.run_name}_{info.config.name}"
    )
    self._initialize_wandb(
        run_name=wandb_run_name, tags=tags, settings=settings.settings
    )

flavors special

Weights & Biases integration flavors.

wandb_experiment_tracker_flavor

Weights & Biases experiment tracker flavor.

WandbExperimentTrackerConfig (BaseExperimentTrackerConfig, WandbExperimentTrackerSettings) pydantic-model

Config for the Wandb experiment tracker.

Attributes:

Name Type Description
entity Optional[str]

Name of an existing wandb entity.
project_name Optional[str]

Name of an existing wandb project to log to.
api_key str

API key to should be authorized to log to the configured wandb entity and project.
Source code in zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py 
class WandbExperimentTrackerConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
    BaseExperimentTrackerConfig, WandbExperimentTrackerSettings
):
    """Config for the Wandb experiment tracker.

    Attributes:
        entity: Name of an existing wandb entity.
        project_name: Name of an existing wandb project to log to.
        api_key: API key to should be authorized to log to the configured wandb
            entity and project.
    """

    api_key: str = SecretField()
    entity: Optional[str] = None
    project_name: Optional[str] = None
WandbExperimentTrackerFlavor (BaseExperimentTrackerFlavor)

Flavor for the Wandb experiment tracker.

Source code in zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py 
class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
    """Flavor for the Wandb experiment tracker."""

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

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

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

        Returns:
                The config class.
        """
        return WandbExperimentTrackerConfig

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

        Returns:
            The implementation class.
        """
        from zenml.integrations.wandb.experiment_trackers import (
            WandbExperimentTracker,
        )

        return WandbExperimentTracker
config_class: Type[zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor.WandbExperimentTrackerConfig] property readonly

Returns WandbExperimentTrackerConfig config class.

Returns:

Type Description
Type[zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor.WandbExperimentTrackerConfig]

The config class.
implementation_class: Type[WandbExperimentTracker] property readonly

Implementation class for this flavor.

Returns:

Type Description
Type[WandbExperimentTracker]

The implementation class.
name: str property readonly

Name of the flavor.

Returns:

Type Description
str

The name of the flavor.
WandbExperimentTrackerSettings (BaseSettings) pydantic-model

Settings for the Wandb experiment tracker.

Attributes:

Name Type Description
run_name Optional[str]

The Wandb run name.
tags List[str]

Tags for the Wandb run.
settings Dict[str, Any]

Settings for the Wandb run.
Source code in zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py 
class WandbExperimentTrackerSettings(BaseSettings):
    """Settings for the Wandb experiment tracker.

    Attributes:
        run_name: The Wandb run name.
        tags: Tags for the Wandb run.
        settings: Settings for the Wandb run.
    """

    run_name: Optional[str] = None
    tags: List[str] = []
    settings: Dict[str, Any] = {}

    @validator("settings", pre=True)
    def _convert_settings(
        cls, value: Union[Dict[str, Any], "wandb.Settings"]
    ) -> Dict[str, Any]:
        """Converts settings to a dictionary.

        Args:
            value: The settings.

        Returns:
            Dict representation of the settings.
        """
        import wandb

        if isinstance(value, wandb.Settings):
            return cast(Dict[str, Any], value.make_static())
        else:
            return value