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

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(

zenml/zen_stores/sql_zen_store.py

@@ -86,6 +86,7 @@ from zenml.config.global_config import GlobalConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.config.server_config import ServerConfiguration
+from zenml.config.step_configurations import StepConfiguration, StepSpec
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
     DEFAULT_PASSWORD,
@@ -118,7 +119,6 @@ from zenml.enums import (
     StackComponentType,
     StackDeploymentProvider,
     StepRunInputArtifactType,
-    StepRunOutputArtifactType,
     StoreType,
     TaggableResourceTypes,
 )
@@ -216,9 +216,7 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -5481,9 +5479,7 @@ class SqlZenStore(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:
@@ -5492,7 +5488,6 @@ class SqlZenStore(BaseZenStore):
         Returns:
             The created run metadata.
         """
-        return_value: List[RunMetadataResponse] = []
         with Session(self.engine) as session:
             for key, value in run_metadata.values.items():
                 type_ = run_metadata.types[key]
@@ -5508,70 +5503,7 @@ class SqlZenStore(BaseZenStore):
                 )
                 session.add(run_metadata_schema)
                 session.commit()
-                return_value.append(
-                    run_metadata_schema.to_model(
-                        include_metadata=True, include_resources=True
-                    )
-                )
-        return return_value
-
-    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.
-
-        Raises:
-            KeyError: if the run metadata doesn't exist.
-        """
-        with Session(self.engine) as session:
-            run_metadata = session.exec(
-                select(RunMetadataSchema).where(
-                    RunMetadataSchema.id == run_metadata_id
-                )
-            ).first()
-            if run_metadata is None:
-                raise KeyError(
-                    f"Unable to get run metadata with ID "
-                    f"{run_metadata_id}: "
-                    f"No run metadata with this ID found."
-                )
-            return run_metadata.to_model(
-                include_metadata=hydrate, include_resources=True
-            )
-
-    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.
-        """
-        with Session(self.engine) as session:
-            query = select(RunMetadataSchema)
-            return self.filter_and_paginate(
-                session=session,
-                query=query,
-                table=RunMetadataSchema,
-                filter_model=run_metadata_filter_model,
-                hydrate=hydrate,
-            )
+        return None
 
     # ----------------------------- Schedules -----------------------------
 
@@ -8189,25 +8121,35 @@ class SqlZenStore(BaseZenStore):
                     session=session,
                 )
 
+            session.commit()
+            session.refresh(step_schema)
+
+            step_model = step_schema.to_model(include_metadata=True)
+
             # Save input artifact IDs into the database.
             for input_name, artifact_version_id in step_run.inputs.items():
+                input_type = self._get_step_run_input_type(
+                    input_name=input_name,
+                    step_config=step_model.config,
+                    step_spec=step_model.spec,
+                )
                 self._set_run_step_input_artifact(
                     run_step_id=step_schema.id,
                     artifact_version_id=artifact_version_id,
                     name=input_name,
-                    input_type=StepRunInputArtifactType.DEFAULT,
+                    input_type=input_type,
                     session=session,
                 )
 
             # Save output artifact IDs into the database.
-            for output_name, artifact_version_id in step_run.outputs.items():
-                self._set_run_step_output_artifact(
-                    step_run_id=step_schema.id,
-                    artifact_version_id=artifact_version_id,
-                    name=output_name,
-                    output_type=StepRunOutputArtifactType.DEFAULT,
-                    session=session,
-                )
+            for output_name, artifact_version_ids in step_run.outputs.items():
+                for artifact_version_id in artifact_version_ids:
+                    self._set_run_step_output_artifact(
+                        step_run_id=step_schema.id,
+                        artifact_version_id=artifact_version_id,
+                        name=output_name,
+                        session=session,
+                    )
 
             if step_run.status != ExecutionStatus.RUNNING:
                 self._update_pipeline_run_status(
@@ -8215,6 +8157,7 @@ class SqlZenStore(BaseZenStore):
                 )
 
             session.commit()
+            session.refresh(step_schema)
 
             return step_schema.to_model(
                 include_metadata=True, include_resources=True
@@ -8307,26 +8250,12 @@ class SqlZenStore(BaseZenStore):
             existing_step_run.update(step_run_update)
             session.add(existing_step_run)
 
-            # Update the output artifacts.
+            # Update the artifacts.
             for name, artifact_version_id in step_run_update.outputs.items():
                 self._set_run_step_output_artifact(
                     step_run_id=step_run_id,
                     artifact_version_id=artifact_version_id,
                     name=name,
-                    output_type=StepRunOutputArtifactType.DEFAULT,
-                    session=session,
-                )
-
-            # Update saved artifacts
-            for (
-                artifact_name,
-                artifact_version_id,
-            ) in step_run_update.saved_artifact_versions.items():
-                self._set_run_step_output_artifact(
-                    step_run_id=step_run_id,
-                    artifact_version_id=artifact_version_id,
-                    name=artifact_name,
-                    output_type=StepRunOutputArtifactType.MANUAL,
                     session=session,
                 )
 
@@ -8355,6 +8284,34 @@ class SqlZenStore(BaseZenStore):
                 include_metadata=True, include_resources=True
             )
 
+    def _get_step_run_input_type(
+        self,
+        input_name: str,
+        step_config: StepConfiguration,
+        step_spec: StepSpec,
+    ) -> StepRunInputArtifactType:
+        """Get the input type of an artifact.
+
+        Args:
+            input_name: The name of the input artifact.
+            step_config: The step config.
+            step_spec: The step spec.
+
+        Returns:
+            The input type of the artifact.
+        """
+        if input_name in step_spec.inputs:
+            return StepRunInputArtifactType.STEP_OUTPUT
+        if input_name in step_config.external_input_artifacts:
+            return StepRunInputArtifactType.EXTERNAL
+        elif (
+            input_name in step_config.model_artifacts_or_metadata
+            or input_name in step_config.client_lazy_loaders
+        ):
+            return StepRunInputArtifactType.LAZY_LOADED
+        else:
+            return StepRunInputArtifactType.MANUAL
+
     @staticmethod
     def _set_run_step_parent_step(
         child_id: UUID, parent_id: UUID, session: Session
@@ -8473,7 +8430,6 @@ class SqlZenStore(BaseZenStore):
         step_run_id: UUID,
         artifact_version_id: UUID,
         name: str,
-        output_type: StepRunOutputArtifactType,
         session: Session,
     ) -> None:
         """Sets an artifact as an output of a step run.
@@ -8482,7 +8438,6 @@ class SqlZenStore(BaseZenStore):
             step_run_id: The ID of the step run.
             artifact_version_id: The ID of the artifact version.
             name: The name of the output in the step run.
-            output_type: In which way the artifact was saved by the step.
             session: The database session to use.
 
         Raises:
@@ -8526,7 +8481,6 @@ class SqlZenStore(BaseZenStore):
             step_id=step_run_id,
             artifact_id=artifact_version_id,
             name=name,
-            type=output_type.value,
         )
         session.add(assignment)
 

zenml/zen_stores/zen_store_interface.py

@@ -90,9 +90,7 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -1633,52 +1631,14 @@ class ZenStoreInterface(ABC):
     # -------------------- Run metadata --------------------
 
     @abstractmethod
-    def create_run_metadata(
-        self, run_metadata: RunMetadataRequest
-    ) -> List[RunMetadataResponse]:
+    def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
         """Creates run metadata.
 
         Args:
             run_metadata: The run metadata to create.
 
         Returns:
-            The created run metadata.
-        """
-
-    @abstractmethod
-    def get_run_metadata(
-        self, run_metadata_id: UUID, hydrate: bool = True
-    ) -> RunMetadataResponse:
-        """Get run metadata by its unique 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 with the given ID.
-
-        Raises:
-            KeyError: if the run metadata doesn't exist.
-        """
-
-    @abstractmethod
-    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.
+            None
         """
 
     # -------------------- Schedules --------------------

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

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