zenml-nightly 0.58.2.dev20240615__py3-none-any.whl → 0.58.2.dev20240623__py3-none-any.whl

zenml/integrations/tensorboard/__init__.py

@@ -35,18 +35,7 @@ class TensorBoardIntegration(Integration):
         Returns:
             A list of requirements.
         """
-        if sys.version_info > (3, 11):
-            tf_version = "2.13"
-        else:
-            # Capping tensorflow to 2.11 for Python 3.10 and below because it
-            # is not compatible with Pytorch
-            # (see https://github.com/pytorch/pytorch/issues/99637).
-            tf_version = "2.11"
-
-        requirements = [
-            f"tensorboard=={tf_version}",
-            "protobuf>=3.6.0,<4.0.0",
-        ]
+        requirements = ["tensorboard>=2.12,<=2.15"]
         return requirements
 
     @classmethod

zenml/integrations/tensorboard/services/tensorboard_service.py

@@ -17,9 +17,6 @@ import uuid
 from typing import Any, Dict, Union
 
 from tensorboard import default, program  # type: ignore [import-untyped]
-from tensorboard.uploader import (  # type: ignore [import-untyped]
-    uploader_subcommand,
-)
 
 from zenml.logger import get_logger
 from zenml.services import (
@@ -104,7 +101,9 @@ class TensorboardService(LocalDaemonService):
                 ),
             )
             attrs["endpoint"] = endpoint
-        super().__init__(config=config, uuid=uuid.uuid4(), **attrs)
+        if "uuid" not in attrs:
+            attrs["uuid"] = uuid.uuid4()
+        super().__init__(config=config, **attrs)
 
     def run(self) -> None:
         """Initialize and run the TensorBoard server."""
@@ -118,7 +117,6 @@ class TensorboardService(LocalDaemonService):
         try:
             tensorboard = program.TensorBoard(
                 plugins=default.get_plugins(),
-                subcommands=[uploader_subcommand.UploaderSubcommand()],
             )
             tensorboard.configure(
                 logdir=self.config.logdir,

zenml/integrations/tensorboard/visualizers/tensorboard_visualizer.py

@@ -194,19 +194,19 @@ def get_step(pipeline_name: str, step_name: str) -> "StepRunResponse":
     Raises:
         RuntimeError: If the step is not found.
     """
-    pipeline = Client().get_pipeline(pipeline_name)
-    if pipeline is None:
+    runs = Client().list_pipeline_runs(pipeline_name=pipeline_name)
+    if runs.total == 0:
         raise RuntimeError(
-            f"No pipeline with name `{pipeline_name}` was found"
+            f"No pipeline runs for pipeline `{pipeline_name}` were found"
         )
 
-    last_run = pipeline.runs[0]
-    step = last_run.steps[step_name]
-    if step is None:
+    last_run = runs[0]
+    if step_name not in last_run.steps:
         raise RuntimeError(
             f"No pipeline step with name `{step_name}` was found in "
             f"pipeline `{pipeline_name}`"
         )
+    step = last_run.steps[step_name]
     return step
 
 

zenml/integrations/tensorflow/__init__.py

