zenml-nightly 0.84.1.dev20250805__py3-none-any.whl → 0.84.1.dev20250806__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.84.1.dev20250805
+0.84.1.dev20250806

zenml/integrations/kubernetes/constants.py

@@ -0,0 +1,27 @@
+#  Copyright (c) ZenML GmbH 2025. 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.
+#
+# Parts of the `prepare_or_run_pipeline()` method of this file are
+# inspired by the Kubernetes dag runner implementation of tfx
+"""Kubernetes orchestrator constants."""
+
+ENV_ZENML_KUBERNETES_RUN_ID = "ZENML_KUBERNETES_RUN_ID"
+KUBERNETES_SECRET_TOKEN_KEY_NAME = "zenml_api_token"
+KUBERNETES_CRON_JOB_METADATA_KEY = "cron_job_name"
+# Annotation keys
+ORCHESTRATOR_ANNOTATION_KEY = "zenml.io/orchestrator"
+RUN_ID_ANNOTATION_KEY = "zenml.io/run-id"
+ORCHESTRATOR_RUN_ID_ANNOTATION_KEY = "zenml.io/orchestrator-run-id"
+STEP_NAME_ANNOTATION_KEY = "zenml.io/step-name"
+STEP_OPERATOR_ANNOTATION_KEY = "zenml.io/step-operator"

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -15,7 +15,13 @@
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
-from pydantic import Field, NonNegativeInt, PositiveInt, field_validator
+from pydantic import (
+    Field,
+    NonNegativeInt,
+    PositiveFloat,
+    PositiveInt,
+    field_validator,
+)
 
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
@@ -23,6 +29,7 @@ from zenml.integrations.kubernetes import KUBERNETES_ORCHESTRATOR_FLAVOR
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.models import ServiceConnectorRequirements
 from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
+from zenml.utils import deprecation_utils
 
 if TYPE_CHECKING:
     from zenml.integrations.kubernetes.orchestrators import (
@@ -42,16 +49,6 @@ class KubernetesOrchestratorSettings(BaseSettings):
         description="Whether to wait for all pipeline steps to complete. "
         "When `False`, the client returns immediately and execution continues asynchronously.",
     )
-    timeout: int = Field(
-        default=0,
-        description="Maximum seconds to wait for synchronous runs. Set to `0` for unlimited duration.",
-    )
-    stream_step_logs: bool = Field(
-        default=True,
-        description="If `True`, the orchestrator pod will stream the logs "
-        "of the step pods. This only has an effect if specified on the "
-        "pipeline, not on individual steps.",
-    )
     service_account_name: Optional[str] = Field(
         default=None,
         description="Kubernetes service account for the orchestrator pod. "
@@ -74,25 +71,9 @@ class KubernetesOrchestratorSettings(BaseSettings):
         default=None,
         description="Pod configuration for the orchestrator container that launches step pods.",
     )
