runnable 0.34.0a1__py3-none-any.whl → 1.0.0__py3-none-any.whl

extensions/nodes/task.py

@@ -0,0 +1,92 @@
+import logging
+from datetime import datetime
+from typing import Any, Dict
+
+from pydantic import ConfigDict, Field
+
+from runnable import datastore, defaults
+from runnable.datastore import StepLog
+from runnable.defaults import MapVariableType
+from runnable.nodes import ExecutableNode
+from runnable.tasks import BaseTaskType, create_task
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+
+class TaskNode(ExecutableNode):
+    """
+    A node of type Task.
+
+    This node does the actual function execution of the graph in all cases.
+    """
+
+    executable: BaseTaskType = Field(exclude=True)
+    node_type: str = Field(default="task", serialization_alias="type")
+
+    # It is technically not allowed as parse_from_config filters them.
+    # This is just to get the task level configuration to be present during serialization.
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def parse_from_config(cls, config: Dict[str, Any]) -> "TaskNode":
+        # separate task config from node config
+        task_config = {
+            k: v for k, v in config.items() if k not in TaskNode.model_fields.keys()
+        }
+        node_config = {
+            k: v for k, v in config.items() if k in TaskNode.model_fields.keys()
+        }
+
+        executable = create_task(task_config)
+        return cls(executable=executable, **node_config, **task_config)
+
+    def get_summary(self) -> Dict[str, Any]:
+        summary = {
+            "name": self.name,
+            "type": self.node_type,
+            "executable": self.executable.get_summary(),
+            "catalog": self._get_catalog_settings(),
+        }
+
+        return summary
+
+    def execute(
+        self,
+        mock=False,
+        map_variable: MapVariableType = None,
+        attempt_number: int = 1,
+    ) -> StepLog:
+        """
+        All that we do in runnable is to come to this point where we actually execute the command.
+
+        Args:
+            executor (_type_): The executor class
+            mock (bool, optional): If we should just mock and not execute. Defaults to False.
+            map_variable (dict, optional): If the node is part of internal branch. Defaults to None.
+
+        Returns:
+            StepAttempt: The attempt object
+        """
+        step_log = self._context.run_log_store.get_step_log(
+            self._get_step_log_name(map_variable), self._context.run_id
+        )
+
+        if not mock:
+            # Do not run if we are mocking the execution, could be useful for caching and dry runs
+            attempt_log = self.executable.execute_command(map_variable=map_variable)
+            attempt_log.attempt_number = attempt_number
+        else:
+            attempt_log = datastore.StepAttempt(
+                status=defaults.SUCCESS,
+                start_time=str(datetime.now()),
+                end_time=str(datetime.now()),
+                attempt_number=attempt_number,
+            )
+
+        logger.info(f"attempt_log: {attempt_log}")
+        logger.info(f"Step {self.name} completed with status: {attempt_log.status}")
+
+        step_log.status = attempt_log.status
+        step_log.attempts.append(attempt_log)
+
+        return step_log

extensions/pipeline_executor/__init__.py

@@ -13,7 +13,7 @@ from runnable import (
     utils,
 )
 from runnable.datastore import DataCatalog, JsonParameter, RunLog, StepLog
-from runnable.defaults import TypeMapVariable
+from runnable.defaults import MapVariableType
 from runnable.executor import BasePipelineExecutor
 from runnable.graph import Graph
 from runnable.nodes import BaseNode
@@ -40,7 +40,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
     @property
     def _context(self):
-        assert context.run_context
+        assert isinstance(context.run_context, context.PipelineContext)
         return context.run_context
 
     def _get_parameters(self) -> Dict[str, JsonParameter]:
@@ -104,7 +104,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         )
 
         # Update run_config
