diff --git a/src/zenml/constants.py b/src/zenml/constants.py index cfdb1cd3477..8102a1508b4 100644 --- a/src/zenml/constants.py +++ b/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) diff --git a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py index 1122ac14814..9afe6817f75 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +++ b/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: diff --git a/src/zenml/pipelines/dynamic/entrypoint_configuration.py b/src/zenml/pipelines/dynamic/entrypoint_configuration.py index f477f938157..d823daca0ce 100644 --- a/src/zenml/pipelines/dynamic/entrypoint_configuration.py +++ b/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 diff --git a/tests/unit/pipelines/dynamic/test_entrypoint_configuration.py b/tests/unit/pipelines/dynamic/test_entrypoint_configuration.py new file mode 100644 index 00000000000..454ddfff1c5 --- /dev/null +++ b/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()