zenml-nightly 0.68.1.dev20241106__tar.gz → 0.68.1.dev20241107__tar.gz

{zenml_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/PKG-INFO

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

{zenml_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.68.1.dev20241106"
+version = "0.68.1.dev20241107"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]
@@ -476,5 +476,8 @@ module = [
     "prodigy.components.db.*",
     "transformers.*",
     "langchain_community.*",
+    "vllm.*",
+    "numba.*",
+    "uvloop.*",
 ]
 ignore_missing_imports = true

zenml_nightly-0.68.1.dev20241107/src/zenml/VERSION

@@ -0,0 +1 @@
+0.68.1.dev20241107

{zenml_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/zenml/artifacts/utils.py

@@ -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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/zenml/integrations/vllm/services/vllm_deployment.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the vLLM Inference Server Service."""
 
+import argparse
 import os
 from typing import Any, List, Optional, Union
 
@@ -137,15 +138,23 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         self.endpoint.prepare_for_start()
 
         import uvloop
-        from vllm.entrypoints.openai.api_server import run_server
-        from vllm.entrypoints.openai.cli_args import make_arg_parser
-        from vllm.utils import FlexibleArgumentParser
+        from vllm.entrypoints.openai.api_server import (
+            run_server,
+        )
+        from vllm.entrypoints.openai.cli_args import (
+            make_arg_parser,
+        )
+        from vllm.utils import (
+            FlexibleArgumentParser,
+        )
 
         try:
-            parser = make_arg_parser(FlexibleArgumentParser())
-            args = parser.parse_args()
+            parser: argparse.ArgumentParser = make_arg_parser(
+                FlexibleArgumentParser()
+            )
+            args: argparse.Namespace = parser.parse_args()
             # Override port with the available port
-            self.config.port = self.endpoint.status.port
+            self.config.port = self.endpoint.status.port or self.config.port
             # Update the arguments in place
             args.__dict__.update(self.config.model_dump())
             uvloop.run(run_server(args=args))
@@ -161,7 +170,7 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
         """
         if not self.is_running:
             return None
-        return self.endpoint.prediction_url_path
+        return self.endpoint.prediction_url
 
     def predict(self, data: "Any") -> "Any":
         """Make a prediction using the service.

{zenml_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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 (
@@ -534,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)
@@ -595,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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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_nightly-0.68.1.dev20241106 → zenml_nightly-0.68.1.dev20241107}/src/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.dev20241106/src/zenml/VERSION

@@ -1 +0,0 @@
-0.68.1.dev20241106