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", step_failed: bool) -> None:
        """Stops the Wandb run.

        Args:
            info: Info about the step that was executed.
            step_failed: Whether the step failed or not.
        """
        wandb.finish(exit_code=1) if step_failed else 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, step_failed)`

Stops the Wandb 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/wandb/experiment_trackers/wandb_experiment_tracker.py

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

    Args:
        info: Info about the step that was executed.
        step_failed: Whether the step failed or not.
    """
    wandb.finish(exit_code=1) if step_failed else 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

Wandb

zenml.integrations.wandb special

WandbIntegration (Integration)

flavors() classmethod

experiment_trackers special

wandb_experiment_tracker

WandbExperimentTracker (BaseExperimentTracker)

config: WandbExperimentTrackerConfig property readonly

settings_class: Optional[Type[BaseSettings]] property readonly

cleanup_step_run(self, info, step_failed)

prepare_step_run(self, info)

flavors special

wandb_experiment_tracker_flavor

WandbExperimentTrackerConfig (BaseExperimentTrackerConfig, WandbExperimentTrackerSettings) pydantic-model

WandbExperimentTrackerFlavor (BaseExperimentTrackerFlavor)

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

implementation_class: Type[WandbExperimentTracker] property readonly

name: str property readonly

WandbExperimentTrackerSettings (BaseSettings) pydantic-model

`zenml.integrations.wandb` `special`

`WandbIntegration (Integration)`

`flavors()` `classmethod`

`experiment_trackers` `special`

`wandb_experiment_tracker`

`WandbExperimentTracker (BaseExperimentTracker)`

`config: WandbExperimentTrackerConfig` `property` `readonly`

`settings_class: Optional[Type[BaseSettings]]` `property` `readonly`

`cleanup_step_run(self, info, step_failed)`

`prepare_step_run(self, info)`

`flavors` `special`

`wandb_experiment_tracker_flavor`

`WandbExperimentTrackerConfig (BaseExperimentTrackerConfig, WandbExperimentTrackerSettings)` `pydantic-model`

`WandbExperimentTrackerFlavor (BaseExperimentTrackerFlavor)`

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

`implementation_class: Type[WandbExperimentTracker]` `property` `readonly`

`name: str` `property` `readonly`

`WandbExperimentTrackerSettings (BaseSettings)` `pydantic-model`