@@ -49,23 +49,15 @@ class TensorflowIntegration(Integration):
         Returns:
             A list of requirements.
         """
-        if sys.version_info > (3, 11):
-            tf_version = "2.13"
-        else:
-            # Capping tensorflow to 2.11 for Python 3.10 and below because it
-            # is not compatible with Pytorch
-            # (see https://github.com/pytorch/pytorch/issues/99637).
-            tf_version = "2.11"
         target_os = target_os or platform.system()
         if target_os == "Darwin" and platform.machine() == "arm64":
             requirements = [
-                f"tensorflow-macos=={tf_version}",
+                f"tensorflow-macos>=2.12,<=2.15",
             ]
         else:
             requirements = [
-                f"tensorflow=={tf_version}",
+                f"tensorflow>=2.12,<=2.15",
                 "tensorflow_io>=0.24.0",
-                "protobuf>=3.6.0,<4.0.0",
             ]
         return requirements
 

zenml/integrations/tensorflow/materializers/keras_materializer.py

@@ -13,10 +13,12 @@
 #  permissions and limitations under the License.
 """Implementation of the TensorFlow Keras materializer."""
 
+import os
 import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
-from tensorflow import keras
+import tensorflow as tf
+from tensorflow.python import keras as tf_keras
 from tensorflow.python.keras.utils.layer_utils import count_params
 
 from zenml.enums import ArtifactType
@@ -31,48 +33,54 @@ if TYPE_CHECKING:
 class KerasMaterializer(BaseMaterializer):
     """Materializer to read/write Keras models."""
 
-    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (keras.Model,)
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (
+        tf.keras.Model,
+        tf_keras.Model,
+    )
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
+    MODEL_FILE_NAME = "model.keras"
 
-    def load(self, data_type: Type[Any]) -> keras.Model:
+    def load(self, data_type: Type[Any]) -> tf_keras.Model:
         """Reads and returns a Keras model after copying it to temporary path.
 
         Args:
             data_type: The type of the data to read.
 
         Returns:
-            A tf.keras.Model model.
+            A keras.Model model.
         """
         # Create a temporary directory to store the model
         temp_dir = tempfile.TemporaryDirectory()
 
         # Copy from artifact store to temporary directory
+        temp_model_file = os.path.join(temp_dir.name, self.MODEL_FILE_NAME)
         io_utils.copy_dir(self.uri, temp_dir.name)
 
         # Load the model from the temporary directory
-        model = keras.models.load_model(temp_dir.name)
+        model = tf.keras.models.load_model(temp_model_file)
 
         # Cleanup and return
         fileio.rmtree(temp_dir.name)
 
         return model
 
-    def save(self, model: keras.Model) -> None:
+    def save(self, model: tf_keras.Model) -> None:
         """Writes a keras model to the artifact store.
 
         Args:
-            model: A tf.keras.Model model.
+            model: A keras.Model model.
         """
         # Create a temporary directory to store the model
         temp_dir = tempfile.TemporaryDirectory()
-        model.save(temp_dir.name)
+        temp_model_file = os.path.join(temp_dir.name, self.MODEL_FILE_NAME)
+        model.save(temp_model_file)
         io_utils.copy_dir(temp_dir.name, self.uri)
 
         # Remove the temporary directory
         fileio.rmtree(temp_dir.name)
 
     def extract_metadata(
-        self, model: keras.Model
+        self, model: tf_keras.Model
     ) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given `Model` object.
 

zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py

@@ -20,11 +20,10 @@ from typing import (
     List,
     Optional,
     Type,
-    Union,
     cast,
 )
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -35,8 +34,6 @@ from zenml.integrations.wandb import WANDB_EXPERIMENT_TRACKER_FLAVOR
 from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
-    from wandb import Settings
-
     from zenml.integrations.wandb.experiment_trackers import (
         WandbExperimentTracker,
     )
@@ -55,11 +52,9 @@ class WandbExperimentTrackerSettings(BaseSettings):
     tags: List[str] = []
     settings: Dict[str, Any] = {}
 
-    @validator("settings", pre=True)
-    def _convert_settings(
-        cls,
-        value: Union[Dict[str, Any], "Settings"],
-    ) -> Dict[str, Any]:
+    @field_validator("settings", mode="before")
+    @classmethod
+    def _convert_settings(cls, value: Any) -> Any:
         """Converts settings to a dictionary.
 
         Args:
@@ -81,7 +76,7 @@ class WandbExperimentTrackerSettings(BaseSettings):
             return value
 
 
