zenml-nightly 0.80.1.dev20250408__py3-none-any.whl → 0.80.2.dev20250410__py3-none-any.whl

zenml/orchestrators/dag_runner.py

@@ -17,7 +17,7 @@ import threading
 import time
 from collections import defaultdict
 from enum import Enum
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 
 from zenml.logger import get_logger
 
@@ -54,6 +54,7 @@ class NodeStatus(Enum):
     WAITING = "Waiting"
     RUNNING = "Running"
     COMPLETED = "Completed"
+    FAILED = "Failed"
 
 
 class ThreadedDagRunner:
@@ -70,6 +71,7 @@ class ThreadedDagRunner:
         self,
         dag: Dict[str, List[str]],
         run_fn: Callable[[str], Any],
+        finalize_fn: Optional[Callable[[Dict[str, NodeStatus]], None]] = None,
         parallel_node_startup_waiting_period: float = 0.0,
     ) -> None:
         """Define attributes and initialize all nodes in waiting state.
@@ -79,6 +81,8 @@ class ThreadedDagRunner:
                 E.g.: [(1->2), (1->3), (2->4), (3->4)] should be represented as
                 `dag={2: [1], 3: [1], 4: [2, 3]}`
             run_fn: A function `run_fn(node)` that runs a single node
+            finalize_fn: A function `finalize_fn(node_states)` that is called
+                when all nodes have completed.
             parallel_node_startup_waiting_period: Delay in seconds to wait in
                 between starting parallel nodes.
         """
@@ -88,6 +92,7 @@ class ThreadedDagRunner:
         self.dag = dag
         self.reversed_dag = reverse_dag(dag)
         self.run_fn = run_fn
+        self.finalize_fn = finalize_fn
         self.nodes = dag.keys()
         self.node_states = {node: NodeStatus.WAITING for node in self.nodes}
         self._lock = threading.Lock()
@@ -123,8 +128,12 @@ class ThreadedDagRunner:
         Args:
             node: The node.
         """
-        self.run_fn(node)
-        self._finish_node(node)
+        try:
+            self.run_fn(node)
+            self._finish_node(node)
+        except Exception as e:
+            self._finish_node(node, failed=True)
+            logger.exception(f"Node `{node}` failed: {e}")
 
     def _run_node_in_thread(self, node: str) -> threading.Thread:
         """Run a single node in a separate thread.
@@ -148,7 +157,7 @@ class ThreadedDagRunner:
         thread.start()
         return thread
 
-    def _finish_node(self, node: str) -> None:
+    def _finish_node(self, node: str, failed: bool = False) -> None:
         """Finish a node run.
 
         First updates the node status to completed.
@@ -156,20 +165,28 @@ class ThreadedDagRunner:
 
         Args:
             node: The node.
+            failed: Whether the node failed.
         """
         # Update node status to completed.
         assert self.node_states[node] == NodeStatus.RUNNING
         with self._lock:
-            self.node_states[node] = NodeStatus.COMPLETED
+            if failed:
+                self.node_states[node] = NodeStatus.FAILED
+            else:
+                self.node_states[node] = NodeStatus.COMPLETED
+
+        if failed:
+            # If the node failed, we don't need to run any downstream nodes.
+            return
 
         # Run downstream nodes.
         threads: List[threading.Thread] = []
-        for downstram_node in self.reversed_dag[node]:
-            if self._can_run(downstram_node):
+        for downstream_node in self.reversed_dag[node]:
+            if self._can_run(downstream_node):
                 if threads and self.parallel_node_startup_waiting_period > 0:
                     time.sleep(self.parallel_node_startup_waiting_period)
 
-                thread = self._run_node_in_thread(downstram_node)
+                thread = self._run_node_in_thread(downstream_node)
                 threads.append(thread)
 
         # Wait for all downstream nodes to complete.
@@ -198,11 +215,27 @@ class ThreadedDagRunner:
         for thread in threads:
             thread.join()
 
-        # Make sure all nodes were run, otherwise print a warning.
+        # Call the finalize function.
+        if self.finalize_fn:
+            self.finalize_fn(self.node_states)
+
+        # Print a status report.
+        failed_nodes = []
+        skipped_nodes = []
         for node in self.nodes:
