runnable 0.34.0a3__py3-none-any.whl → 0.36.0__py3-none-any.whl

runnable/sdk.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
+import inspect
 import logging
-import os
 import re
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -16,26 +16,17 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TextColumn,
-    TimeElapsedColumn,
-)
-from rich.table import Column
 from typing_extensions import Self
 
-from extensions.nodes.nodes import (
-    FailNode,
-    MapNode,
-    ParallelNode,
-    StubNode,
-    SuccessNode,
-    TaskNode,
-)
-from runnable import console, defaults, entrypoints, exceptions, graph, utils
-from runnable.executor import BaseJobExecutor, BasePipelineExecutor
+from extensions.nodes.conditional import ConditionalNode
+from extensions.nodes.fail import FailNode
+from extensions.nodes.map import MapNode
+from extensions.nodes.parallel import ParallelNode
+from extensions.nodes.stub import StubNode
+from extensions.nodes.success import SuccessNode
+from extensions.nodes.task import TaskNode
+from runnable import defaults, graph
+from runnable.executor import BaseJobExecutor
 from runnable.nodes import TraversalNode
 from runnable.tasks import BaseTaskType as RunnableTask
 from runnable.tasks import TaskReturns
@@ -50,6 +41,7 @@ StepType = Union[
     "Parallel",
     "Map",
     "TorchTask",
+    "Conditional",
 ]
 
 
@@ -193,6 +185,9 @@ class BaseTask(BaseTraversal):
             "This method should be implemented in the child class"
         )
 
+    def as_pipeline(self) -> "Pipeline":
+        return Pipeline(steps=[self], name=self.internal_name)  # type: ignore
+
 
 class PythonTask(BaseTask):
     """
@@ -482,6 +477,9 @@ class Stub(BaseTraversal):
 
         return StubNode.parse_from_config(self.model_dump(exclude_none=True))
 
+    def as_pipeline(self) -> "Pipeline":
+        return Pipeline(steps=[self])
+
 
 class Parallel(BaseTraversal):
     """
@@ -521,6 +519,53 @@ class Parallel(BaseTraversal):
         return node
 
 