-class WandbExperimentTrackerConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class WandbExperimentTrackerConfig(
     BaseExperimentTrackerConfig, WandbExperimentTrackerSettings
 ):
     """Config for the Wandb experiment tracker.
@@ -112,7 +107,7 @@ class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
 
     @property
     def docs_url(self) -> Optional[str]:
-        """A url to point at docs explaining this flavor.
+        """A URL to point at docs explaining this flavor.
 
         Returns:
             A flavor docs url.
@@ -121,7 +116,7 @@ class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
 
     @property
     def sdk_docs_url(self) -> Optional[str]:
-        """A url to point at SDK docs explaining this flavor.
+        """A URL to point at SDK docs explaining this flavor.
 
         Returns:
             A flavor SDK docs url.
@@ -130,7 +125,7 @@ class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
 
     @property
     def logo_url(self) -> str:
-        """A url to represent the flavor in the dashboard.
+        """A URL to represent the flavor in the dashboard.
 
         Returns:
             The flavor logo.
@@ -142,7 +137,7 @@ class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
         """Returns `WandbExperimentTrackerConfig` config class.
 
         Returns:
-                The config class.
+            The config class.
         """
         return WandbExperimentTrackerConfig
 

zenml/integrations/whylogs/flavors/whylogs_data_validator_flavor.py

@@ -41,7 +41,7 @@ class WhylogsDataValidatorSettings(BaseSettings):
     dataset_id: Optional[str] = None
 
 
-class WhylogsDataValidatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class WhylogsDataValidatorConfig(
     BaseDataValidatorConfig,
     AuthenticationConfigMixin,
     WhylogsDataValidatorSettings,

zenml/lineage_graph/lineage_graph.py

@@ -168,7 +168,7 @@ class LineageGraph(BaseModel):
             step: The step to add a node for.
             id: The id of the step node.
         """
