From 1e957abeb6026e5c6904113f8f938ed44a3ea4f7 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Thu, 4 Jan 2024 09:27:43 +0100 Subject: [PATCH] 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 --- src/zenml/zen_stores/schemas/pipeline_run_schemas.py | 6 ++---- src/zenml/zen_stores/sql_zen_store.py | 5 +++-- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index b8b0b59869..be6b0d4f03 100644 --- 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 index c566e75ad1..aca8d3f855 100644 --- 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 -----------------------------