zenml-nightly 0.68.1.dev20241103__py3-none-any.whl → 0.70.0.dev20241115__py3-none-any.whl

zenml/client.py

@@ -136,9 +136,7 @@ from zenml.models import (
     PipelineResponse,
     PipelineRunFilter,
     PipelineRunResponse,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -190,6 +188,7 @@ from zenml.models import (
     WorkspaceResponse,
     WorkspaceUpdate,
 )
+from zenml.models.v2.core.step_run import StepRunUpdate
 from zenml.services.service import ServiceConfig
 from zenml.services.service_status import ServiceState
 from zenml.services.service_type import ServiceType
@@ -3427,7 +3426,9 @@ class Client(metaclass=ClientMetaClass):
         Args:
             id_or_prefix: The id or id prefix of the deployment.
         """
-        deployment = self.get_deployment(id_or_prefix=id_or_prefix)
+        deployment = self.get_deployment(
+            id_or_prefix=id_or_prefix, hydrate=False
+        )
         self.zen_store.delete_deployment(deployment_id=deployment.id)
 
     # ------------------------------ Run templates -----------------------------
@@ -3795,6 +3796,7 @@ class Client(metaclass=ClientMetaClass):
         templatable: Optional[bool] = None,
         tag: Optional[str] = None,
         user: Optional[Union[UUID, str]] = None,
+        run_metadata: Optional[Dict[str, str]] = None,
         pipeline: Optional[Union[UUID, str]] = None,
         code_repository: Optional[Union[UUID, str]] = None,
         model: Optional[Union[UUID, str]] = None,
@@ -3834,6 +3836,7 @@ class Client(metaclass=ClientMetaClass):
             templatable: If the runs should be templatable or not.
             tag: Tag to filter by.
             user: The name/ID of the user to filter by.
+            run_metadata: The run_metadata of the run to filter by.
             pipeline: The name/ID of the pipeline to filter by.
             code_repository: Filter by code repository name/ID.
             model: Filter by model name/ID.
@@ -3873,6 +3876,7 @@ class Client(metaclass=ClientMetaClass):
             tag=tag,
             unlisted=unlisted,
             user=user,
+            run_metadata=run_metadata,
             pipeline=pipeline,
             code_repository=code_repository,
             stack=stack,
@@ -4166,6 +4170,8 @@ class Client(metaclass=ClientMetaClass):
         Returns:
             The artifact version.
         """
+        from zenml import get_step_context
+
         if cll := client_lazy_loader(
             method_name="get_artifact_version",
             name_id_or_prefix=name_id_or_prefix,
@@ -4173,13 +4179,26 @@ class Client(metaclass=ClientMetaClass):
             hydrate=hydrate,
         ):
             return cll  # type: ignore[return-value]
-        return self._get_entity_version_by_id_or_name_or_prefix(
+
+        artifact = self._get_entity_version_by_id_or_name_or_prefix(
             get_method=self.zen_store.get_artifact_version,
             list_method=self.list_artifact_versions,
             name_id_or_prefix=name_id_or_prefix,
             version=version,
             hydrate=hydrate,
         )
+        try:
+            step_run = get_step_context().step_run
+            client = Client()
+            client.zen_store.update_run_step(
+                step_run_id=step_run.id,
+                step_run_update=StepRunUpdate(
+                    loaded_artifact_versions={artifact.name: artifact.id}
+                ),
+            )
+        except RuntimeError:
+            pass  # Cannot link to step run if called outside a step
+        return artifact
 
     def list_artifact_versions(
         self,
@@ -4206,6 +4225,7 @@ class Client(metaclass=ClientMetaClass):
         user: Optional[Union[UUID, str]] = None,
         model: Optional[Union[UUID, str]] = None,
         pipeline_run: Optional[Union[UUID, str]] = None,
+        run_metadata: Optional[Dict[str, str]] = None,
         tag: Optional[str] = None,
         hydrate: bool = False,
     ) -> Page[ArtifactVersionResponse]:
@@ -4237,6 +4257,7 @@ class Client(metaclass=ClientMetaClass):
             user: Filter by user name or ID.
             model: Filter by model name or ID.
             pipeline_run: Filter by pipeline run name or ID.
+            run_metadata: Filter by run metadata.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -4417,7 +4438,7 @@ class Client(metaclass=ClientMetaClass):
         resource_id: UUID,
         resource_type: MetadataResourceTypes,
         stack_component_id: Optional[UUID] = None,
-    ) -> List[RunMetadataResponse]:
+    ) -> None:
         """Create run metadata.
 
         Args:
@@ -4430,7 +4451,7 @@ class Client(metaclass=ClientMetaClass):
                 the metadata.
 
         Returns:
-            The created metadata, as string to model dictionary.
+            None
         """
         from zenml.metadata.metadata_types import get_metadata_type
 