-        step_config = step.config.dict()
+        step_config = step.config.model_dump()
         if step_config:
             step_config = {
                 key: value

zenml/logging/step_logging.py

@@ -98,7 +98,7 @@ def prepare_logs_uri(
         artifact_store.makedirs(logs_uri_folder)
         return logs_uri_folder
     else:
-        logs_uri = os.path.join(logs_base_uri, f"{log_key}.log")
+        logs_uri = os.path.join(logs_base_uri, f"{log_key}{LOGS_EXTENSION}")
         if artifact_store.exists(logs_uri):
             logger.warning(
                 f"Logs file {logs_uri} already exists! Removing old log file..."
@@ -238,7 +238,7 @@ class StepLogsStorage:
         # Immutable filesystems state
         self.last_merge_time = time.time()
         self.log_files_not_merged: List[str] = []
-        self.next_merged_file_name: str = f"{time.time()}{LOGS_EXTENSION}"
+        self.next_merged_file_name: str = self._get_timestamped_filename()
 
     @property
     def artifact_store(self) -> "BaseArtifactStore":
@@ -276,6 +276,14 @@ class StepLogsStorage:
             or time.time() - self.last_save_time >= self.time_interval
         )
 
+    def _get_timestamped_filename(self) -> str:
+        """Returns a timestamped filename.
+
+        Returns:
+            The timestamped filename.
+        """
+        return f"{time.time()}{LOGS_EXTENSION}"
+
     def save_to_file(self, force: bool = False) -> None:
         """Method to save the buffer to the given URI.
 
@@ -293,14 +301,14 @@ class StepLogsStorage:
                     if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
                         if not self.log_files_not_merged:
                             self.next_merged_file_name = (
-                                f"{time.time()}{LOGS_EXTENSION}"
+                                self._get_timestamped_filename()
                             )
-                        log_file_ = f"{time.time()}{LOGS_EXTENSION}"
-                        self.log_files_not_merged.append(log_file_)
+                        _logs_uri = self._get_timestamped_filename()
+                        self.log_files_not_merged.append(_logs_uri)
                         with self.artifact_store.open(
                             os.path.join(
                                 self.logs_uri,
-                                log_file_,
+                                _logs_uri,
                             ),
                             "w",
                         ) as file:
@@ -360,7 +368,7 @@ class StepLogsStorage:
         """
         if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
             files_ = files or self.artifact_store.listdir(self.logs_uri)
-            file_name_ = file_name or f"full_log{LOGS_EXTENSION}"
+            file_name_ = file_name or self._get_timestamped_filename()
             if len(files_) > 1:
                 files_.sort()
                 logger.debug("Log files count: %s", len(files_))

zenml/materializers/built_in_materializer.py

@@ -339,11 +339,11 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # Cast the data to the correct type.
         if issubclass(data_type, dict) and not isinstance(outputs, dict):
             keys, values = outputs
-            return dict(zip(keys, values))
+            return data_type(zip(keys, values))
         if issubclass(data_type, tuple) and not isinstance(outputs, tuple):
-            return tuple(outputs)
+            return data_type(outputs)
         if issubclass(data_type, set) and not isinstance(outputs, set):
-            return set(outputs)
+            return data_type(outputs)
         return outputs
 
     def save(self, data: Any) -> None:

zenml/materializers/pydantic_materializer.py

@@ -45,7 +45,7 @@ class PydanticMaterializer(BaseMaterializer):
         """
         data_path = os.path.join(self.uri, DEFAULT_FILENAME)
         contents = yaml_utils.read_json(data_path)
-        return data_type.parse_raw(contents)
+        return data_type.model_validate_json(contents)
 
     def save(self, data: BaseModel) -> None:
         """Serialize a BaseModel to JSON.
@@ -54,7 +54,7 @@ class PydanticMaterializer(BaseMaterializer):
             data: The data to store.
         """
         data_path = os.path.join(self.uri, DEFAULT_FILENAME)
-        yaml_utils.write_json(data_path, data.json())
+        yaml_utils.write_json(data_path, data.model_dump_json())
 
     def extract_metadata(self, data: BaseModel) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given BaseModel object.

zenml/metadata/lazy_load.py

@@ -57,8 +57,8 @@ class RunMetadataLazyGetter:
         from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         return LazyRunMetadataResponse(
-            _lazy_load_model=self._lazy_load_model,
-            _lazy_load_artifact_name=self._lazy_load_artifact_name,
-            _lazy_load_artifact_version=self._lazy_load_artifact_version,
-            _lazy_load_metadata_name=key,
+            lazy_load_model=self._lazy_load_model,
+            lazy_load_artifact_name=self._lazy_load_artifact_name,
+            lazy_load_artifact_version=self._lazy_load_artifact_version,
+            lazy_load_metadata_name=key,
         )

zenml/metadata/metadata_types.py

@@ -13,29 +13,89 @@
 #  permissions and limitations under the License.
 """Custom types that can be used as metadata of ZenML artifacts."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union
+from typing import Any, Dict, List, Set, Tuple, Union
 
-from zenml.utils.enum_utils import StrEnum
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
 
-if TYPE_CHECKING:
-    pass
+from zenml.utils.enum_utils import StrEnum
 
 
 class Uri(str):
     """Special string class to indicate a URI."""
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Additional method for pydantic to recognize it as a valid type.
+
+        Args:
+            source_type: the source type
+            handler: the handler
+
+        Returns:
+            the schema for the custom type.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
 
 class Path(str):
     """Special string class to indicate a path."""
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Additional method for pydantic to recognize it as a valid type.
+
+        Args:
+            source_type: the source type
+            handler: the handler
+
+        Returns:
+            the schema for the custom type.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
 
 class DType(str):
     """Special string class to indicate a data type."""
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Additional method for pydantic to recognize it as a valid type.
+
+        Args:
+            source_type: the source type
+            handler: the handler
+
+        Returns:
+            the schema for the custom type.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
 
 class StorageSize(int):
     """Special int class to indicate the storage size in number of bytes."""
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """Additional method for pydantic to recognize it as a valid type.
+
+        Args:
+            source_type: the source type
+            handler: the handler
+
+        Returns:
+            the schema for the custom type.
+        """
+        return core_schema.no_info_after_validator_function(cls, handler(int))
+
 
 # Union of all types that can be used as metadata. Used for type annotations.
 MetadataType = Union[