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

zenml/orchestrators/input_utils.py

@@ -18,17 +18,19 @@ from uuid import UUID
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
+from zenml.enums import ArtifactSaveType, StepRunInputArtifactType
 from zenml.exceptions import InputResolutionError
 from zenml.utils import pagination_utils
 
 if TYPE_CHECKING:
-    from zenml.models import ArtifactVersionResponse, PipelineRunResponse
+    from zenml.models import PipelineRunResponse
+    from zenml.models.v2.core.step_run import StepRunInputResponse
 
 
 def resolve_step_inputs(
     step: "Step",
     pipeline_run: "PipelineRunResponse",
-) -> Tuple[Dict[str, "ArtifactVersionResponse"], List[UUID]]:
+) -> Tuple[Dict[str, "StepRunInputResponse"], List[UUID]]:
     """Resolves inputs for the current step.
 
     Args:
@@ -45,7 +47,8 @@ def resolve_step_inputs(
         The IDs of the input artifact versions and the IDs of parent steps of
             the current step.
     """
-    from zenml.models import ArtifactVersionResponse, RunMetadataResponse
+    from zenml.models import ArtifactVersionResponse
+    from zenml.models.v2.core.step_run import StepRunInputResponse
 
     current_run_steps = {
         run_step.name: run_step
@@ -54,7 +57,7 @@ def resolve_step_inputs(
         )
     }
 
-    input_artifacts: Dict[str, "ArtifactVersionResponse"] = {}
+    input_artifacts: Dict[str, StepRunInputResponse] = {}
     for name, input_ in step.spec.inputs.items():
         try:
             step_run = current_run_steps[input_.step_name]
@@ -64,22 +67,44 @@ def resolve_step_inputs(
             )
 
         try:
-            artifact = step_run.outputs[input_.output_name]
+            outputs = step_run.outputs[input_.output_name]
         except KeyError:
             raise InputResolutionError(
-                f"No output `{input_.output_name}` found for step "
+                f"No step output `{input_.output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 
-        input_artifacts[name] = artifact
+        step_outputs = [
+            output
+            for output in outputs
+            if output.save_type == ArtifactSaveType.STEP_OUTPUT
+        ]
+        if len(step_outputs) > 2:
+            # This should never happen, there can only be a single regular step
+            # output for a name
+            raise InputResolutionError(
+                f"Too many step outputs for output `{input_.output_name}` of "
+                f"step `{input_.step_name}`."
+            )
+        elif len(step_outputs) == 0:
+            raise InputResolutionError(
+                f"No step output `{input_.output_name}` found for step "
+                f"`{input_.step_name}`."
+            )
+
+        input_artifacts[name] = StepRunInputResponse(
+            input_type=StepRunInputArtifactType.STEP_OUTPUT,
+            **step_outputs[0].model_dump(),
+        )
 
     for (
         name,
         external_artifact,
     ) in step.config.external_input_artifacts.items():
         artifact_version_id = external_artifact.get_artifact_version_id()
-        input_artifacts[name] = Client().get_artifact_version(
-            artifact_version_id
+        input_artifacts[name] = StepRunInputResponse(
+            input_type=StepRunInputArtifactType.EXTERNAL,
+            **Client().get_artifact_version(artifact_version_id).model_dump(),
         )
 
     for name, config_ in step.config.model_artifacts_or_metadata.items():
@@ -98,9 +123,7 @@ def resolve_step_inputs(
             ):
                 # metadata values should go directly in parameters, as primitive types
                 step.config.parameters[name] = (
-                    context_model_version.run_metadata[
-                        config_.metadata_name
-                    ].value
+                    context_model_version.run_metadata[config_.metadata_name]
                 )
             elif config_.artifact_name is None:
                 err_msg = (
@@ -112,14 +135,15 @@ def resolve_step_inputs(
                     config_.artifact_name, config_.artifact_version
                 ):
                     if config_.metadata_name is None:
-                        input_artifacts[name] = artifact_
+                        input_artifacts[name] = StepRunInputResponse(
+                            input_type=StepRunInputArtifactType.LAZY_LOADED,
+                            **artifact_.model_dump(),
+                        )
                     elif config_.metadata_name:
                         # metadata values should go directly in parameters, as primitive types
                         try:
                             step.config.parameters[name] = (
-                                artifact_.run_metadata[
-                                    config_.metadata_name
-                                ].value
+                                artifact_.run_metadata[config_.metadata_name]
                             )
                         except KeyError:
                             err_msg = (
@@ -141,9 +165,10 @@ def resolve_step_inputs(
     for name, cll_ in step.config.client_lazy_loaders.items():
         value_ = cll_.evaluate()
         if isinstance(value_, ArtifactVersionResponse):
-            input_artifacts[name] = value_
-        elif isinstance(value_, RunMetadataResponse):
-            step.config.parameters[name] = value_.value
+            input_artifacts[name] = StepRunInputResponse(
+                input_type=StepRunInputArtifactType.LAZY_LOADED,
+                **value_.model_dump(),
+            )
         else:
             step.config.parameters[name] = value_
 

zenml/orchestrators/step_launcher.py

@@ -33,13 +33,13 @@ from zenml.environment import get_run_environment_dict
 from zenml.logger import get_logger
 from zenml.logging import step_logging
 from zenml.models import (
-    ArtifactVersionResponse,
     LogsRequest,
     PipelineDeploymentResponse,
     PipelineRunRequest,
     PipelineRunResponse,
     StepRunResponse,
 )
+from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators import output_utils, publish_utils, step_run_utils
 from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
@@ -442,7 +442,7 @@ class StepLauncher:
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
         step_run_info: StepRunInfo,
-        input_artifacts: Dict[str, ArtifactVersionResponse],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
         last_retry: bool,
     ) -> None:

zenml/orchestrators/step_run_utils.py

@@ -14,12 +14,12 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Mapping, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple
 
 from zenml.client import Client
 from zenml.config.step_configurations import ArtifactConfiguration, Step
 from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ExecutionStatus
+from zenml.enums import ArtifactSaveType, ExecutionStatus
 from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
@@ -104,6 +104,7 @@ class StepRunRequestFactory:
             input_name: artifact.id
             for input_name, artifact in input_artifacts.items()
         }
+
         request.inputs = input_artifact_ids
         request.parent_step_ids = parent_step_ids
 
@@ -142,8 +143,8 @@ class StepRunRequestFactory:
 
                 request.original_step_run_id = cached_step_run.id
                 request.outputs = {
-                    output_name: artifact.id
-                    for output_name, artifact in cached_step_run.outputs.items()
+                    output_name: [artifact.id for artifact in artifacts]
+                    for output_name, artifacts in cached_step_run.outputs.items()
                 }
 
                 request.status = ExecutionStatus.CACHED
@@ -551,7 +552,7 @@ def link_pipeline_run_to_model_version(
 
 
 def link_output_artifacts_to_model_version(
-    artifacts: Dict[str, ArtifactVersionResponse],
+    artifacts: Dict[str, List[ArtifactVersionResponse]],
     output_configurations: Mapping[str, ArtifactConfiguration],
     model_version: ModelVersionResponse,
 ) -> None:
@@ -562,13 +563,16 @@ def link_output_artifacts_to_model_version(
         output_configurations: The output configurations for the step.
         model_version: The model version to link.
     """
-    for output_name, output_artifact in artifacts.items():
-        artifact_config = None
-        if output_config := output_configurations.get(output_name, None):
-            artifact_config = output_config.artifact_config
-
-        link_artifact_version_to_model_version(
-            artifact_version=output_artifact,
-            model_version=model_version,
-            artifact_config=artifact_config,
-        )
+    for output_name, output_artifacts in artifacts.items():
+        for output_artifact in output_artifacts:
+            artifact_config = None
+            if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT and (
+                output_config := output_configurations.get(output_name, None)
+            ):
+                artifact_config = output_config.artifact_config
+
+            link_artifact_version_to_model_version(
+                artifact_version=output_artifact,
+                model_version=model_version,
+                artifact_config=artifact_config,
+            )

zenml/orchestrators/step_runner.py

@@ -37,10 +37,12 @@ from zenml.constants import (
     ENV_ZENML_IGNORE_FAILURE_HOOK,
     handle_bool_env_var,
 )
+from zenml.enums import ArtifactSaveType
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.logging.step_logging import StepLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
@@ -99,7 +101,7 @@ class StepRunner:
         self,
         pipeline_run: "PipelineRunResponse",
         step_run: "StepRunResponse",
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
         step_run_info: StepRunInfo,
     ) -> None:
@@ -242,7 +244,9 @@ class StepRunner:
                             from zenml.orchestrators import step_run_utils
 
                             step_run_utils.link_output_artifacts_to_model_version(
-                                artifacts=output_artifacts,
+                                artifacts={
+                                    k: [v] for k, v in output_artifacts.items()
+                                },
                                 output_configurations=step_run.config.outputs,
                                 model_version=model_version,
                             )
@@ -303,7 +307,7 @@ class StepRunner:
         self,
         args: List[str],
         annotations: Dict[str, Any],
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
     ) -> Dict[str, Any]:
         """Parses the inputs for a step entrypoint function.
 
@@ -606,6 +610,7 @@ class StepRunner:
                 has_custom_name=has_custom_name,
                 version=version,
                 tags=tags,
+                save_type=ArtifactSaveType.STEP_OUTPUT,
                 metadata=user_metadata,
             )
             artifact_requests.append(artifact_request)

zenml/steps/base_step.py

@@ -327,12 +327,12 @@ class BaseStep:
             The artifacts, external artifacts, model version artifacts/metadata and parameters for the step.
         """
         from zenml.artifacts.external_artifact import ExternalArtifact
+        from zenml.metadata.lazy_load import LazyRunMetadataResponse
         from zenml.model.lazy_load import ModelVersionDataLazyLoader
         from zenml.models.v2.core.artifact_version import (
             ArtifactVersionResponse,
             LazyArtifactVersionResponse,
         )
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         signature = inspect.signature(self.entrypoint, follow_wrapped=True)
 

zenml/steps/entrypoint_function_utils.py

@@ -32,6 +32,7 @@ from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.lazy_load import LazyRunMetadataResponse
 from zenml.steps.utils import (
     OutputSignature,
     parse_return_type_annotations,
@@ -136,10 +137,7 @@ class EntrypointFunctionDefinition(NamedTuple):
             UnmaterializedArtifact,
         )
         from zenml.client_lazy_loader import ClientLazyLoader
-        from zenml.models import (
-            ArtifactVersionResponse,
-            RunMetadataResponse,
-        )
+        from zenml.models import ArtifactVersionResponse
 
         if key not in self.inputs:
             raise KeyError(
@@ -154,8 +152,8 @@ class EntrypointFunctionDefinition(NamedTuple):
                 StepArtifact,
                 ExternalArtifact,
                 ArtifactVersionResponse,
-                RunMetadataResponse,
                 ClientLazyLoader,
+                LazyRunMetadataResponse,
             ),
         ):
             # If we were to do any type validation for artifacts here, we

zenml/steps/step_context.py

@@ -35,11 +35,12 @@ if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType
     from zenml.model.model import Model
     from zenml.models import (
-        ArtifactVersionResponse,
         PipelineResponse,
         PipelineRunResponse,
         StepRunResponse,
     )
+    from zenml.models.v2.core.step_run import StepRunInputResponse
+
 
 logger = get_logger(__name__)
 
@@ -191,7 +192,7 @@ class StepContext(metaclass=SingletonMetaClass):
         return self.model_version.to_model_class()
 
     @property
-    def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def inputs(self) -> Dict[str, "StepRunInputResponse"]:
         """Returns the input artifacts of the current step.
 
         Returns:

zenml/steps/utils.py

@@ -26,7 +26,11 @@ from typing_extensions import Annotated
 
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
-from zenml.enums import ExecutionStatus, MetadataResourceTypes
+from zenml.enums import (
+    ArtifactSaveType,
+    ExecutionStatus,
+    MetadataResourceTypes,
+)
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
@@ -547,8 +551,10 @@ def run_as_single_step_pipeline(
     # 4. Load output artifacts
     step_run = next(iter(run.steps.values()))
     outputs = [
-        step_run.outputs[output_name].load()
+        artifact_version.load()
         for output_name in step_run.config.outputs.keys()
+        for artifact_version in step_run.outputs[output_name]
+        if artifact_version.save_type == ArtifactSaveType.STEP_OUTPUT
     ]
 
     if len(outputs) == 0:

zenml/zen_server/rbac/utils.py

@@ -404,7 +404,6 @@ def get_resource_type_for_model(
         PipelineDeploymentResponse,
         PipelineResponse,
         PipelineRunResponse,
-        RunMetadataResponse,
         RunTemplateResponse,
         SecretResponse,
         ServiceAccountResponse,
@@ -437,7 +436,6 @@ def get_resource_type_for_model(
         ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
         WorkspaceResponse: ResourceType.WORKSPACE,
         UserResponse: ResourceType.USER,
-        RunMetadataResponse: ResourceType.RUN_METADATA,
         PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,

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

@@ -201,9 +201,7 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -2014,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:
@@ -2026,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 -----------------------------