@@ -4465,74 +4486,8 @@ class Client(metaclass=ClientMetaClass):
             values=values,
             types=types,
         )
-        return self.zen_store.create_run_metadata(run_metadata)
-
-    def list_run_metadata(
-        self,
-        sort_by: str = "created",
-        page: int = PAGINATION_STARTING_PAGE,
-        size: int = PAGE_SIZE_DEFAULT,
-        logical_operator: LogicalOperators = LogicalOperators.AND,
-        id: Optional[Union[UUID, str]] = None,
-        created: Optional[Union[datetime, str]] = None,
-        updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[UUID] = None,
-        user_id: Optional[UUID] = None,
-        resource_id: Optional[UUID] = None,
-        resource_type: Optional[MetadataResourceTypes] = None,
-        stack_component_id: Optional[UUID] = None,
-        key: Optional[str] = None,
-        value: Optional["MetadataType"] = None,
-        type: Optional[str] = None,
-        hydrate: bool = False,
-    ) -> Page[RunMetadataResponse]:
-        """List run metadata.
-
-        Args:
-            sort_by: The field to sort the results by.
-            page: The page number to return.
-            size: The number of results to return per page.
-            logical_operator: The logical operator to use for filtering.
-            id: The ID of the metadata.
-            created: The creation time of the metadata.
-            updated: The last update time of the metadata.
-            workspace_id: The ID of the workspace the metadata belongs to.
-            user_id: The ID of the user that created the metadata.
-            resource_id: The ID of the resource the metadata belongs to.
-            resource_type: The type of the resource the metadata belongs to.
-            stack_component_id: The ID of the stack component that produced
-                the metadata.
-            key: The key of the metadata.
-            value: The value of the metadata.
-            type: The type of the metadata.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        metadata_filter_model = RunMetadataFilter(
-            sort_by=sort_by,
-            page=page,
-            size=size,
-            logical_operator=logical_operator,
-            id=id,
-            created=created,
-            updated=updated,
-            workspace_id=workspace_id,
-            user_id=user_id,
-            resource_id=resource_id,
-            resource_type=resource_type,
-            stack_component_id=stack_component_id,
-            key=key,
-            value=value,
-            type=type,
-        )
-        metadata_filter_model.set_scope_workspace(self.active_workspace.id)
-        return self.zen_store.list_run_metadata(
-            metadata_filter_model,
-            hydrate=hydrate,
-        )
+        self.zen_store.create_run_metadata(run_metadata)
+        return None
 
     # -------------------------------- Secrets ---------------------------------
 

zenml/config/server_config.py

@@ -19,13 +19,14 @@ from secrets import token_hex
 from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator
 
 from zenml.constants import (
     DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
     DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
@@ -119,6 +120,8 @@ class ServerConfiguration(BaseModel):
             time of the JWT tokens issued to clients after they have
             authenticated with the ZenML server using an OAuth 2.0 device
             that has been marked as trusted.
+        generic_api_token_lifetime: The lifetime in seconds that generic
+            short-lived API tokens issued for automation purposes are valid.
         external_login_url: The login URL of an external authenticator service
             to use with the `EXTERNAL` authentication scheme.
         external_user_info_url: The user info URL of an external authenticator
@@ -230,6 +233,12 @@ class ServerConfiguration(BaseModel):
             deployment.
         max_request_body_size_in_bytes: The maximum size of the request body in
             bytes. If not specified, the default value of 256 Kb will be used.
+        memcache_max_capacity: The maximum number of entries that the memory
+            cache can hold. If not specified, the default value of 1000 will be
+            used.
+        memcache_default_expiry: The default expiry time in seconds for cache
+            entries. If not specified, the default value of 30 seconds will be
+            used.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -257,6 +266,10 @@ class ServerConfiguration(BaseModel):
     device_expiration_minutes: Optional[int] = None
     trusted_device_expiration_minutes: Optional[int] = None
 
+    generic_api_token_lifetime: PositiveInt = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME
+    )
+
     external_login_url: Optional[str] = None
     external_user_info_url: Optional[str] = None
     external_cookie_name: Optional[str] = None
@@ -321,6 +334,9 @@ class ServerConfiguration(BaseModel):
         DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES
     )
 
+    memcache_max_capacity: int = 1000
+    memcache_default_expiry: int = 30
+
     _deployment_id: Optional[UUID] = None
 
     @model_validator(mode="before")

zenml/constants.py

@@ -167,9 +167,6 @@ ENV_ZENML_SERVER = "ZENML_SERVER"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
-ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES = (
-    "ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES"
-)
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
@@ -270,6 +267,7 @@ ZENML_API_KEY_PREFIX = "ZENKEY_"
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY = 1000
+DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME = 60 * 60  # 1 hour
 
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS = (
     "max-age=63072000; includeSubdomains"
@@ -338,6 +336,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"
@@ -409,10 +408,6 @@ MODEL_METADATA_YAML_FILE_NAME = "model_metadata.yaml"
 
 # orchestrator constants
 ORCHESTRATOR_DOCKER_IMAGE_KEY = "orchestrator"
-PIPELINE_API_TOKEN_EXPIRES_MINUTES = handle_int_env_var(
-    ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES,
-    default=60 * 24,  # 24 hours
-)
 
 # Secret constants
 SECRET_VALUES = "values"

zenml/data_validators/base_data_validator.py

@@ -175,7 +175,7 @@ class BaseDataValidator(StackComponent):
 
         This method should be implemented by data validators that support
         running model validation checks (e.g. confusion matrix validation,
-        performance reports, model error analyses, etc).
+        performance reports, model error analyzes, etc).
 
         Unlike `data_validation`, model validation checks require that a model
         be present as an active component during the validation process.
@@ -184,7 +184,7 @@ class BaseDataValidator(StackComponent):
         accommodate different categories of data validation tests, e.g.:
 
         * single dataset tests: confusion matrix validation,
-        performance reports, model error analyses, etc
+        performance reports, model error analyzes, etc
         * model comparison tests: tests that identify changes in a model
         behavior by comparing how it performs on two different datasets.
 

zenml/enums.py

@@ -34,15 +34,23 @@ class ArtifactType(StrEnum):
 class StepRunInputArtifactType(StrEnum):
     """All possible types of a step run input artifact."""
 
-    DEFAULT = "default"  # input argument that is the output of a previous step
+    STEP_OUTPUT = (
+        "step_output"  # input argument that is the output of a previous step
+    )
     MANUAL = "manual"  # manually loaded via `zenml.load_artifact()`
+    EXTERNAL = "external"  # loaded via `ExternalArtifact(value=...)`
+    LAZY_LOADED = "lazy"  # loaded via various lazy methods
 
 
-class StepRunOutputArtifactType(StrEnum):
-    """All possible types of a step run output artifact."""
+class ArtifactSaveType(StrEnum):
+    """All possible method types of how artifact versions can be saved."""
 
-    DEFAULT = "default"  # output of the current step
+    STEP_OUTPUT = "step_output"  # output of the current step
     MANUAL = "manual"  # manually saved via `zenml.save_artifact()`
+    PREEXISTING = "preexisting"  # register via `zenml.register_artifact()`
+    EXTERNAL = (
+        "external"  # saved via `zenml.ExternalArtifact.upload_by_value()`
+    )
 
 
 class VisualizationType(StrEnum):
@@ -237,6 +245,13 @@ class OAuthDeviceStatus(StrEnum):
     LOCKED = "locked"
 
 
+class APITokenType(StrEnum):
+    """The API token type."""
+
+    GENERIC = "generic"
+    WORKLOAD = "workload"
+
+
 class GenericFilterOps(StrEnum):
     """Ops for all filters for string values on list methods."""
 
@@ -245,6 +260,7 @@ class GenericFilterOps(StrEnum):
     CONTAINS = "contains"
     STARTSWITH = "startswith"
     ENDSWITH = "endswith"
+    ONEOF = "oneof"
     GTE = "gte"
     GT = "gt"
     LTE = "lte"

zenml/exceptions.py

@@ -168,6 +168,10 @@ class EntityExistsError(ZenMLBaseException):
     """Raised when trying to register an entity that already exists."""
 
 
+class EntityCreationError(ZenMLBaseException, RuntimeError):
+    """Raised when failing to create an entity."""
+
+
 class ActionExistsError(EntityExistsError):
     """Raised when registering an action with a name that already exists."""
 

zenml/integrations/__init__.py

@@ -45,6 +45,7 @@ from zenml.integrations.kubernetes import KubernetesIntegration  # noqa
 from zenml.integrations.label_studio import LabelStudioIntegration  # noqa
 from zenml.integrations.langchain import LangchainIntegration  # noqa
 from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
+
 # from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
@@ -52,7 +53,7 @@ from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
 from zenml.integrations.numpy import NumpyIntegration  # noqa
 from zenml.integrations.openai import OpenAIIntegration  # noqa
 from zenml.integrations.pandas import PandasIntegration  # noqa
-from zenml.integrations.pigeon import PigeonIntegration # noqa
+from zenml.integrations.pigeon import PigeonIntegration  # noqa
 from zenml.integrations.pillow import PillowIntegration  # noqa
 from zenml.integrations.polars import PolarsIntegration  # noqa
 from zenml.integrations.prodigy import ProdigyIntegration  # noqa
@@ -78,3 +79,4 @@ from zenml.integrations.tensorflow import TensorflowIntegration  # noqa
 from zenml.integrations.wandb import WandbIntegration  # noqa
 from zenml.integrations.whylogs import WhylogsIntegration  # noqa
 from zenml.integrations.xgboost import XgboostIntegration  # noqa
+from zenml.integrations.vllm import VLLMIntegration  # noqa

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -478,7 +478,9 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         )
 
         # Yield metadata based on the generated execution object
-        yield from self.compute_metadata(execution=execution)
+        yield from self.compute_metadata(
+            execution=execution, settings=settings
+        )
 
         # mainly for testing purposes, we wait for the pipeline to finish
         if settings.synchronous:
@@ -518,19 +520,6 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             "pipeline_execution_arn": pipeline_execution_arn,
         }
 
-        aws_run_id = os.environ[ENV_ZENML_SAGEMAKER_RUN_ID].split("/")[-1]
-
-        region_name, _, _ = dissect_pipeline_execution_arn(
-            pipeline_execution_arn=pipeline_execution_arn
-        )
-
-        orchestrator_logs_url = (
-            f"https://{region_name}.console.aws.amazon.com/"
-            f"cloudwatch/home?region={region_name}#logsV2:log-groups/log-group"
-            f"/$252Faws$252Fsagemaker$252FProcessingJobs$3FlogStreamNameFilter"
-            f"$3Dpipelines-{aws_run_id}-"
-        )
-        run_metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_logs_url)
         return run_metadata
 
     def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
@@ -566,7 +555,7 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
 
         # Fetch the status of the _PipelineExecution
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:
@@ -590,12 +579,15 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             raise ValueError("Unknown status for the pipeline execution.")
 
     def compute_metadata(
-        self, execution: Any
+        self,
+        execution: Any,
+        settings: SagemakerOrchestratorSettings,
     ) -> Iterator[Dict[str, MetadataType]]:
         """Generate run metadata based on the generated Sagemaker Execution.
 
         Args:
             execution: The corresponding _PipelineExecution object.
