zenml-nightly 0.68.1.dev20241105__py3-none-any.whl → 0.68.1.dev20241107__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.68.1.dev20241105
+0.68.1.dev20241107

zenml/artifacts/{load_directory_materializer.py → preexisting_data_materializer.py}

@@ -14,7 +14,6 @@
 """Only-load materializer for directories."""
 
 import os
-import tempfile
 from pathlib import Path
 from typing import Any, ClassVar, Tuple, Type
 
@@ -46,14 +45,14 @@ class PreexistingDataMaterializer(BaseMaterializer):
         Returns:
             Path to the local directory that contains the artifact files.
         """
-        directory = tempfile.mkdtemp(prefix="zenml-artifact")
-        if fileio.isdir(self.uri):
-            self._copy_directory(src=self.uri, dst=directory)
-            return Path(directory)
-        else:
-            dst = os.path.join(directory, os.path.split(self.uri)[-1])
-            fileio.copy(src=self.uri, dst=dst)
-            return Path(dst)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            if fileio.isdir(self.uri):
+                self._copy_directory(src=self.uri, dst=temp_dir)
+                return Path(temp_dir)
+            else:
+                dst = os.path.join(temp_dir, os.path.split(self.uri)[-1])
+                fileio.copy(src=self.uri, dst=dst)
+                return Path(dst)
 
     def save(self, data: Any) -> None:
         """Store the directory in the artifact store.

zenml/artifacts/utils.py

