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

zenml/orchestrators/step_runner.py

@@ -28,7 +28,8 @@ from typing import (
 )
 
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
-from zenml.artifacts.utils import save_artifact
+from zenml.artifacts.utils import _store_artifact_data_and_prepare_request
+from zenml.client import Client
 from zenml.config.step_configurations import StepConfiguration
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
@@ -153,7 +154,7 @@ class StepRunner:
 
             # Initialize the step context singleton
             StepContext._clear()
-            StepContext(
+            step_context = StepContext(
                 pipeline_run=pipeline_run,
                 step_run=step_run,
                 output_materializers=output_materializers,
@@ -246,6 +247,9 @@ class StepRunner:
                                 model_version=model_version,
                             )
                 finally:
+                    step_context._cleanup_registry.execute_callbacks(
+                        raise_on_exception=False
+                    )
                     StepContext._clear()  # Remove the step context singleton
 
             # Update the status and output artifacts of the step run.
@@ -531,7 +535,7 @@ class StepRunner:
             The IDs of the published output artifacts.
         """
         step_context = get_step_context()
-        output_artifacts: Dict[str, "ArtifactVersionResponse"] = {}
+        artifact_requests = []
 
         for output_name, return_value in output_data.items():
             data_type = type(return_value)
@@ -592,22 +596,24 @@ class StepRunner:
             # Get full set of tags
             tags = step_context.get_output_tags(output_name)
 
-            artifact = save_artifact(
+            artifact_request = _store_artifact_data_and_prepare_request(
                 name=artifact_name,
                 data=return_value,
-                materializer=materializer_class,
+                materializer_class=materializer_class,
                 uri=uri,
-                extract_metadata=artifact_metadata_enabled,
-                include_visualizations=artifact_visualization_enabled,
+                store_metadata=artifact_metadata_enabled,
+                store_visualizations=artifact_visualization_enabled,
                 has_custom_name=has_custom_name,
                 version=version,
                 tags=tags,
-                user_metadata=user_metadata,
-                manual_save=False,
+                metadata=user_metadata,
             )
-            output_artifacts[output_name] = artifact
+            artifact_requests.append(artifact_request)
 
-        return output_artifacts
+        responses = Client().zen_store.batch_create_artifact_versions(
+            artifact_requests
+        )
+        return dict(zip(output_data.keys(), responses))
 
     def load_and_run_hook(
         self,

zenml/stack/flavor.py

@@ -16,6 +16,7 @@
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
+from zenml.client import Client
 from zenml.enums import StackComponentType
 from zenml.models import (
     FlavorRequest,
@@ -146,9 +147,6 @@ class Flavor:
         Returns:
             The model.
         """
-        from zenml.client import Client
-
-        client = Client()
         connector_requirements = self.service_connector_requirements
         connector_type = (
             connector_requirements.connector_type
@@ -165,10 +163,16 @@ class Flavor:
             if connector_requirements
             else None
         )
+        user = None
+        workspace = None
+        if is_custom:
+            user = Client().active_user.id
+            workspace = Client().active_workspace.id
+
         model_class = FlavorRequest if is_custom else InternalFlavorRequest
         model = model_class(
-            user=client.active_user.id if is_custom else None,
-            workspace=client.active_workspace.id if is_custom else None,
+            user=user,
+            workspace=workspace,
             name=self.name,
             type=self.type,
             source=source_utils.resolve(self.__class__).import_path,

zenml/steps/step_context.py

@@ -26,6 +26,7 @@ from typing import (
 
 from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
+from zenml.utils.callback_registry import CallbackRegistry
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -145,6 +146,7 @@ class StepContext(metaclass=SingletonMetaClass):
             )
             for key in output_materializers.keys()
         }
+        self._cleanup_registry = CallbackRegistry()
 
     @property
     def pipeline(self) -> "PipelineResponse":

zenml/utils/callback_registry.py

@@ -0,0 +1,71 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Callback registry implementation."""
+
+from typing import Any, Callable, Dict, List, Tuple
+
+from typing_extensions import ParamSpec
+
+from zenml.logger import get_logger
+
+P = ParamSpec("P")
+
+logger = get_logger(__name__)
+
+
+class CallbackRegistry:
+    """Callback registry class."""
+
+    def __init__(self) -> None:
+        """Initializes the callback registry."""
+        self._callbacks: List[
+            Tuple[Callable[P, Any], Tuple[Any], Dict[str, Any]]
+        ] = []
+
+    def register_callback(
+        self, callback: Callable[P, Any], *args: P.args, **kwargs: P.kwargs
+    ) -> None:
+        """Register a callback.
+
+        Args:
+            callback: The callback to register.
+            *args: Arguments to call the callback with.
+            **kwargs: Keyword arguments to call the callback with.
+        """
+        self._callbacks.append((callback, args, kwargs))  # type: ignore[arg-type]
+
+    def reset(self) -> None:
+        """Reset the callbacks."""
+        self._callbacks = []
+
+    def execute_callbacks(self, raise_on_exception: bool) -> None:
+        """Execute all registered callbacks.
+
+        Args:
+            raise_on_exception: If True, exceptions raised during the execution
+                of the callbacks will be raised. If False, a warning with the
+                exception will be logged instead.
+
+        Raises:
+            Exception: Exceptions raised in any of the callbacks if
+                `raise_on_exception` is set to True.
+        """
+        for callback, args, kwargs in self._callbacks:
+            try:
+                callback(*args, **kwargs)
+            except Exception as e:
+                if raise_on_exception:
+                    raise e
+                else:
+                    logger.warning("Failed to run callback: %s", str(e))

zenml/zen_server/rbac/endpoint_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """High-level helper functions to write endpoints with RBAC."""
 
-from typing import Any, Callable, TypeVar, Union
+from typing import Any, Callable, List, TypeVar, Union
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -96,6 +96,48 @@ def verify_permissions_and_create_entity(
     return created
 
 
+def verify_permissions_and_batch_create_entity(
+    batch: List[AnyRequest],
+    resource_type: ResourceType,
+    create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
+) -> List[AnyResponse]:
+    """Verify permissions and create a batch of entities if authorized.
+
+    Args:
+        batch: The batch to create.
+        resource_type: The resource type of the entities to create.
+        create_method: The method to create the entities.
+
+    Raises:
+        IllegalOperationError: If the request model has a different owner then
+            the currently authenticated user.
+        RuntimeError: If the resource type is usage-tracked.
+
+    Returns:
+        The created entities.
+    """
+    auth_context = get_auth_context()
+    assert auth_context
+
+    for request_model in batch:
+        if isinstance(request_model, UserScopedRequest):
+            if request_model.user != auth_context.user.id:
+                raise IllegalOperationError(
+                    f"Not allowed to create resource '{resource_type}' for a "
+                    "different user."
+                )
+
+    verify_permission(resource_type=resource_type, action=Action.CREATE)
+
+    if resource_type in REPORTABLE_RESOURCES:
+        raise RuntimeError(
+            "Batch requests are currently not possible with usage-tracked features."
+        )
+
+    created = create_method(batch)
+    return created
+
+
 def verify_permissions_and_get_entity(
     id: UUIDOrStr,
     get_method: Callable[[UUIDOrStr], AnyResponse],

zenml/zen_server/routers/artifact_version_endpoints.py

@@ -13,12 +13,13 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for artifact versions."""
 
+from typing import List
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
 
 from zenml.artifacts.utils import load_artifact_visualization
-from zenml.constants import API, ARTIFACT_VERSIONS, VERSION_1, VISUALIZE
+from zenml.constants import API, ARTIFACT_VERSIONS, BATCH, VERSION_1, VISUALIZE
 from zenml.models import (
     ArtifactVersionFilter,
     ArtifactVersionRequest,
@@ -30,6 +31,7 @@ from zenml.models import (
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_batch_create_entity,
     verify_permissions_and_create_entity,
     verify_permissions_and_delete_entity,
     verify_permissions_and_get_entity,
@@ -118,6 +120,30 @@ def create_artifact_version(
     )
 
 
+@artifact_version_router.post(
+    BATCH,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def batch_create_artifact_version(
+    artifact_versions: List[ArtifactVersionRequest],
+    _: AuthContext = Security(authorize),
+) -> List[ArtifactVersionResponse]:
+    """Create a batch of artifact versions.
+
+    Args:
+        artifact_versions: The artifact versions to create.
+
+    Returns:
+        The created artifact versions.
+    """
+    return verify_permissions_and_batch_create_entity(
+        batch=artifact_versions,
+        resource_type=ResourceType.ARTIFACT_VERSION,
+        create_method=zen_store().batch_create_artifact_versions,
+    )
+
+
 @artifact_version_router.get(
     "/{artifact_version_id}",
     response_model=ArtifactVersionResponse,

zenml/zen_stores/rest_zen_store.py

@@ -57,6 +57,7 @@ from zenml.constants import (
     ARTIFACT_VERSIONS,
     ARTIFACT_VISUALIZATIONS,
     ARTIFACTS,
+    BATCH,
     CODE_REFERENCES,
     CODE_REPOSITORIES,
     CONFIG,
@@ -991,6 +992,23 @@ class RestZenStore(BaseZenStore):
             route=ARTIFACT_VERSIONS,
         )
 
+    def batch_create_artifact_versions(
+        self, artifact_versions: List[ArtifactVersionRequest]
+    ) -> List[ArtifactVersionResponse]:
+        """Creates a batch of artifact versions.
+
+        Args:
+            artifact_versions: The artifact versions to create.
+
+        Returns:
+            The created artifact versions.
+        """
+        return self._batch_create_resources(
+            resources=artifact_versions,
+            response_model=ArtifactVersionResponse,
+            route=ARTIFACT_VERSIONS,
+        )
+
     def get_artifact_version(
         self, artifact_version_id: UUID, hydrate: bool = True
     ) -> ArtifactVersionResponse:
@@ -4518,6 +4536,40 @@ class RestZenStore(BaseZenStore):
 
         return response_model.model_validate(response_body)
 
+    def _batch_create_resources(
+        self,
+        resources: List[AnyRequest],
+        response_model: Type[AnyResponse],
+        route: str,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> List[AnyResponse]:
+        """Create a new batch of resources.
+
+        Args:
+            resources: The resources to create.
+            response_model: The response model of an individual resource.
+            route: The resource REST route to use.
+            params: Optional query parameters to pass to the endpoint.
+
+        Returns:
+            List of response models.
+        """
+        json_data = [
+            resource.model_dump(mode="json") for resource in resources
+        ]
+        response = self._request(
+            "POST",
+            self.url + API + VERSION_1 + route + BATCH,
+            json=json_data,
+            params=params,
+        )
+        assert isinstance(response, list)
+
+        return [
+            response_model.model_validate(model_data)
+            for model_data in response
+        ]
+
     def _create_workspace_scoped_resource(
         self,
         resource: AnyWorkspaceScopedRequest,

zenml/zen_stores/sql_zen_store.py

@@ -2915,6 +2915,22 @@ class SqlZenStore(BaseZenStore):
                 include_metadata=True, include_resources=True
             )
 
+    def batch_create_artifact_versions(
+        self, artifact_versions: List[ArtifactVersionRequest]
+    ) -> List[ArtifactVersionResponse]:
+        """Creates a batch of artifact versions.
+
+        Args:
+            artifact_versions: The artifact versions to create.
+
+        Returns:
+            The created artifact versions.
+        """
+        return [
+            self.create_artifact_version(artifact_version)
+            for artifact_version in artifact_versions
+        ]
+
     def get_artifact_version(
         self, artifact_version_id: UUID, hydrate: bool = True
     ) -> ArtifactVersionResponse:

zenml/zen_stores/zen_store_interface.py

@@ -663,6 +663,19 @@ class ZenStoreInterface(ABC):
             The created artifact version.
         """
 
+    @abstractmethod
+    def batch_create_artifact_versions(
+        self, artifact_versions: List[ArtifactVersionRequest]
+    ) -> List[ArtifactVersionResponse]:
+        """Creates a batch of artifact versions.
+
+        Args:
+            artifact_versions: The artifact versions to create.
+
+        Returns:
+            The created artifact versions.
+        """
+
     @abstractmethod
     def get_artifact_version(
         self, artifact_version_id: UUID, hydrate: bool = True

{zenml_nightly-0.68.1.dev20241105.dist-info → zenml_nightly-0.68.1.dev20241107.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.68.1.dev20241105
+Version: 0.68.1.dev20241107
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0