zenml-nightly 0.70.0.dev20241201__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/integrations/neptune/experiment_trackers/run_state.py

@@ -20,7 +20,6 @@ import neptune
 
 import zenml
 from zenml.client import Client
-from zenml.integrations.constants import NEPTUNE
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -29,20 +28,38 @@ if TYPE_CHECKING:
 _INTEGRATION_VERSION_KEY = "source_code/integrations/zenml"
 
 
-class InvalidExperimentTrackerSelected(Exception):
-    """Raised if a Neptune run is fetched while using a different experiment tracker."""
-
-
 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: Optional["Run"] = None
-        self._project: Optional[str]
-        self._run_name: Optional[str]
-        self._token: Optional[str]
-        self._tags: Optional[List[str]]
+        self._project: Optional[str] = None
+        self._run_name: Optional[str] = None
+        self._token: Optional[str] = None
+        self._tags: Optional[List[str]] = None
+        self._initialized = False
+
+    def initialize(
+        self,
+        project: Optional[str] = None,
+        token: Optional[str] = None,
+        run_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the run state.
+
+        Args:
+            project: The neptune project.
+            token: The neptune token.
+            run_name: The neptune run name.
+            tags: Tags for the neptune run.
+        """
+        self._project = project
+        self._token = token
+        self._run_name = run_name
+        self._tags = tags
+        self._initialized = True
 
     @property
     def project(self) -> Optional[Any]:
@@ -53,15 +70,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         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.
@@ -71,15 +79,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         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.
@@ -89,15 +88,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         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.
@@ -107,14 +97,14 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._tags
 
-    @tags.setter
-    def tags(self, tags: List[str]) -> None:
-        """Setter for run tags.
+    @property
+    def initialized(self) -> bool:
+        """If the run state is initialized.
 
-        Args:
-            tags: list of tags associated with a Neptune run
+        Returns:
+            If the run state is initialized.
         """
-        self._tags = tags
+        return self._initialized
 
     @property
     def active_run(self) -> "Run":
@@ -137,9 +127,14 @@ class RunProvider(metaclass=SingletonMetaClass):
             self._active_run = run
         return self._active_run
 
-    def reset_active_run(self) -> None:
-        """Resets the active run state to None."""
+    def reset(self) -> None:
+        """Reset the run state."""
         self._active_run = None
+        self._project = None
+        self._run_name = None
+        self._token = None
+        self._tags = None
+        self._initialized = False
 
 
 def get_neptune_run() -> "Run":
@@ -149,14 +144,35 @@ def get_neptune_run() -> "Run":
         Neptune run object
 
     Raises:
-        InvalidExperimentTrackerSelected: when called while using an experiment tracker other than Neptune
+        RuntimeError: When unable to fetch the active neptune run.
     """
-    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
+    from zenml.integrations.neptune.experiment_trackers import (
+        NeptuneExperimentTracker,
     )
+
+    experiment_tracker = Client().active_stack.experiment_tracker
+
+    if not experiment_tracker:
+        raise RuntimeError(
+            "Unable to get neptune run: Missing experiment tracker in the "
+            "active stack."
+        )
+
+    if not isinstance(experiment_tracker, NeptuneExperimentTracker):
+        raise RuntimeError(
+            "Unable to get neptune run: Experiment tracker in the active "
+            f"stack ({experiment_tracker.flavor}) is not a neptune experiment "
+            "tracker."
+        )
+
+    run_state = experiment_tracker.run_state
+    if not run_state.initialized:
+        raise RuntimeError(
+            "Unable to get neptune run: The experiment tracker has not been "
+            "initialized. To solve this, make sure you use the experiment "
+            "tracker in your step. See "
+            "https://docs.zenml.io/stack-components/experiment-trackers/neptune#how-do-you-use-it "
+            "for more information."
+        )
+
+    return experiment_tracker.run_state.active_run

zenml/integrations/registry.py

@@ -111,7 +111,7 @@ class IntegrationRegistry(object):
                 )
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )
@@ -148,7 +148,7 @@ class IntegrationRegistry(object):
                 ].get_uninstall_requirements(target_os=target_os)
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -144,3 +144,15 @@ class SkypilotBaseOrchestratorConfig(
             True if this config is for a local component, False otherwise.
         """
         return False
+
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Skypilot orchestrator runs the entire pipeline in a single VM, or
+        # starts additional VMs from the root VM. Both of those cases are
+        # currently not supported when using client-side caching.
+        return False

zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py

@@ -23,7 +23,7 @@ from typing import (
     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 @@ class WandbExperimentTrackerSettings(BaseSettings):
         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
 

zenml/materializers/built_in_materializer.py

@@ -429,7 +429,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # doesn't work for non-serializable types as they
         # are saved as list of lists in different files
         if _is_serializable(data):
-            return {self.data_path: VisualizationType.JSON}
+            return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
         return {}
 
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:

zenml/model/model.py

@@ -509,22 +509,6 @@ class Model(BaseModel):
             raise ValueError(
                 "`model_version_id` field is for internal use only"
             )
-
-        version = data.get("version", None)
-
-        if (
-            version in [stage.value for stage in ModelStages]
-            and not suppress_class_validation_warnings
-        ):
-            logger.info(
-                f"Version `{version}` matches one of the possible "
-                "`ModelStages` and will be fetched using stage."
-            )
-        if str(version).isnumeric() and not suppress_class_validation_warnings:
-            logger.info(
-                f"`version` `{version}` is numeric and will be fetched "
-                "using version number."
-            )
         data["suppress_class_validation_warnings"] = True
         return data
 
@@ -603,6 +587,18 @@ class Model(BaseModel):
                 hydrate=hydrate,
             )
         else:
+            if self.version in ModelStages.values():
+                logger.info(
+                    f"Version `{self.version}` for model {self.name} matches "
+                    "one of the possible `ModelStages` and will be fetched "
+                    "using stage."
+                )
+            if str(self.version).isnumeric():
+                logger.info(
+                    f"Version `{self.version}` for model {self.name} is "
+                    "numeric and will be fetched using version number."
+                )
+
             mv = zenml_client.get_model_version(
                 model_name_or_id=self.name,
                 model_version_name_or_number_or_id=self.version,

zenml/model/utils.py

@@ -56,7 +56,9 @@ def log_model_metadata(
     """
     logger.warning(
         "The `log_model_metadata` function is deprecated and will soon be "
-        "removed. Please use `log_metadata` instead."
+        "removed. Instead, you can consider using: "
+        "`log_metadata(metadata={...}, infer_model=True)` instead. For more "
+        "info: https://docs.zenml.io/how-to/model-management-metrics/track-metrics-metadata/attach-metadata-to-a-model"
     )
 
     from zenml import log_metadata

zenml/models/v2/base/filter.py

@@ -113,7 +113,7 @@ class Filter(BaseModel, ABC):
     def generate_query_conditions(
         self,
         table: Type[SQLModel],
-    ) -> Union["ColumnElement[bool]"]:
+    ) -> "ColumnElement[bool]":
         """Generate the query conditions for the database.
 
         This method converts the Filter class into an appropriate SQLModel
@@ -291,11 +291,19 @@ class UUIDFilter(StrFilter):
         import sqlalchemy
         from sqlalchemy_utils.functions import cast_if
 
+        from zenml.utils import uuid_utils
+
         # For equality checks, compare the UUID directly
         if self.operation == GenericFilterOps.EQUALS:
+            if not uuid_utils.is_valid_uuid(self.value):
+                return False
+
             return column == self.value
 
         if self.operation == GenericFilterOps.NOT_EQUALS:
+            if not uuid_utils.is_valid_uuid(self.value):
+                return True
+
             return column != self.value
 
         # For all other operations, cast and handle the column as string
@@ -436,7 +444,6 @@ class BaseFilter(BaseModel):
         le=PAGE_SIZE_MAXIMUM,
         description="Page size",
     )
-
     id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Id for this resource",
@@ -491,13 +498,13 @@ class BaseFilter(BaseModel):
                 )
                 value = column
 
-        if column in cls.FILTER_EXCLUDE_FIELDS:
+        if column in cls.CUSTOM_SORTING_OPTIONS:
+            return value
+        elif column in cls.FILTER_EXCLUDE_FIELDS:
             raise ValueError(
                 f"This resource can not be sorted by this field: '{value}'"
             )
-        elif column in cls.model_fields:
-            return value
-        elif column in cls.CUSTOM_SORTING_OPTIONS:
+        if column in cls.model_fields:
             return value
         else:
             raise ValueError(
@@ -703,16 +710,10 @@ class BaseFilter(BaseModel):
 
         conditions = []
 
-        try:
-            filter_ = FilterGenerator(table).define_filter(
-                column="id", value=value, operator=operator
-            )
-            conditions.append(filter_.generate_query_conditions(table=table))
-        except ValueError:
-            # UUID filter with equal operators and no full UUID fail with
-            # a ValueError. In this case, we already know that the filter
-            # will not produce any result and can simply ignore it.
-            pass
+        filter_ = FilterGenerator(table).define_filter(
+            column="id", value=value, operator=operator
+        )
+        conditions.append(filter_.generate_query_conditions(table=table))
 
         filter_ = FilterGenerator(table).define_filter(
             column="name", value=value, operator=operator
@@ -759,7 +760,7 @@ class BaseFilter(BaseModel):
         return self.size * (self.page - 1)
 
     def generate_filter(
-        self, table: Type[SQLModel]
+        self, table: Type["AnySchema"]
     ) -> Union["ColumnElement[bool]"]:
         """Generate the filter for the query.
 
@@ -779,7 +780,7 @@ class BaseFilter(BaseModel):
             filters.append(
                 column_filter.generate_query_conditions(table=table)
             )
-        for custom_filter in self.get_custom_filters():
+        for custom_filter in self.get_custom_filters(table):
             filters.append(custom_filter)
         if self.logical_operator == LogicalOperators.OR:
             return or_(False, *filters)
@@ -788,12 +789,17 @@ class BaseFilter(BaseModel):
         else:
             raise RuntimeError("No valid logical operator was supplied.")
 
-    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
         This can be overridden by subclasses to define custom filters that are
         not based on the columns of the underlying table.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
@@ -1101,18 +1107,8 @@ class FilterGenerator:
             A Filter object.
 
         Raises:
-            ValueError: If the value is not a valid UUID.
+            ValueError: If the value for a oneof filter is not a list.
         """
-        # For equality checks, ensure that the value is a valid UUID.
-        if operator == GenericFilterOps.EQUALS and not isinstance(value, UUID):
-            try:
-                UUID(value)
-            except ValueError as e:
-                raise ValueError(
-                    "Invalid value passed as UUID query parameter."
-                ) from e
-
-        # For equality checks, ensure that the value is a valid UUID.
         if operator == GenericFilterOps.ONEOF and not isinstance(value, list):
             raise ValueError(ONEOF_ERROR)