From df3040f0a3e3ae9c1a16be02ffe66a1b0ab5c6eb Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Fri, 6 Dec 2024 10:35:35 +0100 Subject: [PATCH] Implement wandb settings conversion for latest release (#3246) --- .../flavors/wandb_experiment_tracker_flavor.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py b/src/zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py index 11e84ca2d38..7a1a7321704 100644 --- a/src/zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py +++ b/src/zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py @@ -23,7 +23,7 @@ cast, ) -from pydantic import field_validator +from pydantic import field_validator, BaseModel from zenml.config.base_settings import BaseSettings from zenml.experiment_trackers.base_experiment_tracker import ( @@ -60,18 +60,26 @@ def _convert_settings(cls, value: Any) -> Any: Args: value: The settings. + Raises: + ValueError: If converting the settings failed. + Returns: Dict representation of the settings. """ import wandb if isinstance(value, wandb.Settings): - # Depending on the wandb version, either `make_static` or `to_dict` - # is available to convert the settings to a dictionary - if hasattr(value, "make_static"): + # Depending on the wandb version, either `model_dump`, + # `make_static` or `to_dict` is available to convert the settings + # to a dictionary + if isinstance(value, BaseModel): + return value.model_dump() + elif hasattr(value, "make_static"): return cast(Dict[str, Any], value.make_static()) - else: + elif hasattr(value, "to_dict"): return value.to_dict() + else: + raise ValueError("Unable to convert wandb settings to dict.") else: return value