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

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