+            settings: The Sagemaker orchestrator settings.
 
         Yields:
             A dictionary of metadata related to the pipeline run.
@@ -612,7 +604,9 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_url)
 
         # URL to the corresponding CloudWatch page
-        if logs_url := self._compute_orchestrator_logs_url(execution):
+        if logs_url := self._compute_orchestrator_logs_url(
+            execution, settings
+        ):
             metadata[METADATA_ORCHESTRATOR_LOGS_URL] = Uri(logs_url)
 
         yield metadata
@@ -656,11 +650,13 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
     @staticmethod
     def _compute_orchestrator_logs_url(
         pipeline_execution: Any,
+        settings: SagemakerOrchestratorSettings,
     ) -> Optional[str]:
         """Generate the CloudWatch URL upon pipeline execution.
 
         Args:
             pipeline_execution: The corresponding _PipelineExecution object.
+            settings: The Sagemaker orchestrator settings.
 
         Returns:
             the URL querying the pipeline logs in CloudWatch on AWS.
@@ -670,10 +666,16 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
                 pipeline_execution.arn
             )
 
+            use_training_jobs = True
+            if settings.use_training_step is not None:
+                use_training_jobs = settings.use_training_step
+
+            job_type = "Training" if use_training_jobs else "Processing"
+
             return (
                 f"https://{region_name}.console.aws.amazon.com/"
                 f"cloudwatch/home?region={region_name}#logsV2:log-groups/log-group"
-                f"/$252Faws$252Fsagemaker$252FProcessingJobs$3FlogStreamNameFilter"
+                f"/$252Faws$252Fsagemaker$252F{job_type}Jobs$3FlogStreamNameFilter"
                 f"$3Dpipelines-{execution_id}-"
             )
         except Exception as e:

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -482,7 +482,7 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
 
         # Fetch the status of the PipelineJob
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:

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/constants.py

@@ -76,4 +76,5 @@ WANDB = "wandb"
 VERTEX = "vertex"
 XGBOOST = "xgboost"
 VAULT = "vault"
+VLLM = "vllm"
 LIGHTNING = "lightning"

zenml/integrations/deepchecks/data_validators/deepchecks_data_validator.py

@@ -430,7 +430,7 @@ class DeepchecksDataValidator(BaseDataValidator):
         """Run one or more Deepchecks model validation checks.
 
         Call this method to perform model validation checks (e.g. confusion
-        matrix validation, performance reports, model error analyses, etc).
+        matrix validation, performance reports, model error analyzes, etc).
         A second dataset is required for model performance comparison tests
         (i.e. tests that identify changes in a model behavior by comparing how
         it performs on two different datasets).

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/gcp/orchestrators/vertex_orchestrator.py

@@ -835,7 +835,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
 
         # Fetch the status of the PipelineJob
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:

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]