zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.68.1.dev20241108__py3-none-any.whl

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_server/routers/workspaces_endpoints.py

@@ -74,7 +74,6 @@ from zenml.models import (
     PipelineRunRequest,
     PipelineRunResponse,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -977,7 +976,6 @@ def get_or_create_pipeline_run(
 
 @router.post(
     WORKSPACES + "/{workspace_name_or_id}" + RUN_METADATA,
-    response_model=List[RunMetadataResponse],
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -985,7 +983,7 @@ def create_run_metadata(
     workspace_name_or_id: Union[str, UUID],
     run_metadata: RunMetadataRequest,
     auth_context: AuthContext = Security(authorize),
-) -> List[RunMetadataResponse]:
+) -> None:
     """Creates run metadata.
 
     Args:
@@ -1039,7 +1037,8 @@ def create_run_metadata(
         resource_type=ResourceType.RUN_METADATA, action=Action.CREATE
     )
 
-    return zen_store().create_run_metadata(run_metadata)
+    zen_store().create_run_metadata(run_metadata)
+    return None
 
 
 @router.post(

zenml/zen_server/zen_server_api.py

@@ -70,7 +70,6 @@ from zenml.zen_server.routers import (
     pipeline_deployments_endpoints,
     pipelines_endpoints,
     plugin_endpoints,
-    run_metadata_endpoints,
     run_templates_endpoints,
     runs_endpoints,
     schedule_endpoints,
@@ -423,7 +422,6 @@ app.include_router(pipelines_endpoints.router)
 app.include_router(pipeline_builds_endpoints.router)
 app.include_router(pipeline_deployments_endpoints.router)
 app.include_router(runs_endpoints.router)
-app.include_router(run_metadata_endpoints.router)
 app.include_router(run_templates_endpoints.router)
 app.include_router(schedule_endpoints.router)
 app.include_router(secrets_endpoints.router)

zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py

@@ -0,0 +1,99 @@
+"""move artifact save type [1cb6477f72d6].
+
+Revision ID: 1cb6477f72d6
+Revises: c22561cbb3a9
+Create Date: 2024-10-10 15:44:09.465210
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "1cb6477f72d6"
+down_revision = "c22561cbb3a9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # Step 1: Add nullable save_type column to artifact_version
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("save_type", sa.TEXT(), nullable=True))
+
+    # Step 2: Move data from step_run_output_artifact.type to artifact_version.save_type
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = (
+            SELECT max(step_run_output_artifact.type)
+            FROM step_run_output_artifact
+            WHERE step_run_output_artifact.artifact_id = artifact_version.id
+            GROUP BY artifact_id
+        )
+    """)
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = 'step_output'
+        WHERE artifact_version.save_type = 'default'
+    """)
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = 'external'
+        WHERE save_type is NULL
+    """)
+
+    # # Step 3: Set save_type to non-nullable
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.alter_column(
+            "save_type",
+            existing_type=sa.TEXT(),
+            nullable=False,
+        )
+
+    # Step 4: Remove type column from step_run_output_artifact
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.drop_column("type")
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # Add type column back to step_run_output_artifact
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.add_column(
+            sa.Column("type", sa.TEXT(), nullable=True),
+        )
+
+    # Move data back from artifact_version.save_type to step_run_output_artifact.type
+    op.execute("""
+        UPDATE step_run_output_artifact
+        SET type = (
+            SELECT max(artifact_version.save_type)
+            FROM artifact_version
+            WHERE step_run_output_artifact.artifact_id = artifact_version.id
+            GROUP BY artifact_id
+        )
+    """)
+    op.execute("""
+        UPDATE step_run_output_artifact
+        SET type = 'default'
+        WHERE step_run_output_artifact.type = 'step_output'
+    """)
+
+    # Set type to non-nullable
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.alter_column(
+            "type",
+            existing_type=sa.TEXT(),
+            nullable=False,
+        )
+
+    # Remove save_type column from artifact_version
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.drop_column("save_type")

zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py

@@ -0,0 +1,33 @@
+"""Update step run input types [b557b2871693].
+
+Revision ID: b557b2871693
+Revises: 1cb6477f72d6
+Create Date: 2024-10-30 13:06:55.147202
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b557b2871693"
+down_revision = "1cb6477f72d6"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    op.execute("""
+        UPDATE step_run_input_artifact
+        SET type = 'step_output'
+        WHERE type = 'default'
+    """)
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    op.execute("""
+        UPDATE step_run_input_artifact
+        SET type = 'default'
+        WHERE type = 'step_output'
+    """)

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,
@@ -200,9 +201,7 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -991,6 +990,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:
@@ -1996,9 +2012,7 @@ class RestZenStore(BaseZenStore):
 
     # ----------------------------- Run Metadata -----------------------------
 