-        run_config = utils.get_run_config()
+        run_config = self._context.model_dump()
         logger.debug(f"run_config as seen by executor: {run_config}")
         self._context.run_log_store.set_run_config(
             run_id=self._context.run_id, run_config=run_config
@@ -154,13 +154,15 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         data_catalogs = []
         for name_pattern in node_catalog_settings.get(stage) or []:
             if stage == "get":
-                data_catalog = self._context.catalog_handler.get(
+                data_catalog = self._context.catalog.get(
                     name=name_pattern,
                 )
 
             elif stage == "put":
-                data_catalog = self._context.catalog_handler.put(
-                    name=name_pattern, allow_file_not_found_exc=allow_file_no_found_exc
+                data_catalog = self._context.catalog.put(
+                    name=name_pattern,
+                    allow_file_not_found_exc=allow_file_no_found_exc,
+                    store_copy=node_catalog_settings.get("store_copy", True),
                 )
             else:
                 raise Exception(f"Stage {stage} not supported")
@@ -189,14 +191,15 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             map_variable=map_variable,
         )
         task_console.save_text(log_file_name)
+        task_console.export_text(clear=True)
         # Put the log file in the catalog
-        self._context.catalog_handler.put(name=log_file_name)
+        self._context.catalog.put(name=log_file_name)
         os.remove(log_file_name)
 
     def _execute_node(
         self,
         node: BaseNode,
-        map_variable: TypeMapVariable = None,
+        map_variable: MapVariableType = None,
         mock: bool = False,
     ):
         """
@@ -250,6 +253,10 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         console.print(f"Summary of the step: {step_log.internal_name}")
         console.print(step_log.get_summary(), style=defaults.info_style)
 
+        self.add_task_log_to_catalog(
+            name=self._context_node.internal_name, map_variable=map_variable
+        )
+
         self._context_node = None
 
         self._context.run_log_store.add_step_log(step_log, self._context.run_id)
@@ -266,7 +273,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         """
         step_log.code_identities.append(utils.get_git_code_identity())
 
-    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None):
+    def execute_from_graph(self, node: BaseNode, map_variable: MapVariableType = None):
         """
         This is the entry point to from the graph execution.
 
@@ -315,8 +322,6 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             node.execute_as_graph(map_variable=map_variable)
             return
 
-        task_console.export_text(clear=True)
-
         task_name = node._resolve_map_placeholders(node.internal_name, map_variable)
         console.print(
             f":runner: Executing the node {task_name} ... ", style="bold color(208)"
@@ -324,7 +329,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         self.trigger_node_execution(node=node, map_variable=map_variable)
 
     def trigger_node_execution(
-        self, node: BaseNode, map_variable: TypeMapVariable = None
+        self, node: BaseNode, map_variable: MapVariableType = None
     ):
         """
         Call this method only if we are responsible for traversing the graph via
@@ -342,7 +347,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         pass
 
     def _get_status_and_next_node_name(
-        self, current_node: BaseNode, dag: Graph, map_variable: TypeMapVariable = None
+        self, current_node: BaseNode, dag: Graph, map_variable: MapVariableType = None
     ) -> tuple[str, str]:
         """
         Given the current node and the graph, returns the name of the next node to execute.
@@ -380,7 +385,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
         return step_log.status, next_node_name
 
-    def execute_graph(self, dag: Graph, map_variable: TypeMapVariable = None):
+    def execute_graph(self, dag: Graph, map_variable: MapVariableType = None):
         """
         The parallelization is controlled by the nodes and not by this function.
 
@@ -409,7 +414,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
                 dag.internal_branch_name or "Graph",
                 map_variable,
             )