+class Conditional(BaseTraversal):
+    branches: Dict[str, "Pipeline"]
+    parameter: str  # the name of the parameter should be isalnum
+
+    @field_validator("parameter")
+    @classmethod
+    def validate_parameter(cls, parameter: str) -> str:
+        if not parameter.isalnum():
+            raise AssertionError(
+                "The parameter name should be alphanumeric and not empty"
+            )
+        return parameter
+
+    @field_validator("branches")
+    @classmethod
+    def validate_branches(
+        cls, branches: Dict[str, "Pipeline"]
+    ) -> Dict[str, "Pipeline"]:
+        for branch_name in branches.keys():
+            if not branch_name.isalnum():
+                raise ValueError(f"Branch '{branch_name}' must be alphanumeric.")
+        return branches
+
+    @computed_field  # type: ignore
+    @property
+    def graph_branches(self) -> Dict[str, graph.Graph]:
+        return {
+            name: pipeline._dag.model_copy() for name, pipeline in self.branches.items()
+        }
+
+    def create_node(self) -> ConditionalNode:
+        if not self.next_node:
+            if not (self.terminate_with_failure or self.terminate_with_success):
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
+
+        node = ConditionalNode(
+            name=self.name,
+            branches=self.graph_branches,
+            internal_name="",
+            next_node=self.next_node,
+            parameter=self.parameter,
+        )
+        return node
+
+
 class Map(BaseTraversal):
     """
     A node that iterates over a list of items and executes a pipeline for each item.
@@ -544,7 +589,6 @@ class Map(BaseTraversal):
     iterate_on: str
     iterate_as: str
     reducer: Optional[str] = Field(default=None, alias="reducer")
-    overrides: Dict[str, Any] = Field(default_factory=dict)
 
     @computed_field  # type: ignore
     @property
@@ -565,7 +609,6 @@ class Map(BaseTraversal):
             next_node=self.next_node,
             iterate_on=self.iterate_on,
             iterate_as=self.iterate_as,
-            overrides=self.overrides,
             reducer=self.reducer,
         )
 
@@ -736,6 +779,15 @@ class Pipeline(BaseModel):
             return False
         return True
 
+    def get_caller(self) -> str:
+        caller_stack = inspect.stack()[2]
+        relative_to_root = str(Path(caller_stack.filename).relative_to(Path.cwd()))
+
+        module_name = re.sub(r"\b.py\b", "", relative_to_root.replace("/", "."))
+        module_to_call = f"{module_name}.{caller_stack.function}"
+
+        return module_to_call
+
     def execute(
         self,
         configuration_file: str = "",
@@ -753,106 +805,31 @@ class Pipeline(BaseModel):
             # Immediately return as this call is only for getting the pipeline definition
             return {}
 
-        logger.setLevel(log_level)
-
-        run_id = utils.generate_run_id(run_id=run_id)
-
-        parameters_file = os.environ.get("RUNNABLE_PARAMETERS_FILE", parameters_file)
+        from runnable import context
 
-        tag = os.environ.get("RUNNABLE_tag", tag)
+        logger.setLevel(log_level)
 
-        configuration_file = os.environ.get(
-            "RUNNABLE_CONFIGURATION_FILE", configuration_file
-        )
-        run_context = entrypoints.prepare_configurations(
+        service_configurations = context.ServiceConfigurations(
             configuration_file=configuration_file,
-            run_id=run_id,
-            tag=tag,
-            parameters_file=parameters_file,
+            execution_context=context.ExecutionContext.PIPELINE,
         )
 
-        assert isinstance(run_context.executor, BasePipelineExecutor)
-
-        utils.set_runnable_environment_variables(
-            run_id=run_id, configuration_file=configuration_file, tag=tag
-        )
-
-        dag_definition = self._dag.model_dump(by_alias=True, exclude_none=True)
-        run_context.from_sdk = True
-        run_context.dag = graph.create_graph(dag_definition)
-
-        console.print("Working with context:")
-        console.print(run_context)
-        console.rule(style="[dark orange]")
-
-        if not run_context.executor._is_local:
-            # We are not working with executor that does not work in local environment
-            import inspect
-
-            caller_stack = inspect.stack()[1]
-            relative_to_root = str(Path(caller_stack.filename).relative_to(Path.cwd()))
-
-            module_name = re.sub(r"\b.py\b", "", relative_to_root.replace("/", "."))
-            module_to_call = f"{module_name}.{caller_stack.function}"
-
-            run_context.pipeline_file = f"{module_to_call}.py"
-            run_context.from_sdk = True
-
-        # Prepare for graph execution
-        run_context.executor._set_up_run_log(exists_ok=False)
-
-        with Progress(
-            SpinnerColumn(spinner_name="runner"),
-            TextColumn(
-                "[progress.description]{task.description}", table_column=Column(ratio=2)
-            ),
-            BarColumn(table_column=Column(ratio=1), style="dark_orange"),
-            TimeElapsedColumn(table_column=Column(ratio=1)),
-            console=console,
-            expand=True,
-        ) as progress:
-            pipeline_execution_task = progress.add_task(
-                "[dark_orange] Starting execution .. ", total=1
-            )
-            try:
-                run_context.progress = progress
-
-                run_context.executor.execute_graph(dag=run_context.dag)
+        configurations = {
+            "pipeline_definition_file": self.get_caller(),
+            "parameters_file": parameters_file,
+            "tag": tag,
+            "run_id": run_id,
+            "execution_mode": context.ExecutionMode.PYTHON,
+            "configuration_file": configuration_file,
+            **service_configurations.services,
+        }
 
-                if not run_context.executor._is_local:
-                    # non local executors just traverse the graph and do nothing
-                    return {}
+        run_context = context.PipelineContext.model_validate(configurations)
+        context.run_context = run_context
 
-                run_log = run_context.run_log_store.get_run_log_by_id(
-                    run_id=run_context.run_id, full=False
-                )
+        assert isinstance(run_context, context.PipelineContext)
 
-                if run_log.status == defaults.SUCCESS:
-                    progress.update(
-                        pipeline_execution_task,
-                        description="[green] Success",
-                        completed=True,
-                    )
-                else:
-                    progress.update(
-                        pipeline_execution_task,
-                        description="[red] Failed",
-                        completed=True,
-                    )
-                    raise exceptions.ExecutionFailedError(run_context.run_id)
-            except Exception as e:  # noqa: E722
-                console.print(e, style=defaults.error_style)
-                progress.update(
-                    pipeline_execution_task,
-                    description="[red] Errored execution",
-                    completed=True,
-                )
-                raise
-
-        if run_context.executor._is_local:
-            return run_context.run_log_store.get_run_log_by_id(
-                run_id=run_context.run_id
-            )
+        run_context.execute()
 
 
 class BaseJob(BaseModel):
@@ -876,6 +853,15 @@ class BaseJob(BaseModel):
     def get_task(self) -> RunnableTask:
         raise NotImplementedError
 
+    def get_caller(self) -> str:
+        caller_stack = inspect.stack()[2]
+        relative_to_root = str(Path(caller_stack.filename).relative_to(Path.cwd()))
+
+        module_name = re.sub(r"\b.py\b", "", relative_to_root.replace("/", "."))
+        module_to_call = f"{module_name}.{caller_stack.function}"
+
+        return module_to_call
+
     def return_catalog_settings(self) -> Optional[List[str]]:
         if self.catalog is None:
             return []
@@ -902,65 +888,32 @@ class BaseJob(BaseModel):
         if self._is_called_for_definition():
             # Immediately return as this call is only for getting the job definition
             return {}
-        logger.setLevel(log_level)
-
-        run_id = utils.generate_run_id(run_id=job_id)
-
-        parameters_file = os.environ.get("RUNNABLE_PARAMETERS_FILE", parameters_file)
-
-        tag = os.environ.get("RUNNABLE_tag", tag)
+        from runnable import context
 
-        configuration_file = os.environ.get(
-            "RUNNABLE_CONFIGURATION_FILE", configuration_file
-        )
+        logger.setLevel(log_level)
 
-        run_context = entrypoints.prepare_configurations(
+        service_configurations = context.ServiceConfigurations(
             configuration_file=configuration_file,
-            run_id=run_id,
-            tag=tag,
-            parameters_file=parameters_file,
-            is_job=True,
-        )
-
-        assert isinstance(run_context.executor, BaseJobExecutor)
-        run_context.from_sdk = True
-
-        utils.set_runnable_environment_variables(
-            run_id=run_id, configuration_file=configuration_file, tag=tag
+            execution_context=context.ExecutionContext.JOB,
         )
 
-        console.print("Working with context:")
-        console.print(run_context)
-        console.rule(style="[dark orange]")
-
-        if not run_context.executor._is_local:
-            # We are not working with executor that does not work in local environment
-            import inspect
-
-            caller_stack = inspect.stack()[1]
-            relative_to_root = str(Path(caller_stack.filename).relative_to(Path.cwd()))
-
-            module_name = re.sub(r"\b.py\b", "", relative_to_root.replace("/", "."))
-            module_to_call = f"{module_name}.{caller_stack.function}"
-
-            run_context.job_definition_file = f"{module_to_call}.py"
-
-        job = self.get_task()
-        catalog_settings = self.return_catalog_settings()
+        configurations = {
+            "job_definition_file": self.get_caller(),
+            "parameters_file": parameters_file,
+            "tag": tag,
+            "run_id": job_id,
+            "execution_mode": context.ExecutionMode.PYTHON,
+            "configuration_file": configuration_file,
+            "job": self.get_task(),
+            "catalog_settings": self.return_catalog_settings(),
+            **service_configurations.services,
+        }
 
-        try:
-            run_context.executor.submit_job(job, catalog_settings=catalog_settings)
-        finally:
-            run_context.executor.add_task_log_to_catalog("job")
+        run_context = context.JobContext.model_validate(configurations)
 
-        logger.info(
-            "Executing the job from the user. We are still in the caller's compute environment"
-        )
+        assert isinstance(run_context.job_executor, BaseJobExecutor)
 
-        if run_context.executor._is_local:
-            return run_context.run_log_store.get_run_log_by_id(
-                run_id=run_context.run_id
-            )
+        run_context.execute()
 
 
 class PythonJob(BaseJob):

runnable/tasks.py

@@ -25,7 +25,7 @@ from runnable.datastore import (
     Parameter,
     StepAttempt,
 )
-from runnable.defaults import TypeMapVariable
+from runnable.defaults import MapVariableType
 
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
@@ -48,7 +48,29 @@ class TeeIO(io.StringIO):
         self.output_stream.flush()
 
 
-sys.stdout = TeeIO()
+@contextlib.contextmanager
+def redirect_output():
+    # Set the stream handlers to use the custom TeeIO class
+
+    # Backup the original stdout and stderr
+    original_stdout = sys.stdout
+    original_stderr = sys.stderr
+
+    # Redirect stdout and stderr to custom TeeStream objects
+    sys.stdout = TeeIO(sys.stdout)
+    sys.stderr = TeeIO(sys.stderr)
+
+    # Replace stream for all StreamHandlers to use the new sys.stdout
+    for handler in logging.getLogger().handlers:
+        if isinstance(handler, logging.StreamHandler):
+            handler.stream = sys.stdout
+
+    try:
+        yield sys.stdout, sys.stderr
+    finally:
+        # Restore the original stdout and stderr
+        sys.stdout = original_stdout
+        sys.stderr = original_stderr
 
 
 class TaskReturns(BaseModel):
@@ -79,7 +101,7 @@ class BaseTaskType(BaseModel):
     def set_secrets_as_env_variables(self):
         # Preparing the environment for the task execution
         for key in self.secrets:
-            secret_value = context.run_context.secrets_handler.get(key)
+            secret_value = context.run_context.secrets.get(key)
             os.environ[key] = secret_value
 
     def delete_secrets_from_env_variables(self):
@@ -90,7 +112,7 @@ class BaseTaskType(BaseModel):
 
     def execute_command(
         self,
-        map_variable: TypeMapVariable = None,
+        map_variable: MapVariableType = None,
     ) -> StepAttempt:
         """The function to execute the command.
 
