Improved Kubernetes dynamic pipeline retries

src/zenml/constants.py

@@ -315,6 +315,7 @@ def handle_float_env_var(var: str, default: float = 0.0) -> float:
     "ZENML_STEP_RUNS_FETCH_MAX_CHUNK_LENGTH"
 )
 ENV_ZENML_STREAM_PUBLISHER_BATCH_SIZE = "ZENML_STREAM_PUBLISHER_BATCH_SIZE"
+DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE = 45
 
 # Logging variables
 IS_DEBUG_ENV: bool = handle_bool_env_var(ENV_ZENML_DEBUG, default=False)

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -37,6 +37,7 @@
 from contextlib import contextmanager
 from typing import (
     TYPE_CHECKING,
+    Any,
     Dict,
     Generator,
     List,
@@ -54,6 +55,7 @@
 from zenml.client import Client
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import (
+    DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE,
     METADATA_ORCHESTRATOR_RUN_ID,
     ORCHESTRATOR_DOCKER_IMAGE_KEY,
 )
@@ -546,6 +548,7 @@ def _prepare_job_manifest(
         settings: KubernetesOrchestratorSettings,
         pod_settings: Optional[KubernetesPodSettings] = None,
         backoff_limit: Optional[int] = None,
+        pod_failure_policy: Optional[Dict[str, Any]] = None,
     ) -> k8s_client.V1Job:
         """Prepares the job manifest for a Kubernetes job.
 
@@ -560,6 +563,8 @@ def _prepare_job_manifest(
             settings: Component settings for the orchestrator.
             pod_settings: Optional settings for the pod.
             backoff_limit: The backoff limit for the job.
+            pod_failure_policy: Default pod failure policy for the job. The
+                `pod_failure_policy` setting takes precedence over this value.
 
         Returns:
             The job manifest.
@@ -592,34 +597,39 @@ def _prepare_job_manifest(
             termination_grace_period_seconds=settings.pod_stop_grace_period,
         )
 
-        pod_failure_policy = settings.pod_failure_policy or {
-            # These rules are applied sequentially. This means any failure in
-            # the main container will count towards the max retries. Any other
-            # disruption will not count towards the max retries.
-            "rules": [
-                # If the main container fails, we count it towards the max
+        pod_failure_policy = (
+            settings.pod_failure_policy
+            or pod_failure_policy
+            or {
+                # These rules are applied sequentially. This means any failure
+                # in the main container will count towards the max retries.
+                # Any other disruption will not count towards the max
                 # retries.
-                {
-                    "action": "Count",
-                    "onExitCodes": {
-                        "containerName": "main",
-                        "operator": "NotIn",
-                        "values": [0],
+                "rules": [
+                    # If the main container fails, we count it towards the
+                    # max retries.
+                    {
+                        "action": "Count",
+                        "onExitCodes": {
+                            "containerName": "main",
+                            "operator": "NotIn",
+                            "values": [0],
+                        },
                     },
-                },
-                # If the pod is interrupted at any other time, we don't count
-                # it as a retry
-                {
-                    "action": "Ignore",
-                    "onPodConditions": [
-                        {
-                            "type": "DisruptionTarget",
-                            "status": "True",
-                        }
-                    ],
-                },
-            ]
-        }
+                    # If the pod is interrupted at any other time, we don't
+                    # count it as a retry
+                    {
+                        "action": "Ignore",
+                        "onPodConditions": [
+                            {
+                                "type": "DisruptionTarget",
+                                "status": "True",
+                            }
+                        ],
+                    },
+                ]
+            }
+        )
 
         return build_job_manifest(
             job_name=name,
@@ -786,6 +796,44 @@ def _submit_orchestrator_job(
             settings, pipeline_name=snapshot.pipeline_configuration.name
         )
 
+        pod_failure_policy = {
+            # Rule order matters: a disrupted pod can also exit with code
+            # 137/143, so the DisruptionTarget rule must be evaluated before
+            # the exit-code rules below.
+            "rules": [
+                # If the pod is disrupted, we don't count it as a retry.
+                {
+                    "action": "Ignore",
+                    "onPodConditions": [
+                        {
+                            "type": "DisruptionTarget",
+                            "status": "True",
+                        }
+                    ],
+                },
+                # If the pod exits with the dedicated exit code, the run
+                # already reached a terminal status and retrying is useless.
+                {
+                    "action": "FailJob",
+                    "onExitCodes": {
+                        "containerName": "main",
+                        "operator": "In",
+                        "values": [DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE],
+                    },
+                },
+                # Any other failure of the main container counts towards the
+                # max retries.
+                {
+                    "action": "Count",
+                    "onExitCodes": {
+                        "containerName": "main",
+                        "operator": "NotIn",
+                        "values": [0],
+                    },
+                },
+            ]
+        }
+
         try:
             with self._create_auth_secret_if_necessary(
                 snapshot, environment, orchestrator_pod_settings
@@ -801,6 +849,7 @@ def _submit_orchestrator_job(
                     settings=settings,
                     pod_settings=orchestrator_pod_settings,
                     backoff_limit=settings.orchestrator_job_backoff_limit,
+                    pod_failure_policy=pod_failure_policy,
                 )
 
                 if snapshot.schedule:

src/zenml/pipelines/dynamic/entrypoint_configuration.py

@@ -13,15 +13,20 @@
 #  permissions and limitations under the License.
 """Entrypoint configuration to run a dynamic pipeline."""
 
+import sys
 from typing import Any, Dict, List
 from uuid import UUID
 
 from zenml.client import Client
+from zenml.constants import DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE
 from zenml.entrypoints.base_entrypoint_configuration import (
     BaseEntrypointConfiguration,
 )
 from zenml.execution.pipeline.dynamic.runner import DynamicPipelineRunner
 from zenml.integrations.registry import integration_registry
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
 
 RUN_ID_OPTION = "run_id"
 
@@ -57,7 +62,12 @@ def get_entrypoint_arguments(
         return args
 
     def run(self) -> None:
-        """Prepares the environment and runs the configured dynamic pipeline."""
+        """Prepares the environment and runs the configured dynamic pipeline.
+
+        Raises:
+            BaseException: If the pipeline run did not reach a terminal
+                status.
+        """
         snapshot = self.snapshot
 
         # Activate all the integrations. This makes sure that all materializers
@@ -71,4 +81,14 @@ def run(self) -> None:
             run = Client().get_pipeline_run(UUID(run_id))
 
         runner = DynamicPipelineRunner(snapshot=snapshot, run=run)
-        runner.run_pipeline()
+        try:
+            runner.run_pipeline()
+        except BaseException:
+            if runner.run.status.is_finished:
+                logger.error(
+                    "Pipeline run `%s` failed.",
+                    runner.run.id,
+                    exc_info=True,
+                )
+                sys.exit(DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE)
+            raise

tests/unit/pipelines/dynamic/test_entrypoint_configuration.py

@@ -0,0 +1,115 @@
+#  Copyright (c) ZenML GmbH 2026. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Tests for the dynamic pipeline entrypoint configuration."""
+
+from typing import Any
+
+import pytest
+
+from zenml.constants import DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE
+from zenml.enums import ExecutionStatus
+from zenml.models import PipelineRequest, PipelineSnapshotRequest
+from zenml.pipelines.dynamic.entrypoint_configuration import (
+    DynamicPipelineEntrypointConfiguration,
+)
+
+
+class _StubRun:
+    """Stub run exposing only the status the entrypoint configuration reads."""
+
+    def __init__(self, status: ExecutionStatus) -> None:
+        self.status = status
+        self.id = "stub-run-id"
+
+
+class _FailingStubRunner:
+    """Stub runner whose `run_pipeline` always raises."""
+
+    def __init__(self, status: ExecutionStatus, **kwargs: Any) -> None:
+        self.run = _StubRun(status=status)
+
+    def run_pipeline(self) -> None:
+        raise RuntimeError("boom")
+
+
+def _create_snapshot(clean_client: Any) -> Any:
+    pipeline = clean_client.zen_store.create_pipeline(
+        PipelineRequest(
+            name="pipeline",
+            project=clean_client.active_project.id,
+        )
+    )
+    request = PipelineSnapshotRequest(
+        user=clean_client.active_user.id,
+        project=clean_client.active_project.id,
+        run_name_template="",
+        pipeline_configuration={"name": "pipeline"},
+        stack=clean_client.active_stack.id,
+        client_version="0.12.3",
+        server_version="0.12.3",
+        pipeline=pipeline.id,
+    )
+    return clean_client.zen_store.create_snapshot(request)
+
+
+def _entrypoint_config_with_stub_runner(
+    clean_client: Any,
+    monkeypatch: pytest.MonkeyPatch,
+    status: ExecutionStatus,
+) -> DynamicPipelineEntrypointConfiguration:
+    snapshot = _create_snapshot(clean_client)
+
+    monkeypatch.setattr(
+        "zenml.pipelines.dynamic.entrypoint_configuration.DynamicPipelineRunner",
+        lambda **kwargs: _FailingStubRunner(status=status, **kwargs),
+    )
+    monkeypatch.setattr(
+        "zenml.pipelines.dynamic.entrypoint_configuration.integration_registry.activate_integrations",
+        lambda: None,
+    )
+    monkeypatch.setattr(
+        DynamicPipelineEntrypointConfiguration,
+        "prepare_code_environment",
+        lambda self: None,
+    )
+
+    return DynamicPipelineEntrypointConfiguration(
+        arguments=["--snapshot_id", str(snapshot.id)]
+    )
+
+
+def test_run_exits_with_dedicated_code_when_run_reached_terminal_status(
+    clean_client: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Tests that `run()` exits with the dedicated code once the run is terminal."""
+    entrypoint_config = _entrypoint_config_with_stub_runner(
+        clean_client, monkeypatch, status=ExecutionStatus.FAILED
+    )
+
+    with pytest.raises(SystemExit) as exc_info:
+        entrypoint_config.run()
+
+    assert exc_info.value.code == DYNAMIC_PIPELINE_RUN_FAILED_EXIT_CODE
+
+
+def test_run_reraises_when_run_did_not_reach_terminal_status(
+    clean_client: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Tests that `run()` re-raises the original exception when the run is not terminal."""
+    entrypoint_config = _entrypoint_config_with_stub_runner(
+        clean_client, monkeypatch, status=ExecutionStatus.RUNNING
+    )
+
+    with pytest.raises(RuntimeError, match="boom"):
+        entrypoint_config.run()