-    pod_name_prefix: Optional[str] = Field(
+    job_name_prefix: Optional[str] = Field(
         default=None,
-        description="Custom prefix for generated pod names. Helps identify pods in the cluster.",
-    )
-    pod_startup_timeout: int = Field(
-        default=600,
-        description="Maximum seconds to wait for step pods to start. Default is 10 minutes.",
-    )
-    pod_failure_max_retries: int = Field(
-        default=3,
-        description="Maximum retry attempts when step pods fail to start.",
-    )
-    pod_failure_retry_delay: int = Field(
-        default=10,
-        description="Delay in seconds between pod failure retry attempts.",
-    )
-    pod_failure_backoff: float = Field(
-        default=1.0,
-        description="Exponential backoff factor for retry delays. Values > 1.0 increase delay with each retry.",
+        description="Prefix for the job name.",
     )
     max_parallelism: Optional[PositiveInt] = Field(
         default=None,
@@ -100,19 +81,20 @@ class KubernetesOrchestratorSettings(BaseSettings):
     )
     successful_jobs_history_limit: Optional[NonNegativeInt] = Field(
         default=None,
-        description="Number of successful scheduled jobs to retain in cluster history.",
+        description="Number of successful scheduled jobs to retain in history.",
     )
     failed_jobs_history_limit: Optional[NonNegativeInt] = Field(
         default=None,
-        description="Number of failed scheduled jobs to retain in cluster history.",
+        description="Number of failed scheduled jobs to retain in history.",
     )
     ttl_seconds_after_finished: Optional[NonNegativeInt] = Field(
         default=None,
-        description="Seconds to keep finished scheduled jobs before automatic cleanup.",
+        description="Seconds to keep finished jobs before automatic cleanup.",
     )
     active_deadline_seconds: Optional[NonNegativeInt] = Field(
         default=None,
-        description="Deadline in seconds for the active pod. If the pod is inactive for this many seconds, it will be terminated.",
+        description="Job deadline in seconds. If the job doesn't finish "
+        "within this time, it will be terminated.",
     )
     backoff_limit_margin: NonNegativeInt = Field(
         default=0,
@@ -131,6 +113,10 @@ class KubernetesOrchestratorSettings(BaseSettings):
         "the chance of the server receiving the maximum amount of retry "
         "requests.",
     )
+    orchestrator_job_backoff_limit: NonNegativeInt = Field(
+        default=3,
+        description="The backoff limit for the orchestrator job.",
+    )
     fail_on_container_waiting_reasons: Optional[List[str]] = Field(
         default=[
             "InvalidImageName",
@@ -144,11 +130,21 @@ class KubernetesOrchestratorSettings(BaseSettings):
         "`pod.status.containerStatuses[*].state.waiting.reason` of a job pod, "
         "should cause the job to fail immediately.",
     )
-    job_monitoring_interval: int = Field(
+    job_monitoring_interval: PositiveFloat = Field(
         default=3,
         description="The interval in seconds to monitor the job. Each interval "
-        "is used to check for container issues and streaming logs for the "
-        "job pods.",
+        "is used to check for container issues for the job pods.",
+    )
+    job_monitoring_delay: PositiveFloat = Field(
+        default=0.0,
+        description="The delay in seconds to wait between monitoring active "
+        "step jobs. This can be used to reduce load on the Kubernetes API "
+        "server.",
+    )
+    interrupt_check_interval: PositiveFloat = Field(
+        default=1.0,
+        description="The interval in seconds to check for run interruptions.",
+        ge=0.5,
     )
     pod_failure_policy: Optional[Dict[str, Any]] = Field(
         default=None,
@@ -169,6 +165,53 @@ class KubernetesOrchestratorSettings(BaseSettings):
         description="When stopping a pipeline run, the amount of seconds to wait for a step pod to shutdown gracefully.",
     )
 
+    # Deprecated fields
+    timeout: Optional[int] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    stream_step_logs: Optional[bool] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    pod_startup_timeout: Optional[int] = Field(
+        default=None,
+        description="DEPRECATED/UNUSED.",
+        deprecated=True,
+    )
+    pod_failure_max_retries: Optional[int] = Field(
+        default=None,
+        description="DEPRECATED/UNUSED.",
+        deprecated=True,
+    )
+    pod_failure_retry_delay: Optional[int] = Field(
+        default=None,
+        description="DEPRECATED/UNUSED.",
+        deprecated=True,
+    )
+    pod_failure_backoff: Optional[float] = Field(
+        default=None,
+        description="DEPRECATED/UNUSED.",
+        deprecated=True,
+    )
+    pod_name_prefix: Optional[str] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        "timeout",
+        "stream_step_logs",
+        "pod_startup_timeout",
+        "pod_failure_max_retries",
+        "pod_failure_retry_delay",
+        "pod_failure_backoff",
+        ("pod_name_prefix", "job_name_prefix"),
+    )
+
     @field_validator("pod_failure_policy", mode="before")
     @classmethod
     def _convert_pod_failure_policy(cls, value: Any) -> Any:

zenml/integrations/kubernetes/flavors/kubernetes_step_operator_flavor.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Kubernetes step operator flavor."""
 
-from typing import TYPE_CHECKING, Optional, Type
+from typing import TYPE_CHECKING, List, Optional, Type
 
-from pydantic import Field
+from pydantic import Field, NonNegativeInt
 
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
@@ -23,6 +23,7 @@ from zenml.integrations.kubernetes import KUBERNETES_STEP_OPERATOR_FLAVOR
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.models import ServiceConnectorRequirements
 from zenml.step_operators import BaseStepOperatorConfig, BaseStepOperatorFlavor
+from zenml.utils import deprecation_utils
 
 if TYPE_CHECKING:
     from zenml.integrations.kubernetes.step_operators import (
@@ -31,11 +32,7 @@ if TYPE_CHECKING:
 
 
 class KubernetesStepOperatorSettings(BaseSettings):
-    """Settings for the Kubernetes step operator.
-
-    Configuration options for individual step execution on Kubernetes.
-    Field descriptions are defined inline using Field() descriptors.
-    """
+    """Settings for the Kubernetes step operator."""
 
     pod_settings: Optional[KubernetesPodSettings] = Field(
         default=None,
@@ -49,32 +46,66 @@ class KubernetesStepOperatorSettings(BaseSettings):
         default=False,
         description="Whether to run step containers in privileged mode with extended permissions.",
     )
-    pod_startup_timeout: int = Field(
-        default=600,
-        description="Maximum seconds to wait for step pods to start. Default is 10 minutes.",
+    job_name_prefix: Optional[str] = Field(
+        default=None,
+        description="Prefix for the job name.",
     )
-    pod_failure_max_retries: int = Field(
-        default=3,
-        description="Maximum retry attempts when step pods fail to start.",
+    ttl_seconds_after_finished: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Seconds to keep finished jobs before automatic cleanup.",
     )
-    pod_failure_retry_delay: int = Field(
-        default=10,
-        description="Delay in seconds between pod failure retry attempts.",
+    active_deadline_seconds: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Job deadline in seconds. If the job doesn't finish "
+        "within this time, it will be terminated.",
     )
-    pod_failure_backoff: float = Field(
-        default=1.0,
-        description="Exponential backoff factor for retry delays. Values > 1.0 increase delay with each retry.",
+    fail_on_container_waiting_reasons: Optional[List[str]] = Field(
+        default=[
+            "InvalidImageName",
+            "ErrImagePull",
+            "ImagePullBackOff",
+            "CreateContainerConfigError",
+        ],
+        description="List of container waiting reasons that should cause the "
+        "job to fail immediately. This should be set to a list of "
+        "nonrecoverable reasons, which if found in any "
+        "`pod.status.containerStatuses[*].state.waiting.reason` of a job pod, "
+        "should cause the job to fail immediately.",
+    )
+
+    # Deprecated fields
+    pod_startup_timeout: Optional[int] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    pod_failure_max_retries: Optional[int] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    pod_failure_retry_delay: Optional[int] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    pod_failure_backoff: Optional[float] = Field(
+        default=None,
+        deprecated=True,
+        description="DEPRECATED/UNUSED.",
+    )
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        "pod_startup_timeout",
+        "pod_failure_max_retries",
+        "pod_failure_retry_delay",
+        "pod_failure_backoff",
     )
 
 
 class KubernetesStepOperatorConfig(
     BaseStepOperatorConfig, KubernetesStepOperatorSettings
 ):
-    """Configuration for the Kubernetes step operator.
-
-    Defines cluster connection and execution settings.
-    Field descriptions are defined inline using Field() descriptors.
-    """
+    """Configuration for the Kubernetes step operator."""
 
     kubernetes_namespace: str = Field(
         default="zenml",

zenml/integrations/kubernetes/orchestrators/dag_runner.py

@@ -0,0 +1,367 @@
+#  Copyright (c) ZenML GmbH 2025. 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.
+"""DAG runner."""
+
+import queue
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic import BaseModel
+
+from zenml.logger import get_logger
+from zenml.utils.enum_utils import StrEnum
+
+logger = get_logger(__name__)
+
+
+class NodeStatus(StrEnum):
+    """Status of a DAG node."""
+
+    NOT_READY = "not_ready"  # Can not be started yet
+    READY = "ready"  # Can be started but is still waiting in the queue
+    STARTING = "starting"  # Is being started, but not yet running
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+    CANCELLED = "cancelled"
+
+
+class InterruptMode(StrEnum):
+    """Interrupt mode."""
+
+    GRACEFUL = "graceful"
+    FORCE = "force"
+
+
+class Node(BaseModel):
+    """DAG node."""
+
+    id: str
+    status: NodeStatus = NodeStatus.NOT_READY
+    upstream_nodes: List[str] = []
+    metadata: Dict[str, Any] = {}
+
+    @property
+    def is_finished(self) -> bool:
+        """Whether the node is finished.
+
+        Returns:
+            Whether the node is finished.
+        """
+        return self.status in {
+            NodeStatus.COMPLETED,
+            NodeStatus.FAILED,
+            NodeStatus.SKIPPED,
+            NodeStatus.CANCELLED,
+        }
+
+
+class DagRunner:
+    """DAG runner.
+
+    This class does the orchestration of running the nodes of a DAG. It is
+    running two loops in separate threads:
+    The main thread
+      - checks if any nodes should be skipped or are ready to
+        run, in which case the node will be added to the startup queue
+      - creates a worker thread to start the node and executes it in a thread
+        pool if there are nodes in the startup queue and the maximum
+        parallelism is not reached
+      - periodically checks if the DAG should be interrupted
+    The monitoring thread
+      - monitors the running nodes and updates their status
+    """
+
+    def __init__(
+        self,
+        nodes: List[Node],
+        node_startup_function: Callable[[Node], NodeStatus],
+        node_monitoring_function: Callable[[Node], NodeStatus],
+        node_stop_function: Optional[Callable[[Node], None]] = None,
+        interrupt_function: Optional[
+            Callable[[], Optional[InterruptMode]]
+        ] = None,
+        monitoring_interval: float = 1.0,
+        monitoring_delay: float = 0.0,
+        interrupt_check_interval: float = 1.0,
+        max_parallelism: Optional[int] = None,
+    ) -> None:
+        """Initialize the DAG runner.
+
+        Args:
+            nodes: The nodes of the DAG.
+            node_startup_function: The function to start a node.
+            node_monitoring_function: The function to monitor a node.
+            node_stop_function: The function to stop a node.
+            interrupt_function: Will be periodically called to check if the
+                DAG should be interrupted.
+            monitoring_interval: The interval in which the nodes are monitored.
+            monitoring_delay: The delay in seconds to wait between monitoring
+                different nodes.
+            interrupt_check_interval: The interval in which the interrupt
+                function is called.
+            max_parallelism: The maximum number of nodes to run in parallel.
+        """
+        self.nodes = {node.id: node for node in nodes}
+        self.startup_queue: queue.Queue[Node] = queue.Queue()
+        self.node_startup_function = node_startup_function
+        self.node_monitoring_function = node_monitoring_function
+        self.node_stop_function = node_stop_function
+        self.interrupt_function = interrupt_function
+        self.monitoring_thread = threading.Thread(
+            name="DagRunner-Monitoring-Loop",
+            target=self._monitoring_loop,
+            daemon=True,
+        )
+        self.monitoring_interval = monitoring_interval
+        self.monitoring_delay = monitoring_delay
+        self.interrupt_check_interval = interrupt_check_interval
+        self.max_parallelism = max_parallelism
+        self.shutdown_event = threading.Event()
+        self.startup_executor = ThreadPoolExecutor(
+            max_workers=10, thread_name_prefix="DagRunner-Startup"
+        )
+
+    @property
+    def running_nodes(self) -> List[Node]:
+        """Running nodes.
+
+        Returns:
+            Running nodes.
+        """
+        return [
+            node
+            for node in self.nodes.values()
+            if node.status == NodeStatus.RUNNING
+        ]
+
+    @property
+    def active_nodes(self) -> List[Node]:
+        """Active nodes.
+
+        Active nodes are nodes that are either running or starting.
+
+        Returns:
+            Active nodes.
+        """
+        return [
+            node
+            for node in self.nodes.values()
+            if node.status in {NodeStatus.RUNNING, NodeStatus.STARTING}
+        ]
+
+    def _initialize_startup_queue(self) -> None:
+        """Initialize the startup queue.
+
+        The startup queue contains all nodes that are ready to be started.
+        """
+        for node in self.nodes.values():
+            if node.status in {NodeStatus.READY, NodeStatus.STARTING}:
+                self.startup_queue.put(node)
+
+    def _can_start_node(self, node: Node) -> bool:
+        """Check if a node can be started.
+
+        Args:
+            node: The node to check.
+
+        Returns:
+            Whether the node can be started.
+        """
+        return all(
+            self.nodes[upstream_node_id].status == NodeStatus.COMPLETED
+            for upstream_node_id in node.upstream_nodes
+        )
+
+    def _should_skip_node(self, node: Node) -> bool:
+        """Check if a node should be skipped.
+
+        Args:
+            node: The node to check.
+
+        Returns:
+            Whether the node should be skipped.
+        """
+        return any(
+            self.nodes[upstream_node_id].status
+            in {NodeStatus.FAILED, NodeStatus.SKIPPED, NodeStatus.CANCELLED}
+            for upstream_node_id in node.upstream_nodes
+        )
+
+    def _start_node(self, node: Node) -> None:
+        """Start a node.
+
+        This will start of a thread that will run the startup function.
+
+        Args:
+            node: The node to start.
+        """
+        node.status = NodeStatus.STARTING
+
+        def _start_node_task() -> None:
+            if self.shutdown_event.is_set():
+                logger.debug(
+                    "Cancelling startup of node `%s` because shutdown was "
+                    "requested.",
+                    node.id,
+                )
+                return
+
+            try:
+                node.status = self.node_startup_function(node)
+            except Exception:
+                node.status = NodeStatus.FAILED
+                logger.exception("Node `%s` failed to start.", node.id)
+            else:
+                logger.info(
+                    "Node `%s` started (status: %s)", node.id, node.status
+                )
+
+        self.startup_executor.submit(_start_node_task)
+
+    def _stop_node(self, node: Node) -> None:
+        """Stop a node.
+
+        Args:
+            node: The node to stop.
+
+        Raises:
+            RuntimeError: If the node stop function is not set.
+        """
+        if not self.node_stop_function:
+            raise RuntimeError("Node stop function is not set.")
+
+        self.node_stop_function(node)
+
+    def _stop_all_nodes(self) -> None:
+        """Stop all running nodes."""
+        for node in self.running_nodes:
+            self._stop_node(node)
+            node.status = NodeStatus.CANCELLED
+
+    def _process_nodes(self) -> bool:
+        """Process the nodes.
+
+        This method will check if any nodes should be skipped or are ready to
+        run, in which case the node will be added to the startup queue.
+
+        Returns:
+            Whether the DAG is finished.
+        """
+        finished = True
+
+        for node in self.nodes.values():
+            if node.status == NodeStatus.NOT_READY:
+                if self._should_skip_node(node):
+                    node.status = NodeStatus.SKIPPED
+                    logger.warning(
+                        "Skipping node `%s` because upstream node failed.",
+                        node.id,
+                    )
+                elif self._can_start_node(node):
+                    node.status = NodeStatus.READY
+                    self.startup_queue.put(node)
+
+            if not node.is_finished:
+                finished = False
+
+        # Start nodes until we reach the maximum configured parallelism
+        max_parallelism = self.max_parallelism or len(self.nodes)
+        while len(self.active_nodes) < max_parallelism:
+            try:
+                node = self.startup_queue.get_nowait()
+            except queue.Empty:
+                break
+            else:
+                self.startup_queue.task_done()
+                self._start_node(node)
+
+        return finished
+
+    def _monitoring_loop(self) -> None:
+        """Monitoring loop.
+
+        This should run in a separate thread and monitors the running nodes.
+        """
+        while not self.shutdown_event.is_set():
+            start_time = time.time()
+            for node in self.running_nodes:
+                try:
+                    node.status = self.node_monitoring_function(node)
+                except Exception:
+                    node.status = NodeStatus.FAILED
+                    logger.exception("Node `%s` failed.", node.id)
+                else:
+                    logger.debug(
+                        "Node `%s` status updated to `%s`",
+                        node.id,
+                        node.status,
+                    )
+                    if node.status == NodeStatus.FAILED:
+                        logger.error("Node `%s` failed.", node.id)
+                    elif node.status == NodeStatus.COMPLETED:
+                        logger.info("Node `%s` completed.", node.id)
+
+                time.sleep(self.monitoring_delay)
+
+            duration = time.time() - start_time
+            time_to_sleep = max(0, self.monitoring_interval - duration)
+            self.shutdown_event.wait(timeout=time_to_sleep)
+
+    def run(self) -> Dict[str, NodeStatus]:
+        """Run the DAG.
+
+        Returns:
+            The final node states.
+        """
+        self._initialize_startup_queue()
+
+        self.monitoring_thread.start()
+
+        interrupt_mode = None
+        last_interrupt_check = time.time()
+
+        while True:
+            if self.interrupt_function is not None:
+                if (
+                    time.time() - last_interrupt_check
+                    >= self.interrupt_check_interval
+                ):
+                    if interrupt_mode := self.interrupt_function():
+                        logger.warning("DAG execution interrupted.")
+                        break
+                    last_interrupt_check = time.time()
+
+            is_finished = self._process_nodes()
+            if is_finished:
+                break
+
+            time.sleep(0.5)
+
+        self.shutdown_event.set()
+        if interrupt_mode == InterruptMode.FORCE:
+            # If a force interrupt was requested, we stop all running nodes.
+            self._stop_all_nodes()
+
+        self.monitoring_thread.join()
+
+        node_statuses = {
+            node_id: node.status for node_id, node in self.nodes.items()
+        }
+        logger.debug("Finished with node statuses: %s", node_statuses)
+
+        return node_statuses