@@ -31,7 +31,7 @@ from typing import (
 from uuid import UUID, uuid4
 
 from zenml.artifacts.artifact_config import ArtifactConfig
-from zenml.artifacts.load_directory_materializer import (
+from zenml.artifacts.preexisting_data_materializer import (
     PreexistingDataMaterializer,
 )
 from zenml.client import Client
@@ -82,6 +82,114 @@ logger = get_logger(__name__)
 # ----------
 
 
+def _save_artifact_visualizations(
+    data: Any, materializer: "BaseMaterializer"
+) -> List[ArtifactVisualizationRequest]:
+    """Save artifact visualizations.
+
+    Args:
+        data: The data for which to save the visualizations.
+        materializer: The materializer that should be used to generate and
+            save the visualizations.
+
+    Returns:
+        List of requests for the saved visualizations.
+    """
+    try:
+        visualizations = materializer.save_visualizations(data)
+    except Exception as e:
+        logger.warning("Failed to save artifact visualizations: %s", e)
+        return []
+
+    return [
+        ArtifactVisualizationRequest(
+            type=type,
+            uri=uri,
+        )
+        for uri, type in visualizations.items()
+    ]
+
+
+def _store_artifact_data_and_prepare_request(
+    data: Any,
+    name: str,
+    uri: str,
+    materializer_class: Type["BaseMaterializer"],
+    version: Optional[Union[int, str]] = None,
+    tags: Optional[List[str]] = None,
+    store_metadata: bool = True,
+    store_visualizations: bool = True,
+    has_custom_name: bool = True,
+    metadata: Optional[Dict[str, "MetadataType"]] = None,
+) -> ArtifactVersionRequest:
+    """Store artifact data and prepare a request to the server.
+
+    Args:
+        data: The artifact data.
+        name: The artifact name.
+        uri: The artifact URI.
+        materializer_class: The materializer class to use for storing the
+            artifact data.
+        version: The artifact version.
+        tags: Tags for the artifact version.
+        store_metadata: Whether to store metadata for the artifact version.
+        store_visualizations: Whether to store visualizations for the artifact
+            version.
+        has_custom_name: Whether the artifact has a custom name.
+        metadata: Metadata to store for the artifact version. This will be
+            ignored if `store_metadata` is set to `False`.
+
+    Returns:
+        Artifact version request for the artifact data that was stored.
+    """
+    artifact_store = Client().active_stack.artifact_store
+    artifact_store.makedirs(uri)
+
+    materializer = materializer_class(uri=uri, artifact_store=artifact_store)
+    materializer.uri = materializer.uri.replace("\\", "/")
+
+    data_type = type(data)
+    materializer.validate_save_type_compatibility(data_type)
+    materializer.save(data)
+
+    visualizations = (
+        _save_artifact_visualizations(data=data, materializer=materializer)
+        if store_visualizations
+        else None
+    )
+
+    combined_metadata: Dict[str, "MetadataType"] = {}
+    if store_metadata:
+        try:
+            combined_metadata = materializer.extract_full_metadata(data)
+        except Exception as e:
+            logger.warning("Failed to extract materializer metadata: %s", e)
+
+        # Update with user metadata to potentially overwrite values coming from
+        # the materializer
+        combined_metadata.update(metadata or {})
+
+    artifact_version_request = ArtifactVersionRequest(
+        artifact_name=name,
+        version=version,
+        tags=tags,
+        type=materializer.ASSOCIATED_ARTIFACT_TYPE,
+        uri=materializer.uri,
+        materializer=source_utils.resolve(materializer.__class__),
+        data_type=source_utils.resolve(data_type),
+        user=Client().active_user.id,
+        workspace=Client().active_workspace.id,
+        artifact_store_id=artifact_store.id,
+        visualizations=visualizations,
+        has_custom_name=has_custom_name,
+        metadata=validate_metadata(combined_metadata)
+        if combined_metadata
+        else None,
+    )
+
+    return artifact_version_request
+
+
 def save_artifact(
     data: Any,
     name: str,
@@ -89,13 +197,14 @@ def save_artifact(
     tags: Optional[List[str]] = None,
     extract_metadata: bool = True,
     include_visualizations: bool = True,
-    has_custom_name: bool = True,
     user_metadata: Optional[Dict[str, "MetadataType"]] = None,
     materializer: Optional["MaterializerClassOrSource"] = None,
     uri: Optional[str] = None,
     is_model_artifact: bool = False,
     is_deployment_artifact: bool = False,
+    # TODO: remove these once external artifact does not use this function anymore
     manual_save: bool = True,
+    has_custom_name: bool = True,
 ) -> "ArtifactVersionResponse":
     """Upload and publish an artifact.
 
@@ -107,8 +216,6 @@ def save_artifact(
         tags: Tags to associate with the artifact.
         extract_metadata: If artifact metadata should be extracted and returned.
         include_visualizations: If artifact visualizations should be generated.
-        has_custom_name: If the artifact name is custom and should be listed in
-            the dashboard "Artifacts" tab.
         user_metadata: User-provided metadata to store with the artifact.
         materializer: The materializer to use for saving the artifact to the
             artifact store.
@@ -119,6 +226,8 @@ def save_artifact(
         is_deployment_artifact: If the artifact is a deployment artifact.
         manual_save: If this function is called manually and should therefore
             link the artifact to the current step run.
+        has_custom_name: If the artifact name is custom and should be listed in
+            the dashboard "Artifacts" tab.
 
     Returns:
         The saved artifact response.
@@ -129,11 +238,8 @@ def save_artifact(
     from zenml.utils import source_utils
 
     client = Client()
-
-    # Get the current artifact store
     artifact_store = client.active_stack.artifact_store
 
-    # Build and check the artifact URI
     if not uri:
         uri = os.path.join("custom_artifacts", name, str(uuid4()))
     if not uri.startswith(artifact_store.path):
@@ -147,9 +253,7 @@ def save_artifact(
             uri=uri,
             name=name,
         )
-    artifact_store.makedirs(uri)
 
-    # Find and initialize the right materializer class
     if isinstance(materializer, type):
         materializer_class = materializer
     elif materializer:
@@ -158,60 +262,18 @@ def save_artifact(
         )
     else:
         materializer_class = materializer_registry[type(data)]
-    materializer_object = materializer_class(uri)
-
-    # Force URIs to have forward slashes
-    materializer_object.uri = materializer_object.uri.replace("\\", "/")
-
-    # Save the artifact to the artifact store
-    data_type = type(data)
-    materializer_object.validate_save_type_compatibility(data_type)
-    materializer_object.save(data)
 
-    # Save visualizations of the artifact
-    visualizations: List[ArtifactVisualizationRequest] = []
-    if include_visualizations:
-        try:
-            vis_data = materializer_object.save_visualizations(data)
-            for vis_uri, vis_type in vis_data.items():
-                vis_model = ArtifactVisualizationRequest(
-                    type=vis_type,
-                    uri=vis_uri,
-                )
-                visualizations.append(vis_model)
-        except Exception as e:
-            logger.warning(
-                f"Failed to save visualization for output artifact '{name}': "
-                f"{e}"
-            )
-
-    # Save metadata of the artifact
-    artifact_metadata: Dict[str, "MetadataType"] = {}
-    if extract_metadata:
-        try:
-            artifact_metadata = materializer_object.extract_full_metadata(data)
-            artifact_metadata.update(user_metadata or {})
-        except Exception as e:
-            logger.warning(
-                f"Failed to extract metadata for output artifact '{name}': {e}"
-            )
-
-    artifact_version_request = ArtifactVersionRequest(
-        artifact_name=name,
+    artifact_version_request = _store_artifact_data_and_prepare_request(
+        data=data,
+        name=name,
+        uri=uri,
+        materializer_class=materializer_class,
         version=version,
         tags=tags,
-        type=materializer_object.ASSOCIATED_ARTIFACT_TYPE,
-        uri=materializer_object.uri,
-        materializer=source_utils.resolve(materializer_object.__class__),
-        data_type=source_utils.resolve(data_type),
-        user=Client().active_user.id,
-        workspace=Client().active_workspace.id,
-        artifact_store_id=artifact_store.id,
-        visualizations=visualizations,
+        store_metadata=extract_metadata,
+        store_visualizations=include_visualizations,
         has_custom_name=has_custom_name,
-        metadata=validate_metadata(artifact_metadata)
-        if artifact_metadata
-        else None,
+        metadata=user_metadata,
     )
     artifact_version = client.zen_store.create_artifact_version(
         artifact_version=artifact_version_request

zenml/constants.py

@@ -338,6 +338,7 @@ ARTIFACTS = "/artifacts"
 ARTIFACT_VERSIONS = "/artifact_versions"
 ARTIFACT_VISUALIZATIONS = "/artifact_visualizations"
 AUTH = "/auth"
+BATCH = "/batch"
 CODE_REFERENCES = "/code_references"
 CODE_REPOSITORIES = "/code_repositories"
 COMPONENT_TYPES = "/component-types"

zenml/integrations/bentoml/materializers/bentoml_bento_materializer.py

@@ -14,7 +14,6 @@
 """Materializer for BentoML Bento objects."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import bentoml
@@ -23,7 +22,6 @@ from bentoml.exceptions import BentoMLException
 
 from zenml.enums import ArtifactType
 from zenml.integrations.bentoml.constants import DEFAULT_BENTO_FILENAME
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
@@ -49,23 +47,21 @@ class BentoMaterializer(BaseMaterializer):
         Returns:
             An bento.Bento object.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the Bento from the temporary directory
-        imported_bento = Bento.import_from(
-            os.path.join(temp_dir.name, DEFAULT_BENTO_FILENAME)
-        )
-
-        # Try save the Bento to the local BentoML store
-        try:
-            _ = bentoml.get(imported_bento.tag)
-        except BentoMLException:
-            imported_bento.save()
-        return imported_bento
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            # Copy from artifact store to temporary directory
+            io_utils.copy_dir(self.uri, temp_dir)
+
+            # Load the Bento from the temporary directory
+            imported_bento = Bento.import_from(
+                os.path.join(temp_dir, DEFAULT_BENTO_FILENAME)
+            )
+
+            # Try save the Bento to the local BentoML store
+            try:
+                _ = bentoml.get(imported_bento.tag)
+            except BentoMLException:
+                imported_bento.save()
+            return imported_bento
 
     def save(self, bento: bento.Bento) -> None:
         """Write to artifact store.
@@ -73,18 +69,10 @@ class BentoMaterializer(BaseMaterializer):
         Args:
             bento: An bento.Bento object.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        temp_bento_path = os.path.join(temp_dir.name, DEFAULT_BENTO_FILENAME)
-
-        # save the image in a temporary directory
-        bentoml.export_bento(bento.tag, temp_bento_path)
-
-        # copy the saved image to the artifact store
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_bento_path = os.path.join(temp_dir, DEFAULT_BENTO_FILENAME)
+            bentoml.export_bento(bento.tag, temp_bento_path)
+            io_utils.copy_dir(temp_dir, self.uri)
 
     def extract_metadata(
         self, bento: bento.Bento

zenml/integrations/evidently/__init__.py

@@ -33,7 +33,7 @@ from zenml.stack import Flavor
 
 # Fix numba errors in Docker and suppress logs and deprecation warning spam
 try:
-    from numba.core.errors import (  # type: ignore[import-not-found]
+    from numba.core.errors import (
         NumbaDeprecationWarning,
         NumbaPendingDeprecationWarning,
     )

zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py

@@ -15,7 +15,6 @@
 
 import os
 from collections import defaultdict
-from tempfile import TemporaryDirectory, mkdtemp
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -88,12 +87,12 @@ class HFDatasetMaterializer(BaseMaterializer):
         Returns:
             The dataset read from the specified dir.
         """
-        temp_dir = mkdtemp()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_DATASET_DIR),
-            temp_dir,
-        )
-        return load_from_disk(temp_dir)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_DATASET_DIR),
+                temp_dir,
+            )
+            return load_from_disk(temp_dir)
 
     def save(self, ds: Union[Dataset, DatasetDict]) -> None:
         """Writes a Dataset to the specified dir.
@@ -101,16 +100,13 @@ class HFDatasetMaterializer(BaseMaterializer):
         Args:
             ds: The Dataset to write.
         """
