Extend pipeline in_progress checks

Author: Json-Andriopoulos

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

@@ -239,6 +239,7 @@ class StepRunResponseBody(ProjectScopedResponseBody):
         title="The substitutions of the step run.",
         default={},
     )
+    cached_heartbeat_threshold: Optional[int] = Field(title="", default=None)
     model_config = ConfigDict(protected_namespaces=())
 
 
@@ -609,6 +610,15 @@ def latest_heartbeat(self) -> Optional[datetime]:
         """
         return self.get_body().latest_heartbeat
 
+    @property
+    def cached_heartbeat_threshold(self) -> Optional[int]:
+        """The `cached_heartbeat_threshold` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().cached_heartbeat_threshold
+
     @property
     def logs(self) -> Optional["LogsResponse"]:
         """The `logs` property.

src/zenml/steps/heartbeat.py

@@ -22,6 +22,7 @@
 from uuid import UUID
 
 from zenml.enums import ExecutionStatus
+from zenml.utils.time_utils import to_local_tz
 
 if TYPE_CHECKING:
     from zenml.models import StepRunResponse
@@ -173,36 +174,64 @@ def _heartbeat(self) -> None:
             self._terminated = True
 
 
-def is_heartbeat_unhealthy(step_run: "StepRunResponse") -> bool:
+def cached_is_heartbeat_unhealthy(
+    step_run_id: UUID,
+    status: ExecutionStatus,
+    latest_heartbeat: datetime | None,
+    start_time: datetime | None = None,
+    heartbeat_threshold: int | None = None,
+) -> bool:
     """Utility function - Checks if step heartbeats indicate un-healthy execution.
 
     Args:
-        step_run: Information regarding a step run.
+        step_run_id: The run step id.
+        status: The run step status.
+        latest_heartbeat: The run step latest heartbeat.
+        start_time: The run step start time.
+        heartbeat_threshold: If heartbeat enabled the max minutes without heartbeat
+            for healthy, running tasks.
 
     Returns:
         True if the step heartbeat is unhealthy, False otherwise.
     """
-    if not step_run.spec.enable_heartbeat:
+    if not heartbeat_threshold:
         return False
 
-    if step_run.status.is_finished:
+    if status.is_finished:
         return False
 
-    if step_run.latest_heartbeat:
-        heartbeat_diff = (
-            datetime.now(tz=timezone.utc) - step_run.latest_heartbeat
+    if latest_heartbeat:
+        heartbeat_diff = datetime.now(tz=timezone.utc) - to_local_tz(
+            latest_heartbeat
+        )
+    elif start_time:
+        heartbeat_diff = datetime.now(tz=timezone.utc) - to_local_tz(
+            start_time
         )
-    elif step_run.start_time:
-        heartbeat_diff = datetime.now(tz=timezone.utc) - step_run.start_time
     else:
         return False
 
-    logger.info("%s heartbeat diff=%s", step_run.name, heartbeat_diff)
+    logger.debug("Step %s heartbeat diff=%s", step_run_id, heartbeat_diff)
 
-    if (
-        heartbeat_diff.total_seconds()
-        > step_run.config.heartbeat_healthy_threshold * 60
-    ):
+    if heartbeat_diff.total_seconds() > heartbeat_threshold * 60:
         return True
 
     return False
+
+
+def is_heartbeat_unhealthy(step_run: "StepRunResponse") -> bool:
+    """Utility function - Checks if step heartbeats indicate un-healthy execution.
+
+    Args:
+        step_run: Information regarding a step run.
+
+    Returns:
+        True if the step heartbeat is unhealthy, False otherwise.
+    """
+    return cached_is_heartbeat_unhealthy(
+        step_run_id=step_run.id,
+        status=step_run.status,
+        start_time=step_run.start_time,
+        heartbeat_threshold=step_run.cached_heartbeat_threshold,
+        latest_heartbeat=step_run.latest_heartbeat,
+    )

src/zenml/zen_stores/migrations/versions/2837a7ab3d77_add_cached_heartbeat_threshold_column.py

@@ -0,0 +1,32 @@
+"""Add cached heartbeat threshold column [2837a7ab3d77].
+
+Revision ID: 2837a7ab3d77
+Revises: d203788f82b9
+Create Date: 2025-11-28 12:27:14.341553
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2837a7ab3d77"
+down_revision = "d203788f82b9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("step_run", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "cached_heartbeat_threshold", sa.Integer(), nullable=True
+            )
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("step_run", schema=None) as batch_op:
+        batch_op.drop_column("cached_heartbeat_threshold")

src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -20,7 +20,7 @@
 
 from pydantic import ConfigDict
 from sqlalchemy import UniqueConstraint
-from sqlalchemy.orm import object_session, selectinload
+from sqlalchemy.orm import Session, object_session, selectinload
 from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import TEXT, Column, Field, Relationship, select
 