-    def create_run_metadata(
-        self, run_metadata: RunMetadataRequest
-    ) -> List[RunMetadataResponse]:
+    def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
         """Creates run metadata.
 
         Args:
@@ -2008,55 +2022,8 @@ class RestZenStore(BaseZenStore):
             The created run metadata.
         """
         route = f"{WORKSPACES}/{str(run_metadata.workspace)}{RUN_METADATA}"
-        response_body = self.post(f"{route}", body=run_metadata)
-        result: List[RunMetadataResponse] = []
-        if isinstance(response_body, list):
-            for metadata in response_body or []:
-                result.append(RunMetadataResponse.model_validate(metadata))
-        return result
-
-    def get_run_metadata(
-        self, run_metadata_id: UUID, hydrate: bool = True
-    ) -> RunMetadataResponse:
-        """Gets run metadata with the given ID.
-
-        Args:
-            run_metadata_id: The ID of the run metadata to get.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._get_resource(
-            resource_id=run_metadata_id,
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            params={"hydrate": hydrate},
-        )
-
-    def list_run_metadata(
-        self,
-        run_metadata_filter_model: RunMetadataFilter,
-        hydrate: bool = False,
-    ) -> Page[RunMetadataResponse]:
-        """List run metadata.
-
-        Args:
-            run_metadata_filter_model: All filter parameters including
-                pagination params.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._list_paginated_resources(
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            filter_model=run_metadata_filter_model,
-            params={"hydrate": hydrate},
-        )
+        self.post(f"{route}", body=run_metadata)
+        return None
 
     # ----------------------------- Schedules -----------------------------
 
@@ -4518,6 +4485,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/schemas/artifact_schemas.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of artifact table."""
 
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
@@ -23,6 +24,7 @@ from sqlmodel import Field, Relationship
 
 from zenml.config.source import Source
 from zenml.enums import (
+    ArtifactSaveType,
     ArtifactType,
     ExecutionStatus,
     MetadataResourceTypes,
@@ -196,6 +198,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             overlaps="tags",
         ),
     )
+    save_type: str = Field(sa_column=Column(TEXT, nullable=False))
 
     # Foreign keys
     artifact_id: UUID = build_foreign_key_field(
@@ -300,6 +303,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             uri=artifact_version_request.uri,
             materializer=artifact_version_request.materializer.model_dump_json(),
             data_type=artifact_version_request.data_type.model_dump_json(),
+            save_type=artifact_version_request.save_type.value,
         )
 
     def to_model(
@@ -360,6 +364,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             updated=self.updated,
             tags=[t.tag.to_model() for t in self.tags],
             producer_pipeline_run_id=producer_pipeline_run_id,
+            save_type=ArtifactSaveType(self.save_type),
             artifact_store_id=self.artifact_store_id,
         )
 
@@ -370,7 +375,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
                 workspace=self.workspace.to_model(),
                 producer_step_run_id=producer_step_run_id,
                 visualizations=[v.to_model() for v in self.visualizations],
-                run_metadata={m.key: m.to_model() for m in self.run_metadata},
+                run_metadata={
+                    m.key: json.loads(m.value) for m in self.run_metadata
+                },
             )
 
         resources = None

zenml/zen_stores/schemas/model_schemas.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
@@ -379,8 +380,7 @@ class ModelVersionSchema(NamedSchema, table=True):
                 workspace=self.workspace.to_model(),
                 description=self.description,
                 run_metadata={
-                    rm.key: rm.to_model(include_metadata=True)
-                    for rm in self.run_metadata
+                    rm.key: json.loads(rm.value) for rm in self.run_metadata
                 },
             )
 

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -271,7 +271,7 @@ class PipelineRunSchema(NamedSchema, table=True):
         )
 
         run_metadata = {
-            metadata_schema.key: metadata_schema.to_model()
+            metadata_schema.key: json.loads(metadata_schema.value)
             for metadata_schema in self.run_metadata
         }
 

zenml/zen_stores/schemas/run_metadata_schemas.py

@@ -13,20 +13,13 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of pipeline run metadata tables."""
 