-        temp_dir = TemporaryDirectory()
-        path = os.path.join(temp_dir.name, DEFAULT_DATASET_DIR)
-        try:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            path = os.path.join(temp_dir, DEFAULT_DATASET_DIR)
             ds.save_to_disk(path)
             io_utils.copy_dir(
                 path,
                 os.path.join(self.uri, DEFAULT_DATASET_DIR),
             )
-        finally:
-            fileio.rmtree(temp_dir.name)
 
     def extract_metadata(
         self, ds: Union[Dataset, DatasetDict]

zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFPTModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: PreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFPTModelMaterializer(BaseMaterializer):
         Args:
             model: The Torch Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: PreTrainedModel

zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface t5 materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Type, Union
 
 from transformers import (
@@ -52,8 +51,7 @@ class HFT5Materializer(BaseMaterializer):
             ValueError: Unsupported data type used
         """
         filepath = self.uri
-
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Copy files from artifact store to temporary directory
             for file in fileio.listdir(filepath):
                 src = os.path.join(filepath, file)
@@ -86,8 +84,7 @@ class HFT5Materializer(BaseMaterializer):
         Args:
             obj: A T5ForConditionalGeneration model or T5Tokenizer.
         """
-        # Create a temporary directory
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Save the model or tokenizer
             obj.save_pretrained(temp_dir)
 

zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFTFModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = "TF" + config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = "TF" + config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: TFPreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFTFModelMaterializer(BaseMaterializer):
         Args:
             model: The TF Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: TFPreTrainedModel

zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface tokenizer materializer."""
 
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Tuple, Type
 
 from transformers import AutoTokenizer
@@ -46,7 +45,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Returns:
             The tokenizer read from the specified dir.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             io_utils.copy_dir(
                 os.path.join(self.uri, DEFAULT_TOKENIZER_DIR), temp_dir
             )
@@ -58,7 +57,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Args:
             tokenizer: The HFTokenizer to write.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             tokenizer.save_pretrained(temp_dir)
             io_utils.copy_dir(
                 temp_dir,

zenml/integrations/lightgbm/materializers/lightgbm_booster_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM booster materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type
 
 import lightgbm as lgb
@@ -42,18 +41,13 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             A lightgbm Booster object.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        booster = lgb.Booster(model_file=temp_file)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return booster
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            booster = lgb.Booster(model_file=temp_file)
+            return booster
 
     def save(self, booster: lgb.Booster) -> None:
         """Creates a JSON serialization for a lightgbm Booster model.
@@ -62,8 +56,7 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             booster: A lightgbm Booster model.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_path = os.path.join(tmp_dir, "model.txt")
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            tmp_path = os.path.join(temp_dir, "model.txt")
             booster.save_model(tmp_path)
             fileio.copy(tmp_path, filepath)