Refactor pipeline run updates (#2117)

* Refactor pipeline run updates * Auto-update of Starter template * Auto-update of E2E template * Auto-update of NLP template --------- Co-authored-by: GitHub Actions <actions@github.com>
zenml-io · Jan 4, 2024 · 1e957ab · 1e957ab
1 parent 0a9e655
commit 1e957ab
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py
@@ -223,8 +223,6 @@ def to_model(self, hydrate: bool = False) -> "PipelineRunResponse":
         }
 
         if self.deployment is not None:
-            steps = {s.name: s.to_model() for s in self.step_runs}
-
             deployment = self.deployment.to_model()
 
             config = deployment.pipeline_configuration
@@ -237,8 +235,6 @@ def to_model(self, hydrate: bool = False) -> "PipelineRunResponse":
             code_reference = deployment.code_reference
 
         elif self.pipeline_configuration is not None:
-            steps = {step.name: step.to_model() for step in self.step_runs}
-
             config = PipelineConfiguration.parse_raw(
                 self.pipeline_configuration
             )
@@ -274,6 +270,8 @@ def to_model(self, hydrate: bool = False) -> "PipelineRunResponse":
         )
         metadata = None
         if hydrate:
+            steps = {step.name: step.to_model() for step in self.step_runs}
+
             metadata = PipelineRunResponseMetadata(
                 workspace=self.workspace.to_model(),
                 run_metadata=run_metadata,

diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py
@@ -5453,13 +5453,14 @@ def _update_pipeline_run_status(
         )
 
         if new_status != pipeline_run.status:
-            pipeline_run.status = new_status
+            run_update = PipelineRunUpdate(status=new_status)
             if new_status in {
                 ExecutionStatus.COMPLETED,
                 ExecutionStatus.FAILED,
             }:
-                pipeline_run.end_time = datetime.utcnow()
+                run_update.end_time = datetime.utcnow()
 
+            pipeline_run.update(run_update)
             session.add(pipeline_run)
 
     # ----------------------------- Users -----------------------------