-import json
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, List, Optional
 from uuid import UUID
 
 from sqlalchemy import TEXT, VARCHAR, Column
 from sqlmodel import Field, Relationship
 
 from zenml.enums import MetadataResourceTypes
-from zenml.metadata.metadata_types import MetadataTypeEnum
-from zenml.models import (
-    RunMetadataResponse,
-    RunMetadataResponseBody,
-    RunMetadataResponseMetadata,
-)
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
@@ -110,43 +103,3 @@ class RunMetadataSchema(BaseSchema, table=True):
     key: str
     value: str = Field(sa_column=Column(TEXT, nullable=False))
     type: str
-
-    def to_model(
-        self,
-        include_metadata: bool = False,
-        include_resources: bool = False,
-        **kwargs: Any,
-    ) -> "RunMetadataResponse":
-        """Convert a `RunMetadataSchema` to a `RunMetadataResponse`.
-
-        Args:
-            include_metadata: Whether the metadata will be filled.
-            include_resources: Whether the resources will be filled.
-            **kwargs: Keyword arguments to allow schema specific logic
-
-
-        Returns:
-            The created `RunMetadataResponse`.
-        """
-        body = RunMetadataResponseBody(
-            user=self.user.to_model() if self.user else None,
-            key=self.key,
-            created=self.created,
-            updated=self.updated,
-            value=json.loads(self.value),
-            type=MetadataTypeEnum(self.type),
-        )
-        metadata = None
-        if include_metadata:
-            metadata = RunMetadataResponseMetadata(
-                workspace=self.workspace.to_model(),
-                resource_id=self.resource_id,
-                resource_type=MetadataResourceTypes(self.resource_type),
-                stack_component_id=self.stack_component_id,
-            )
-
-        return RunMetadataResponse(
-            id=self.id,
-            body=body,
-            metadata=metadata,
-        )

zenml/zen_stores/schemas/step_run_schemas.py

@@ -15,7 +15,7 @@
 
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
 from pydantic import ConfigDict
@@ -28,6 +28,7 @@ from zenml.constants import MEDIUMTEXT_MAX_LENGTH
 from zenml.enums import (
     ExecutionStatus,
     MetadataResourceTypes,
+    StepRunInputArtifactType,
 )
 from zenml.models import (
     StepRunRequest,
@@ -36,7 +37,11 @@ from zenml.models import (
     StepRunResponseMetadata,
     StepRunUpdate,
 )
-from zenml.models.v2.core.step_run import StepRunResponseResources
+from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+from zenml.models.v2.core.step_run import (
+    StepRunInputResponse,
+    StepRunResponseResources,
+)
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
@@ -214,21 +219,25 @@ class StepRunSchema(NamedSchema, table=True):
                 or a step_configuration.
         """
         run_metadata = {
-            metadata_schema.key: metadata_schema.to_model()
+            metadata_schema.key: json.loads(metadata_schema.value)
             for metadata_schema in self.run_metadata
         }
 
         input_artifacts = {
-            artifact.name: artifact.artifact_version.to_model()
+            artifact.name: StepRunInputResponse(
+                input_type=StepRunInputArtifactType(artifact.type),
+                **artifact.artifact_version.to_model().model_dump(),
+            )
             for artifact in self.input_artifacts
         }
 
-        output_artifacts = {
-            artifact.name: artifact.artifact_version.to_model(
-                pipeline_run_id_in_context=self.pipeline_run_id
+        output_artifacts: Dict[str, List["ArtifactVersionResponse"]] = {}
+        for artifact in self.output_artifacts:
+            if artifact.name not in output_artifacts:
+                output_artifacts[artifact.name] = []
+            output_artifacts[artifact.name].append(
+                artifact.artifact_version.to_model()
             )
-            for artifact in self.output_artifacts
-        }
 
         full_step_config = None
         if self.deployment is not None:
@@ -398,7 +407,6 @@ class StepRunOutputArtifactSchema(SQLModel, table=True):
 
     # Fields
     name: str
-    type: str
 
     # Foreign keys
     step_id: UUID = build_foreign_key_field(