-            branch_execution_task = self._context.progress.add_task(
+            branch_execution_task = context.progress.add_task(
                 f"[dark_orange]Executing {branch_task_name}",
                 total=1,
             )
@@ -429,7 +434,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
             depth = " " * ((task_name.count(".")) or 1 - 1)
 
-            task_execution = self._context.progress.add_task(
+            task_execution = context.progress.add_task(
                 f"{depth}Executing {task_name}", total=1
             )
 
@@ -440,20 +445,20 @@ class GenericPipelineExecutor(BasePipelineExecutor):
                 )
 
                 if status == defaults.SUCCESS:
-                    self._context.progress.update(
+                    context.progress.update(
                         task_execution,
                         description=f"{depth}[green] {task_name} Completed",
                         completed=True,
                         overflow="fold",
                     )
                 else:
-                    self._context.progress.update(
+                    context.progress.update(
                         task_execution,
                         description=f"{depth}[red] {task_name} Failed",
                         completed=True,
                     )  # type ignore
             except Exception as e:  # noqa: E722
-                self._context.progress.update(
+                context.progress.update(
                     task_execution,
                     description=f"{depth}[red] {task_name} Errored",
                     completed=True,
@@ -461,11 +466,6 @@ class GenericPipelineExecutor(BasePipelineExecutor):
                 console.print(e, style=defaults.error_style)
                 logger.exception(e)
                 raise
-            finally:
-                # Add task log to the catalog
-                self.add_task_log_to_catalog(
-                    name=working_on.internal_name, map_variable=map_variable
-                )
 
             console.rule(style="[dark orange]")
 
@@ -475,7 +475,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
             current_node = next_node_name
 
         if branch_execution_task:
-            self._context.progress.update(
+            context.progress.update(
                 branch_execution_task,
                 description=f"[green3] {branch_task_name} completed",
                 completed=True,
@@ -567,7 +567,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
         return effective_node_config
 
-    def fan_out(self, node: BaseNode, map_variable: TypeMapVariable = None):
+    def fan_out(self, node: BaseNode, map_variable: MapVariableType = None):
         """
         This method is used to appropriately fan-out the execution of a composite node.
         This is only useful when we want to execute a composite node during 3rd party orchestrators.
@@ -599,7 +599,7 @@ class GenericPipelineExecutor(BasePipelineExecutor):
 
         node.fan_out(map_variable=map_variable)
 
-    def fan_in(self, node: BaseNode, map_variable: TypeMapVariable = None):
+    def fan_in(self, node: BaseNode, map_variable: MapVariableType = None):
         """
         This method is used to appropriately fan-in after the execution of a composite node.
         This is only useful when we want to execute a composite node during 3rd party orchestrators.

extensions/pipeline_executor/argo.py

@@ -20,13 +20,13 @@ from pydantic import (
 from pydantic.alias_generators import to_camel
 from ruamel.yaml import YAML
 
-from extensions.nodes.nodes import MapNode, ParallelNode, TaskNode
-
-# TODO: Should be part of a wider refactor
-# from extensions.nodes.torch import TorchNode
+from extensions.nodes.conditional import ConditionalNode
+from extensions.nodes.map import MapNode
+from extensions.nodes.parallel import ParallelNode
+from extensions.nodes.task import TaskNode
 from extensions.pipeline_executor import GenericPipelineExecutor
-from runnable import defaults, utils
-from runnable.defaults import TypeMapVariable
+from runnable import defaults
+from runnable.defaults import MapVariableType
 from runnable.graph import Graph, search_node_by_internal_name
 from runnable.nodes import BaseNode
 
@@ -307,6 +307,7 @@ class DagTask(BaseModelWIthConfig):
     template: str  # Should be name of a container template or dag template
     arguments: Optional[Arguments] = Field(default=None)
     with_param: Optional[str] = Field(default=None)
+    when_param: Optional[str] = Field(default=None, serialization_alias="when")
     depends: Optional[str] = Field(default=None)
 
 
@@ -451,7 +452,7 @@ class ArgoExecutor(GenericPipelineExecutor):
     """
 
     service_name: str = "argo"
-    _is_local: bool = False
+    _should_setup_run_log_at_traversal: bool = PrivateAttr(default=False)
     mock: bool = False
 
     model_config = ConfigDict(
@@ -533,13 +534,13 @@ class ArgoExecutor(GenericPipelineExecutor):
         parameters: Optional[list[Parameter]],
         task_name: str,
     ):
-        map_variable: TypeMapVariable = {}
+        map_variable: MapVariableType = {}
         for parameter in parameters or []:
             map_variable[parameter.name] = (  # type: ignore
                 "{{inputs.parameters." + str(parameter.name) + "}}"
             )
 
-        fan_command = utils.get_fan_command(
+        fan_command = self._context.get_fan_command(
             mode=mode,
             node=node,
             run_id=self._run_id_as_parameter,
@@ -563,6 +564,8 @@ class ArgoExecutor(GenericPipelineExecutor):
         outputs: Optional[Outputs] = None
         if mode == "out" and node.node_type == "map":
             outputs = Outputs(parameters=[OutputParameter(name="iterate-on")])
+        if mode == "out" and node.node_type == "conditional":
+            outputs = Outputs(parameters=[OutputParameter(name="case")])
 
         container_template = ContainerTemplate(
             name=task_name,
@@ -586,7 +589,7 @@ class ArgoExecutor(GenericPipelineExecutor):
         task_name: str,
         inputs: Optional[Inputs] = None,
     ) -> ContainerTemplate:
-        assert node.node_type in ["task", "torch", "success", "stub", "fail"]
+        assert node.node_type in ["task", "success", "stub", "fail"]
 
         node_override = None
         if hasattr(node, "overrides"):
@@ -602,17 +605,17 @@ class ArgoExecutor(GenericPipelineExecutor):
 
         inputs = inputs or Inputs(parameters=[])
 
-        map_variable: TypeMapVariable = {}
+        map_variable: MapVariableType = {}
         for parameter in inputs.parameters or []:
             map_variable[parameter.name] = (  # type: ignore
                 "{{inputs.parameters." + str(parameter.name) + "}}"
             )
 
         # command = "runnable execute-single-node"
-        command = utils.get_node_execution_command(
+        command = self._context.get_node_callable_command(
             node=node,
-            over_write_run_id=self._run_id_as_parameter,
             map_variable=map_variable,
+            over_write_run_id=self._run_id_as_parameter,
             log_level=self._log_level_as_parameter,
         )
 
@@ -649,7 +652,7 @@ class ArgoExecutor(GenericPipelineExecutor):
     def _set_env_vars_to_task(
         self, working_on: BaseNode, container_template: CoreContainerTemplate
     ):
-        if working_on.node_type not in ["task", "torch"]:
+        if working_on.node_type not in ["task"]:
             return
 
         global_envs: dict[str, str] = {}
@@ -711,6 +714,7 @@ class ArgoExecutor(GenericPipelineExecutor):
             assert parent_dag_template.dag
 
             parent_dag_template.dag.tasks.append(on_failure_task)
+
             self._gather_tasks_for_dag_template(
                 on_failure_dag,
                 dag=dag,
@@ -722,6 +726,7 @@ class ArgoExecutor(GenericPipelineExecutor):
     # - We are using withParam and arguments of the map template to send that value in
     # - The map template should receive that value as a parameter into the template.
     # - The task then start to use it as inputs.parameters.iterate-on
+    # the when param should be an evaluation
 
     def _gather_tasks_for_dag_template(
         self,
@@ -757,7 +762,7 @@ class ArgoExecutor(GenericPipelineExecutor):
             depends = task_name
 
             match working_on.node_type:
-                case "task" | "success" | "stub":
+                case "task" | "success" | "stub" | "fail":
                     template_of_container = self._create_container_template(
                         working_on,
                         task_name=task_name,
@@ -767,9 +772,11 @@ class ArgoExecutor(GenericPipelineExecutor):
 
                     self._templates.append(template_of_container)
 
-                case "map" | "parallel":
-                    assert isinstance(working_on, MapNode) or isinstance(
-                        working_on, ParallelNode
+                case "map" | "parallel" | "conditional":
+                    assert (
+                        isinstance(working_on, MapNode)
+                        or isinstance(working_on, ParallelNode)
+                        or isinstance(working_on, ConditionalNode)
                     )
                     node_type = working_on.node_type
 
@@ -792,7 +799,8 @@ class ArgoExecutor(GenericPipelineExecutor):
                     )
 
                     # Add the composite task
-                    with_param = None
+                    with_param: Optional[str] = None
+                    when_param: Optional[str] = None
                     added_parameters = parameters or []
                     branches = {}
                     if node_type == "map":
@@ -807,22 +815,34 @@ class ArgoExecutor(GenericPipelineExecutor):
                     elif node_type == "parallel":
                         assert isinstance(working_on, ParallelNode)
                         branches = working_on.branches
+                    elif node_type == "conditional":
+                        assert isinstance(working_on, ConditionalNode)
+                        branches = working_on.branches
+                        when_param = (
+                            f"{{{{tasks.{task_name}-fan-out.outputs.parameters.case}}}}"
+                        )
                     else:
                         raise ValueError("Invalid node type")
 
                     fan_in_depends = ""
 
                     for name, branch in branches.items():
+                        match_when = branch.internal_branch_name.split(".")[-1]
                         name = (
                             name.replace(" ", "-").replace(".", "-").replace("_", "-")
                         )
 
+                        if node_type == "conditional":
+                            assert isinstance(working_on, ConditionalNode)
+                            when_param = f"'{match_when}' == {{{{tasks.{task_name}-fan-out.outputs.parameters.case}}}}"
+
                         branch_task = DagTask(
                             name=f"{task_name}-{name}",
                             template=f"{task_name}-{name}",
                             depends=f"{task_name}-fan-out.Succeeded",
                             arguments=Arguments(parameters=added_parameters),
                             with_param=with_param,
+                            when_param=when_param,
                         )
                         composite_template.dag.tasks.append(branch_task)
 
@@ -836,6 +856,8 @@ class ArgoExecutor(GenericPipelineExecutor):
                             ),
                         )
 
+                        assert isinstance(branch, Graph)
+
                         self._gather_tasks_for_dag_template(
                             dag_template=branch_template,
                             dag=branch,
@@ -862,28 +884,6 @@ class ArgoExecutor(GenericPipelineExecutor):
 
                     self._templates.append(composite_template)
 
-                case "torch":
-                    from extensions.nodes.torch import TorchNode
-
-                    assert isinstance(working_on, TorchNode)
-                    # TODO: Need to add multi-node functionality
-                    # Check notes on the torch node
-
-                    template_of_container = self._create_container_template(
-                        working_on,
-                        task_name=task_name,
-                        inputs=Inputs(parameters=parameters),
-                    )
-                    assert template_of_container.container is not None
-
-                    if working_on.node_type == "task":
-                        self._expose_secrets_to_task(
-                            working_on,
-                            container_template=template_of_container.container,
-                        )
-
-                    self._templates.append(template_of_container)
-
             self._handle_failures(
                 working_on,
                 dag,
@@ -958,7 +958,7 @@ class ArgoExecutor(GenericPipelineExecutor):
                 f,
             )
 
-    def _implicitly_fail(self, node: BaseNode, map_variable: TypeMapVariable):
+    def _implicitly_fail(self, node: BaseNode, map_variable: MapVariableType):
         assert self._context.dag
         _, current_branch = search_node_by_internal_name(
             dag=self._context.dag, internal_name=node.internal_name
@@ -1005,7 +1005,7 @@ class ArgoExecutor(GenericPipelineExecutor):
 
         self._implicitly_fail(node, map_variable)
 
-    def fan_out(self, node: BaseNode, map_variable: TypeMapVariable = None):
+    def fan_out(self, node: BaseNode, map_variable: MapVariableType = None):
         # This could be the first step of the graph
         self._use_volumes()
 
@@ -1025,7 +1025,13 @@ class ArgoExecutor(GenericPipelineExecutor):
             with open("/tmp/output.txt", mode="w", encoding="utf-8") as myfile:
                 json.dump(iterate_on.get_value(), myfile, indent=4)
 
-    def fan_in(self, node: BaseNode, map_variable: TypeMapVariable = None):
+        if node.node_type == "conditional":
+            assert isinstance(node, ConditionalNode)
+
+            with open("/tmp/output.txt", mode="w", encoding="utf-8") as myfile:
+                json.dump(node.get_parameter_value(), myfile, indent=4)
+
+    def fan_in(self, node: BaseNode, map_variable: MapVariableType = None):
         self._use_volumes()
         super().fan_in(node, map_variable)
 
@@ -1036,9 +1042,9 @@ class ArgoExecutor(GenericPipelineExecutor):
             case "chunked-fs":
                 self._context.run_log_store.log_folder = self._container_log_location
 
-        match self._context.catalog_handler.service_name:
+        match self._context.catalog.service_name:
             case "file-system":
-                self._context.catalog_handler.catalog_location = (
+                self._context.catalog.catalog_location = (
                     self._container_catalog_location
                 )
 

extensions/pipeline_executor/emulate.py

@@ -0,0 +1,112 @@
+import logging
+import shlex
+import subprocess
+import sys
+
+from pydantic import PrivateAttr
+
+from extensions.pipeline_executor import GenericPipelineExecutor
+from runnable import defaults
+from runnable.defaults import MapVariableType
+from runnable.nodes import BaseNode
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+
+class Emulator(GenericPipelineExecutor):
+    """
+    In the mode of local execution, we run everything on the local computer.
+
+    This has some serious implications on the amount of time it would take to complete the run.
+    Also ensure that the local compute is good enough for the compute to happen of all the steps.
+
+    Example config:
+
+    ```yaml
+    pipeline-executor:
+      type: local
+    ```
+
+    """
+
+    service_name: str = "emulator"
+
+    _should_setup_run_log_at_traversal: bool = PrivateAttr(default=True)
+
+    def trigger_node_execution(
+        self, node: BaseNode, map_variable: MapVariableType = None
+    ):
+        """
+        In this mode of execution, we prepare for the node execution and execute the node
+
+        Args:
+            node (BaseNode): [description]
+            map_variable (str, optional): [description]. Defaults to ''.
+        """
+        command = self._context.get_node_callable_command(
+            node, map_variable=map_variable
+        )
+
+        self.run_click_command(command)
+        # execute the command in a forked process
+
+        step_log = self._context.run_log_store.get_step_log(
+            node._get_step_log_name(map_variable), self._context.run_id
+        )
+        if step_log.status != defaults.SUCCESS:
+            msg = "Node execution inside the emulate failed. Please check the logs.\n"
+            logger.error(msg)
+            step_log.status = defaults.FAIL
+            self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+
+    def execute_node(self, node: BaseNode, map_variable: MapVariableType = None):
+        """
+        For local execution, we just execute the node.
+
+        Args:
+            node (BaseNode): _description_
+            map_variable (dict[str, str], optional): _description_. Defaults to None.
+        """
+        self._execute_node(node=node, map_variable=map_variable)
+
+    def run_click_command(self, command: str) -> str:
+        """
+        Execute a Click-based CLI command in the current virtual environment.
+
+        Args:
+            args: List of Click command arguments (including subcommands and options)
+
+        Returns:
+            Combined stdout/stderr output as string
+        """
+        # For Click commands installed via setup.py entry_points
+        # command = [sys.executable, '-m', 'your_package.cli'] + args
+
+        # For direct module execution
+        sub_command = [sys.executable, "-m", "runnable.cli"] + shlex.split(command)[1:]
+
+        process = subprocess.Popen(
+            sub_command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+            bufsize=1,
+        )
+
+        output = []
+        try:
+            while True:
+                line = process.stdout.readline()  # type: ignore
+                if not line and process.poll() is not None:
+                    break
+                print(line, end="")
+                output.append(line)
+        finally:
+            process.stdout.close()  # type: ignore
+
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(
+                process.returncode, command, "".join(output)
+            )
+
+        return "".join(output)