-            if self.node_states[node] == NodeStatus.WAITING:
-                upstream_nodes = self.dag[node]
-                logger.warning(
-                    f"Node `{node}` was never run, because it was still"
-                    f" waiting for the following nodes: `{upstream_nodes}`."
-                )
+            if self.node_states[node] == NodeStatus.FAILED:
+                failed_nodes.append(node)
+            elif self.node_states[node] == NodeStatus.WAITING:
+                skipped_nodes.append(node)
+
+        if failed_nodes:
+            logger.error(
+                "The following nodes failed: " + ", ".join(failed_nodes)
+            )
+        if skipped_nodes:
+            logger.warning(
+                "The following nodes were not run because they depend on other "
+                "nodes that didn't complete: " + ", ".join(skipped_nodes)
+            )
+        if not failed_nodes and not skipped_nodes:
+            logger.info("All nodes completed successfully.")

zenml/orchestrators/local/local_orchestrator.py

@@ -27,7 +27,7 @@ from zenml.stack import Stack
 from zenml.utils import string_utils
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
 
 logger = get_logger(__name__)
 
@@ -46,6 +46,7 @@ class LocalOrchestrator(BaseOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Iterates through all steps and executes them sequentially.
 
@@ -54,6 +55,7 @@ class LocalOrchestrator(BaseOrchestrator):
             stack: The stack on which the pipeline is deployed.
             environment: Environment variables to set in the orchestration
                 environment.
+            placeholder_run: An optional placeholder run for the deployment.
         """
         if deployment.schedule:
             logger.warning(

zenml/orchestrators/local_docker/local_docker_orchestrator.py

@@ -39,7 +39,7 @@ from zenml.stack import Stack, StackValidator
 from zenml.utils import docker_utils, string_utils
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
 
 logger = get_logger(__name__)
 
@@ -96,6 +96,7 @@ class LocalDockerOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Sequentially runs all pipeline steps in local Docker containers.
 
@@ -104,6 +105,7 @@ class LocalDockerOrchestrator(ContainerizedOrchestrator):
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
                 environment.
+            placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             RuntimeError: If a step fails.

zenml/orchestrators/step_launcher.py

@@ -154,8 +154,8 @@ class StepLauncher:
         if step_logging_enabled:
             # Configure the logs
             logs_uri = step_logging.prepare_logs_uri(
-                self._stack.artifact_store,
-                self._step.config.name,
+                artifact_store=self._stack.artifact_store,
+                step_name=self._step_name,
             )
 
             logs_context = step_logging.StepLogsStorageContext(

zenml/utils/secret_utils.py

@@ -102,7 +102,7 @@ def SecretField(*args: Any, **kwargs: Any) -> Any:
     """
     json_schema_extra = kwargs.get("json_schema_extra", {})
     json_schema_extra.update({PYDANTIC_SENSITIVE_FIELD_MARKER: True})
-    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)
 
 
 def ClearTextField(*args: Any, **kwargs: Any) -> Any:
@@ -119,7 +119,7 @@ def ClearTextField(*args: Any, **kwargs: Any) -> Any:
     """
     json_schema_extra = kwargs.get("json_schema_extra", {})
     json_schema_extra.update({PYDANTIC_CLEAR_TEXT_FIELD_MARKER: True})
-    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)
 
 
 def is_secret_field(field: "FieldInfo") -> bool:

zenml/zen_server/dashboard/assets/{404-2I8egBQu.js → 404-CYPi9d8E.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-AvWw-1nd.js";import{h as s,r as t}from"./index-DPjvk73v.js";import{E as r}from"./EmptyState-DpbfQBDE.js";import{S as a}from"./help-CfT0tY2I.js";import{L as o}from"./@react-router-BUo5vhN4.js";import"./@tanstack-CcI3lvwB.js";import"./@reactflow-BHoFKFSZ.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-C7hRs6Kx.js";import{h as s,r as t}from"./index-Be2jLIYM.js";import{E as r}from"./EmptyState-M1jafpg6.js";import{S as a}from"./help-DyMolRxD.js";import{L as o}from"./@react-router-BGgzhn5D.js";import"./@tanstack-CSxjHCME.js";import"./@reactflow-CQi1Z1Wq.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};