Fix potential race condition during run status update (#3720)

* Fix potential race condition during run status update

* Release lock earlier when not changing status

* Linting

Author: schustmi

src/zenml/zen_stores/sql_zen_store.py

@@ -8925,6 +8925,7 @@ def _update_pipeline_run_status(
 
         pipeline_run = session.exec(
             select(PipelineRunSchema)
+            .with_for_update()
             .options(
                 joinedload(
                     jl_arg(PipelineRunSchema.deployment), innerjoin=True
@@ -8950,96 +8951,100 @@ def _update_pipeline_run_status(
             num_steps=num_steps,
         )
 
-        if pipeline_run.is_placeholder_run() and not new_status.is_finished:
-            # If the pipeline run is a placeholder run, no step has been started
-            # for the run yet. This means the orchestrator hasn't started
+        if new_status == pipeline_run.status or (
+            pipeline_run.is_placeholder_run() and not new_status.is_finished
+        ):
+            # The status hasn't changed -> no need to update the status.
+            # If the pipeline run is a placeholder run (=no step has been started
+            # for the run yet), this means the orchestrator hasn't started
             # running yet, and this method is most likely being called as
             # part of the creation of some cached steps. In this case, we don't
             # update the status unless the run is finished.
+
+            # Commit so that we release the lock on the pipeline run.
+            session.commit()
             return
 
-        if new_status != pipeline_run.status:
-            run_update = PipelineRunUpdate(status=new_status)
-            if new_status in {
-                ExecutionStatus.COMPLETED,
-                ExecutionStatus.FAILED,
-            }:
-                run_update.end_time = utc_now()
-                if pipeline_run.start_time and isinstance(
-                    pipeline_run.start_time, datetime
-                ):
-                    duration_time = (
-                        run_update.end_time - pipeline_run.start_time
-                    )
-                    duration_seconds = duration_time.total_seconds()
-                    start_time_str = pipeline_run.start_time.strftime(
+        run_update = PipelineRunUpdate(status=new_status)
+        if new_status.is_finished:
+            run_update.end_time = utc_now()
+
+        pipeline_run.update(run_update)
+        session.add(pipeline_run)
+        # Commit so that we release the lock on the pipeline run.
+        session.commit()
+
+        if new_status.is_finished:
+            assert run_update.end_time
+            if pipeline_run.start_time:
+                duration_time = run_update.end_time - pipeline_run.start_time
+                duration_seconds = duration_time.total_seconds()
+                start_time_str = pipeline_run.start_time.strftime(
+                    "%Y-%m-%dT%H:%M:%S.%fZ"
+                )
+            else:
+                start_time_str = None
+                duration_seconds = None
+
+            stack = pipeline_run.deployment.stack
+            assert stack
+            stack_metadata = {
+                str(component.type): component.flavor
+                for component in stack.components
+            }
+            with track_handler(
+                AnalyticsEvent.RUN_PIPELINE_ENDED
+            ) as analytics_handler:
+                analytics_handler.metadata = {
+                    "project_id": pipeline_run.project_id,
+                    "pipeline_run_id": pipeline_run_id,
+                    "template_id": pipeline_run.deployment.template_id,
+                    "status": new_status,
+                    "num_steps": num_steps,
+                    "start_time": start_time_str,
+                    "end_time": run_update.end_time.strftime(
                         "%Y-%m-%dT%H:%M:%S.%fZ"
-                    )
-                else:
-                    start_time_str = None
-                    duration_seconds = None
-
-                stack = pipeline_run.deployment.stack
-                assert stack
-                stack_metadata = {
-                    str(component.type): component.flavor
-                    for component in stack.components
+                    ),
+                    "duration_seconds": duration_seconds,
+                    **stack_metadata,
                 }
-                with track_handler(
-                    AnalyticsEvent.RUN_PIPELINE_ENDED
-                ) as analytics_handler:
-                    analytics_handler.metadata = {
-                        "project_id": pipeline_run.project_id,
-                        "pipeline_run_id": pipeline_run_id,
-                        "template_id": pipeline_run.deployment.template_id,
-                        "status": new_status,
-                        "num_steps": num_steps,
-                        "start_time": start_time_str,
-                        "end_time": run_update.end_time.strftime(
-                            "%Y-%m-%dT%H:%M:%S.%fZ"
-                        ),
-                        "duration_seconds": duration_seconds,
-                        **stack_metadata,
-                    }
 
-                completed_onboarding_steps: Set[str] = {
-                    OnboardingStep.PIPELINE_RUN,
-                    OnboardingStep.STARTER_SETUP_COMPLETED,
-                }
-                if stack_metadata["orchestrator"] not in {
-                    "local",
-                    "local_docker",
-                }:
-                    completed_onboarding_steps.update(
-                        {
-                            OnboardingStep.PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR,
-                        }
+            completed_onboarding_steps: Set[str] = {
+                OnboardingStep.PIPELINE_RUN,
+                OnboardingStep.STARTER_SETUP_COMPLETED,
+            }
+            if stack_metadata["orchestrator"] not in {
+                "local",
+                "local_docker",
+            }:
+                completed_onboarding_steps.update(
+                    {
+                        OnboardingStep.PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR,
+                    }
+                )
+            if stack_metadata["artifact_store"] != "local":
+                completed_onboarding_steps.update(
+                    {
+                        OnboardingStep.PIPELINE_RUN_WITH_REMOTE_ARTIFACT_STORE,
+                        OnboardingStep.PRODUCTION_SETUP_COMPLETED,
+                    }
+                )
+            if OnboardingStep.THIRD_PIPELINE_RUN not in (
+                self._cached_onboarding_state or {}
+            ):
+                onboarding_state = self.get_onboarding_state()
+                if OnboardingStep.PIPELINE_RUN in onboarding_state:
+                    completed_onboarding_steps.add(
+                        OnboardingStep.SECOND_PIPELINE_RUN
                     )
-                if stack_metadata["artifact_store"] != "local":
-                    completed_onboarding_steps.update(
-                        {
-                            OnboardingStep.PIPELINE_RUN_WITH_REMOTE_ARTIFACT_STORE,
-                            OnboardingStep.PRODUCTION_SETUP_COMPLETED,
-                        }
+                if OnboardingStep.SECOND_PIPELINE_RUN in onboarding_state:
+                    completed_onboarding_steps.add(
+                        OnboardingStep.THIRD_PIPELINE_RUN
                     )
-                if OnboardingStep.THIRD_PIPELINE_RUN not in (
-                    self._cached_onboarding_state or {}
-                ):
-                    onboarding_state = self.get_onboarding_state()
-                    if OnboardingStep.PIPELINE_RUN in onboarding_state:
-                        completed_onboarding_steps.add(
-                            OnboardingStep.SECOND_PIPELINE_RUN
-                        )
-                    if OnboardingStep.SECOND_PIPELINE_RUN in onboarding_state:
-                        completed_onboarding_steps.add(
-                            OnboardingStep.THIRD_PIPELINE_RUN
-                        )
 
-                self._update_onboarding_state(
-                    completed_steps=completed_onboarding_steps, session=session
-                )
-            pipeline_run.update(run_update)
-            session.add(pipeline_run)
+            self._update_onboarding_state(
+                completed_steps=completed_onboarding_steps, session=session
+            )
 
     # --------------------------- Triggers ---------------------------