@@ -130,7 +152,7 @@ class BaseTaskType(BaseModel):
         finally:
             self.delete_secrets_from_env_variables()
 
-    def resolve_unreduced_parameters(self, map_variable: TypeMapVariable = None):
+    def resolve_unreduced_parameters(self, map_variable: MapVariableType = None):
         """Resolve the unreduced parameters."""
         params = self._context.run_log_store.get_parameters(
             run_id=self._context.run_id
@@ -153,7 +175,7 @@ class BaseTaskType(BaseModel):
 
     @contextlib.contextmanager
     def execution_context(
-        self, map_variable: TypeMapVariable = None, allow_complex: bool = True
+        self, map_variable: MapVariableType = None, allow_complex: bool = True
     ):
         params = self.resolve_unreduced_parameters(map_variable=map_variable)
         logger.info(f"Parameters available for the execution: {params}")
@@ -267,7 +289,7 @@ class PythonTaskType(BaseTaskType):  # pylint: disable=too-few-public-methods
 
     def execute_command(
         self,
-        map_variable: TypeMapVariable = None,
+        map_variable: MapVariableType = None,
     ) -> StepAttempt:
         """Execute the notebook as defined by the command."""
         attempt_log = StepAttempt(status=defaults.FAIL, start_time=str(datetime.now()))
@@ -289,13 +311,21 @@ class PythonTaskType(BaseTaskType):  # pylint: disable=too-few-public-methods
                     logger.info(
                         f"Calling {func} from {module} with {filtered_parameters}"
                     )
-
-                    out_file = TeeIO()
-                    with contextlib.redirect_stdout(out_file):
+                    context.progress.stop()  # redirecting stdout clashes with rich progress
+                    with redirect_output() as (buffer, stderr_buffer):
                         user_set_parameters = f(
                             **filtered_parameters
                         )  # This is a tuple or single value
-                    task_console.print(out_file.getvalue())
+
+                        print(
+                            stderr_buffer.getvalue()
+                        )  # To print the logging statements
+
+                    # TODO: Avoid double print!!
+                    with task_console.capture():
+                        task_console.log(buffer.getvalue())
+                        task_console.log(stderr_buffer.getvalue())
+                    context.progress.start()
                 except Exception as e:
                     raise exceptions.CommandCallError(
                         f"Function call: {self.command} did not succeed.\n"
@@ -478,14 +508,15 @@ class NotebookTaskType(BaseTaskType):
 
         return command
 
-    def get_notebook_output_path(self, map_variable: TypeMapVariable = None) -> str:
+    def get_notebook_output_path(self, map_variable: MapVariableType = None) -> str:
         tag = ""
         map_variable = map_variable or {}
         for key, value in map_variable.items():
             tag += f"{key}_{value}_"
 
-        if hasattr(self._context.executor, "_context_node"):
-            tag += self._context.executor._context_node.name
+        if isinstance(self._context, context.PipelineContext):
+            assert self._context.pipeline_executor._context_node
+            tag += self._context.pipeline_executor._context_node.name
 
         tag = "".join(x for x in tag if x.isalnum()).strip("-")
 
@@ -496,7 +527,7 @@ class NotebookTaskType(BaseTaskType):
 
     def execute_command(
         self,
-        map_variable: TypeMapVariable = None,
+        map_variable: MapVariableType = None,
     ) -> StepAttempt:
         """Execute the python notebook as defined by the command.
 
@@ -551,12 +582,20 @@ class NotebookTaskType(BaseTaskType):
                 }
                 kwds.update(ploomber_optional_args)
 
-                out_file = TeeIO()
-                with contextlib.redirect_stdout(out_file):
+                context.progress.stop()  # redirecting stdout clashes with rich progress
+
+                with redirect_output() as (buffer, stderr_buffer):
                     pm.execute_notebook(**kwds)
-                task_console.print(out_file.getvalue())
 
-                context.run_context.catalog_handler.put(name=notebook_output_path)
+                    print(stderr_buffer.getvalue())  # To print the logging statements
+
+                with task_console.capture():
+                    task_console.log(buffer.getvalue())
+                    task_console.log(stderr_buffer.getvalue())
+
+                context.progress.start()
+
+                context.run_context.catalog.put(name=notebook_output_path)
 
                 client = PloomberClient.from_path(path=notebook_output_path)
                 namespace = client.get_namespace()
@@ -674,7 +713,7 @@ class ShellTaskType(BaseTaskType):
 
     def execute_command(
         self,
-        map_variable: TypeMapVariable = None,
+        map_variable: MapVariableType = None,
     ) -> StepAttempt:
         # Using shell=True as we want to have chained commands to be executed in the same shell.
         """Execute the shell command as defined by the command.
@@ -698,7 +737,7 @@ class ShellTaskType(BaseTaskType):
         # Expose secrets as environment variables
         if self.secrets:
             for key in self.secrets:
-                secret_value = context.run_context.secrets_handler.get(key)
+                secret_value = context.run_context.secrets.get(key)
                 subprocess_env[key] = secret_value
 
         try:
@@ -724,6 +763,7 @@ class ShellTaskType(BaseTaskType):
                 capture = False
                 return_keys = {x.name: x for x in self.returns}
 
+                context.progress.stop()  # redirecting stdout clashes with rich progress
                 proc = subprocess.Popen(
                     command,
                     shell=True,
@@ -747,6 +787,7 @@ class ShellTaskType(BaseTaskType):
                         continue
                     task_console.print(line, style=defaults.warning_style)
 
+                context.progress.start()
                 output_parameters: Dict[str, Parameter] = {}
                 metrics: Dict[str, Parameter] = {}