@@ -815,6 +815,30 @@ def is_placeholder_run(self) -> bool:
             ExecutionStatus.PROVISIONING.value,
         }
 
+    @staticmethod
+    def _step_in_progress(step_id: UUID, session: Session) -> bool:
+        from zenml.steps.heartbeat import cached_is_heartbeat_unhealthy
+
+        step_run = session.execute(
+            select(StepRunSchema).where(StepRunSchema.id == step_id)
+        ).scalar_one_or_none()
+
+        if not step_run:
+            return False
+
+        status: ExecutionStatus = ExecutionStatus(step_run.status)
+
+        return not (
+            status.is_finished
+            or cached_is_heartbeat_unhealthy(
+                step_run_id=step_run.id,
+                status=status,
+                heartbeat_threshold=step_run.cached_heartbeat_threshold,
+                start_time=step_run.start_time,
+                latest_heartbeat=step_run.latest_heartbeat,
+            )
+        )
+
     def _check_if_run_in_progress(self) -> bool:
         """Checks whether the run is in progress.
 
@@ -843,19 +867,25 @@ def _check_if_run_in_progress(self) -> bool:
 
                 if session := object_session(self):
                     step_run_statuses = session.execute(
-                        select(StepRunSchema.name, StepRunSchema.status).where(
-                            StepRunSchema.pipeline_run_id == self.id
-                        )
+                        select(
+                            StepRunSchema.id,
+                            StepRunSchema.name,
+                            StepRunSchema.status,
+                        ).where(StepRunSchema.pipeline_run_id == self.id)
                     ).all()
 
                     if self.snapshot and self.snapshot.pipeline_spec:
                         step_dict = self.get_upstream_steps()
 
                         dag = build_dag(step_dict)
 
+                        step_name_to_id = {
+                            name: id_ for id_, name, _ in step_run_statuses
+                        }
+
                         failed_steps = {
                             name
-                            for name, status in step_run_statuses
+                            for _, name, status in step_run_statuses
                             if ExecutionStatus(status).is_failed
                         }
 
@@ -874,12 +904,21 @@ def _check_if_run_in_progress(self) -> bool:
 
                         for step_name, _ in step_dict.items():
                             if step_name in steps_to_skip:
+                                # failed steps downstream
+                                continue
+
+                            if steps_statuses[step_name].is_finished:
+                                # completed steps
                                 continue
 
                             if step_name not in steps_statuses:
+                                # steps that haven't started yet
                                 return True
 
-                            elif not steps_statuses[step_name].is_finished:
+                            if self._step_in_progress(
+                                step_name_to_id[step_name], session=session
+                            ):
+                                # running steps without unhealthy heartbeats
                                 return True
 
                         return False

src/zenml/zen_stores/schemas/step_run_schemas.py

@@ -117,7 +117,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
             nullable=True,
         )
     )
-
+    cached_heartbeat_threshold: Optional[int] = Field(nullable=True)
     # Foreign keys
     original_step_run_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
@@ -427,6 +427,7 @@ def to_model(
             start_time=self.start_time,
             end_time=self.end_time,
             latest_heartbeat=self.latest_heartbeat,
+            cached_heartbeat_threshold=self.cached_heartbeat_threshold,
             created=self.created,
             updated=self.updated,
             model_version_id=self.model_version_id,

src/zenml/zen_stores/sql_zen_store.py

@@ -10095,6 +10095,13 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                 is_retriable=is_retriable,
             )
 
+            # cached top-level heartbeat config property (for fast validation).
+            step_schema.cached_heartbeat_threshold = (
+                step_config.config.heartbeat_healthy_threshold
+                if step_config.spec.enable_heartbeat
+                else None
+            )
+
             session.add(step_schema)
             try:
                 session.commit()

tests/unit/steps/test_heartbeat_worker.py

@@ -140,6 +140,7 @@ def test_heartbeat_healthiness_check(monkeypatch):
             is_retriable=False,
             start_time=None,
             latest_heartbeat=None,
+            cached_heartbeat_threshold=5,
         ),
         metadata=StepRunResponseMetadata(
             config=StepConfiguration(
@@ -189,15 +190,5 @@ def now(cls, tz=None):
         assert not heartbeat.is_heartbeat_unhealthy(step_run)
 
         # if step heartbeat not enabled = healthy (default response)
-        step_run.metadata.config = StepConfiguration(
-            name="test", heartbeat_healthy_threshold=5
-        )
-        assert heartbeat.is_heartbeat_unhealthy(step_run)
-        step_run.metadata.spec = StepSpec(
-            enable_heartbeat=False,
-            source=Source(module="test", type=SourceType.BUILTIN),
-            upstream_steps=["test"],
-            inputs={"test": InputSpec(step_name="test", output_name="test")},
-            invocation_id="test",
-        )
+        step_run.body.cached_heartbeat_threshold = None
         assert not heartbeat.is_heartbeat_unhealthy(step_run)