Better input artifacts typing

src/zenml/artifacts/utils.py

@@ -387,17 +387,6 @@ def load_artifact(
         The loaded artifact.
     """
     artifact = Client().get_artifact_version(name_or_id, version)
-    try:
-        step_run = get_step_context().step_run
-        client = Client()
-        client.zen_store.update_run_step(
-            step_run_id=step_run.id,
-            step_run_update=StepRunUpdate(
-                loaded_artifact_versions={artifact.name: artifact.id}
-            ),
-        )
-    except RuntimeError:
-        pass  # Cannot link to step run if called outside of a step
     return load_artifact_from_response(artifact)
 
 

src/zenml/client.py

@@ -188,6 +188,7 @@
     WorkspaceResponse,
     WorkspaceUpdate,
 )
+from zenml.models.v2.core.step_run import StepRunUpdate
 from zenml.services.service import ServiceConfig
 from zenml.services.service_status import ServiceState
 from zenml.services.service_type import ServiceType
@@ -4164,20 +4165,35 @@ def get_artifact_version(
         Returns:
             The artifact version.
         """
+        from zenml import get_step_context
+
         if cll := client_lazy_loader(
             method_name="get_artifact_version",
             name_id_or_prefix=name_id_or_prefix,
             version=version,
             hydrate=hydrate,
         ):
             return cll  # type: ignore[return-value]
-        return self._get_entity_version_by_id_or_name_or_prefix(
+
+        artifact = self._get_entity_version_by_id_or_name_or_prefix(
             get_method=self.zen_store.get_artifact_version,
             list_method=self.list_artifact_versions,
             name_id_or_prefix=name_id_or_prefix,
             version=version,
             hydrate=hydrate,
         )
+        try:
+            step_run = get_step_context().step_run
+            client = Client()
+            client.zen_store.update_run_step(
+                step_run_id=step_run.id,
+                step_run_update=StepRunUpdate(
+                    loaded_artifact_versions={artifact.name: artifact.id}
+                ),
+            )
+        except RuntimeError:
+            pass  # Cannot link to step run if called outside of a step
+        return artifact
 
     def list_artifact_versions(
         self,

src/zenml/enums.py

@@ -34,8 +34,12 @@ class ArtifactType(StrEnum):
 class StepRunInputArtifactType(StrEnum):
     """All possible types of a step run input artifact."""
 
-    DEFAULT = "default"  # input argument that is the output of a previous step
+    STEP_OUTPUT = (
+        "step_output"  # input argument that is the output of a previous step
+    )
     MANUAL = "manual"  # manually loaded via `zenml.load_artifact()`
+    EXTERNAL = "external"  # loaded via `ExternalArtifact(value=...)`
+    LAZY_LOADED = "lazy"  # loaded via various lazy methods
 
 
 class ArtifactSaveType(StrEnum):

src/zenml/models/v2/core/step_run.py

@@ -21,7 +21,7 @@
 
 from zenml.config.step_configurations import StepConfiguration, StepSpec
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ExecutionStatus
+from zenml.enums import ExecutionStatus, StepRunInputArtifactType
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
@@ -31,18 +31,37 @@
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
+from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
-    from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.logs import (
         LogsRequest,
         LogsResponse,
     )
 
 
+class StepRunInputResponse(ArtifactVersionResponse):
+    """Response model for step run inputs."""
+
+    input_type: StepRunInputArtifactType
+
+    def get_hydrated_version(self) -> "StepRunInputResponse":
+        """Get the hydrated version of this step run input.
+
+        Returns:
+            an instance of the same entity with the metadata field attached.
+        """
+        from zenml.client import Client
+
+        return StepRunInputResponse(
+            input_type=self.input_type,
+            **Client().zen_store.get_artifact_version(self.id).model_dump(),
+        )
+
+
 # ------------------ Request Model ------------------
 
 
@@ -160,11 +179,11 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
         title="The end time of the step run.",
         default=None,
     )
-    inputs: Dict[str, "ArtifactVersionResponse"] = Field(
+    inputs: Dict[str, StepRunInputResponse] = Field(
         title="The input artifact versions of the step run.",
         default_factory=dict,
     )
-    outputs: Dict[str, List["ArtifactVersionResponse"]] = Field(
+    outputs: Dict[str, List[ArtifactVersionResponse]] = Field(
         title="The output artifact versions of the step run.",
         default_factory=dict,
     )
@@ -268,7 +287,7 @@ def get_hydrated_version(self) -> "StepRunResponse":
 
     # Helper properties
     @property
-    def input(self) -> "ArtifactVersionResponse":
+    def input(self) -> ArtifactVersionResponse:
         """Returns the input artifact that was used to run this step.
 
         Returns:
@@ -287,7 +306,7 @@ def input(self) -> "ArtifactVersionResponse":
         return next(iter(self.inputs.values()))
 
     @property
-    def output(self) -> "ArtifactVersionResponse":
+    def output(self) -> ArtifactVersionResponse:
         """Returns the output artifact that was written by this step.
 
         Returns:
@@ -319,7 +338,7 @@ def status(self) -> ExecutionStatus:
         return self.get_body().status
 
     @property
-    def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def inputs(self) -> Dict[str, StepRunInputResponse]:
         """The `inputs` property.
 
         Returns:
@@ -328,7 +347,7 @@ def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
         return self.get_body().inputs
 
     @property
-    def outputs(self) -> Dict[str, List["ArtifactVersionResponse"]]:
+    def outputs(self) -> Dict[str, List[ArtifactVersionResponse]]:
         """The `outputs` property.
 
         Returns:

src/zenml/orchestrators/input_utils.py

@@ -18,18 +18,19 @@
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
-from zenml.enums import ArtifactSaveType
+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:
@@ -47,6 +48,7 @@ def resolve_step_inputs(
             the current step.
     """
     from zenml.models import ArtifactVersionResponse
+    from zenml.models.v2.core.step_run import StepRunInputResponse
 
     current_run_steps = {
         run_step.name: run_step
@@ -55,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]
@@ -90,15 +92,19 @@ def resolve_step_inputs(
                 f"`{input_.step_name}`."
             )
 
-        input_artifacts[name] = step_outputs[0]
+        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():
@@ -129,7 +135,10 @@ 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:
@@ -156,7 +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_
+            input_artifacts[name] = StepRunInputResponse(
+                input_type=StepRunInputArtifactType.LAZY_LOADED,
+                **value_.model_dump(),
+            )
         else:
             step.config.parameters[name] = value_
 

src/zenml/orchestrators/step_launcher.py

@@ -33,13 +33,13 @@
 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 @@ def _run_step_without_step_operator(
         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:

src/zenml/orchestrators/step_run_utils.py

@@ -104,6 +104,7 @@ def populate_request(self, request: StepRunRequest) -> None:
             input_name: artifact.id
             for input_name, artifact in input_artifacts.items()
         }
+
         request.inputs = input_artifact_ids
         request.parent_step_ids = parent_step_ids
 

src/zenml/orchestrators/step_runner.py

@@ -42,6 +42,7 @@
 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,
@@ -100,7 +101,7 @@ def run(
         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:
@@ -306,7 +307,7 @@ def _parse_inputs(
         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.
 

src/zenml/steps/step_context.py

@@ -35,11 +35,12 @@
     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 @@ def model(self) -> "Model":
         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:

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

@@ -26,9 +26,10 @@ def upgrade() -> None:
     op.execute("""
         UPDATE artifact_version
         SET save_type = (
-            SELECT step_run_output_artifact.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("""
@@ -71,9 +72,10 @@ def downgrade() -> None:
     op.execute("""
         UPDATE step_run_output_artifact
         SET type = (
-            SELECT artifact_version.save_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("""