Server-side parent step computation (#3762)

Author: schustmi

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

@@ -129,6 +129,7 @@ class StepRunRequest(ProjectScopedRequest):
     parent_step_ids: List[UUID] = Field(
         title="The IDs of the parent steps of this step run.",
         default_factory=list,
+        deprecated=True,
     )
     inputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the input artifact versions of the step run.",

src/zenml/orchestrators/input_utils.py

@@ -14,8 +14,7 @@
 """Utilities for inputs."""
 
 import json
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
-from uuid import UUID
+from typing import TYPE_CHECKING, Dict, Optional
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
@@ -32,7 +31,7 @@ def resolve_step_inputs(
     step: "Step",
     pipeline_run: "PipelineRunResponse",
     step_runs: Optional[Dict[str, "StepRunResponse"]] = None,
-) -> Tuple[Dict[str, "StepRunInputResponse"], List[UUID]]:
+) -> Dict[str, "StepRunInputResponse"]:
     """Resolves inputs for the current step.
 
     Args:
@@ -49,16 +48,14 @@ def resolve_step_inputs(
             resolved in runtime due to missing object.
 
     Returns:
-        The IDs of the input artifact versions and the IDs of parent steps of
-            the current step.
+        The input artifact versions.
     """
     from zenml.models import ArtifactVersionResponse
     from zenml.models.v2.core.step_run import StepRunInputResponse
 
     step_runs = step_runs or {}
 
-    steps_to_fetch = set(step.spec.upstream_steps)
-    steps_to_fetch.update(
+    steps_to_fetch = set(
         input_.step_name for input_ in step.spec.inputs.values()
     )
     # Remove all the step runs that we've already fetched.
@@ -205,9 +202,4 @@ def resolve_step_inputs(
         else:
             step.config.parameters[name] = value_
 
-    parent_step_ids = [
-        step_runs[upstream_step].id
-        for upstream_step in step.spec.upstream_steps
-    ]
-
-    return input_artifacts, parent_step_ids
+    return input_artifacts

src/zenml/orchestrators/step_run_utils.py

@@ -95,7 +95,7 @@ def populate_request(
         """
         step = self.deployment.step_configurations[request.name]
 
-        input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
+        input_artifacts = input_utils.resolve_step_inputs(
             step=step,
             pipeline_run=self.pipeline_run,
             step_runs=step_runs,
@@ -108,7 +108,6 @@ def populate_request(
         request.inputs = {
             name: [artifact.id] for name, artifact in input_artifacts.items()
         }
-        request.parent_step_ids = parent_step_ids
 
         cache_key = cache_utils.generate_cache_key(
             step=step,

src/zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -90,7 +90,7 @@ def _run_step(
         )
 
         stack = Client().active_stack
-        input_artifacts, _ = input_utils.resolve_step_inputs(
+        input_artifacts = input_utils.resolve_step_inputs(
             step=step, pipeline_run=pipeline_run
         )
         output_artifact_uris = output_utils.prepare_output_artifact_uris(

src/zenml/zen_stores/sql_zen_store.py

@@ -8779,19 +8779,18 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                 session.commit()
                 session.refresh(step_schema)
 
-            # Save parent step IDs into the database.
-            for parent_step_id in step_run.parent_step_ids:
-                self._set_run_step_parent_step(
-                    child_step_run=step_schema,
-                    parent_id=parent_step_id,
-                    session=session,
-                )
-
             session.commit()
             session.refresh(step_schema)
 
             step_model = step_schema.to_model(include_metadata=True)
 
+            for upstream_step in step_model.spec.upstream_steps:
+                self._set_run_step_parent_step(
+                    child_step_run=step_schema,
+                    parent_step_name=upstream_step,
+                    session=session,
+                )
+
             # Save input artifact IDs into the database.
             for input_name, artifact_version_ids in step_run.inputs.items():
                 for artifact_version_id in artifact_version_ids:
@@ -9047,22 +9046,33 @@ def _get_step_run_input_type_from_config(
             return StepRunInputArtifactType.MANUAL
 
     def _set_run_step_parent_step(
-        self, child_step_run: StepRunSchema, parent_id: UUID, session: Session
+        self,
+        child_step_run: StepRunSchema,
+        parent_step_name: str,
+        session: Session,
     ) -> None:
         """Sets the parent step run for a step run.
 
         Args:
             child_step_run: The child step run to set the parent for.
-            parent_id: The ID of the parent step run to set a child for.
+            parent_step_name: The name of the parent step run to set a child for.
             session: The database session to use.
+
+        Raises:
+            RuntimeError: If the parent step run is not found.
         """
-        parent_step_run = self._get_reference_schema_by_id(
-            resource=child_step_run,
-            reference_schema=StepRunSchema,
-            reference_id=parent_id,
-            session=session,
-            reference_type="parent step",
-        )
+        parent_step_run = session.exec(
+            select(StepRunSchema)
+            .where(StepRunSchema.name == parent_step_name)
+            .where(
+                StepRunSchema.pipeline_run_id == child_step_run.pipeline_run_id
+            )
+        ).first()
+        if parent_step_run is None:
+            raise RuntimeError(
+                f"Parent step run `{parent_step_name}` not found for step run "
+                f"`{child_step_run.name}`."
+            )
 
         # Check if the parent step is already set.
         assignment = session.exec(

tests/unit/orchestrators/test_input_utils.py

@@ -59,7 +59,7 @@ def test_input_resolution(
         }
     )
 
-    input_artifacts, parent_ids = input_utils.resolve_step_inputs(
+    input_artifacts = input_utils.resolve_step_inputs(
         step=step, pipeline_run=sample_pipeline_run
     )
     assert input_artifacts == {
@@ -68,7 +68,6 @@ def test_input_resolution(
             **sample_artifact_version_model.model_dump(),
         )
     }
-    assert parent_ids == [step_run.id]
 
 
 def test_input_resolution_with_missing_step_run(mocker, sample_pipeline_run):