runnable 0.13.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

runnable/parameters.py

@@ -15,6 +15,8 @@ from runnable.utils import remove_prefix
 
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
+# TODO: Revisit this, it might be a bit too complicated than required
+
 
 def get_user_set_parameters(remove: bool = False) -> Dict[str, JsonParameter]:
     """
@@ -34,9 +36,13 @@ def get_user_set_parameters(remove: bool = False) -> Dict[str, JsonParameter]:
         if env_var.startswith(defaults.PARAMETER_PREFIX):
             key = remove_prefix(env_var, defaults.PARAMETER_PREFIX)
             try:
-                parameters[key.lower()] = JsonParameter(kind="json", value=json.loads(value))
+                parameters[key.lower()] = JsonParameter(
+                    kind="json", value=json.loads(value)
+                )
             except json.decoder.JSONDecodeError:
-                logger.warning(f"Parameter {key} could not be JSON decoded, adding the literal value")
+                logger.warning(
+                    f"Parameter {key} could not be JSON decoded, adding the literal value"
+                )
                 parameters[key.lower()] = JsonParameter(kind="json", value=value)
 
             if remove:
@@ -52,7 +58,9 @@ def serialize_parameter_as_str(value: Any) -> str:
 
 
 def filter_arguments_for_func(
-    func: Callable[..., Any], params: Dict[str, Any], map_variable: TypeMapVariable = None
+    func: Callable[..., Any],
+    params: Dict[str, Any],
+    map_variable: TypeMapVariable = None,
 ) -> Dict[str, Any]:
     """
     Inspects the function to be called as part of the pipeline to find the arguments of the function.
@@ -96,11 +104,16 @@ def filter_arguments_for_func(
             # No parameter of this name was provided
             if value.default == inspect.Parameter.empty:
                 # No default value is given in the function signature. error as parameter is required.
-                raise ValueError(f"Parameter {name} is required for {func.__name__} but not provided")
+                raise ValueError(
+                    f"Parameter {name} is required for {func.__name__} but not provided"
+                )
             # default value is given in the function signature, nothing further to do.
             continue
 
-        if type(value.annotation) in [BaseModel, pydantic._internal._model_construction.ModelMetaclass]:
+        if type(value.annotation) in [
+            BaseModel,
+            pydantic._internal._model_construction.ModelMetaclass,
+        ]:
             # We try to cast it as a pydantic model if asked
             named_param = params[name].get_value()
 
@@ -110,7 +123,9 @@ def filter_arguments_for_func(
 
             bound_model = bind_args_for_pydantic_model(named_param, value.annotation)
             bound_args[name] = bound_model
-            unassigned_params = unassigned_params.difference(bound_model.model_fields.keys())
+            unassigned_params = unassigned_params.difference(
+                bound_model.model_fields.keys()
+            )
 
         elif value.annotation in [str, int, float, bool]:
             # Cast it if its a primitive type. Ensure the type matches the annotation.
@@ -120,12 +135,16 @@ def filter_arguments_for_func(
 
         unassigned_params.remove(name)
 
-        params = {key: params[key] for key in unassigned_params}  # remove keys from params if they are assigned
+        params = {
+            key: params[key] for key in unassigned_params
+        }  # remove keys from params if they are assigned
 
     return bound_args
 
 
-def bind_args_for_pydantic_model(params: Dict[str, Any], model: Type[BaseModel]) -> BaseModel:
+def bind_args_for_pydantic_model(
+    params: Dict[str, Any], model: Type[BaseModel]
+) -> BaseModel:
     class EasyModel(model):  # type: ignore
         model_config = ConfigDict(extra="ignore")
 

runnable/pickler.py

@@ -9,7 +9,7 @@ import runnable.context as context
 
 class BasePickler(ABC, BaseModel):
     """
-    The base class for all picklers.
+    The base class for all pickler.
 
     We are still in the process of hardening the design of this class.
     For now, we are just going to use pickle.

runnable/sdk.py

@@ -26,8 +26,7 @@ from rich.progress import (
 from rich.table import Column
 from typing_extensions import Self
 
-from runnable import console, defaults, entrypoints, exceptions, graph, utils
-from runnable.extensions.nodes import (
+from extensions.nodes.nodes import (
     FailNode,
     MapNode,
     ParallelNode,
@@ -35,9 +34,14 @@ from runnable.extensions.nodes import (
     SuccessNode,
     TaskNode,
 )
+from runnable import console, defaults, entrypoints, exceptions, graph, utils
+from runnable.executor import BaseJobExecutor, BasePipelineExecutor
 from runnable.nodes import TraversalNode
+from runnable.tasks import BaseTaskType as RunnableTask
 from runnable.tasks import TaskReturns
 
+# TODO: This might have to be an extension
+
 logger = logging.getLogger(defaults.LOGGER_NAME)
 
 StepType = Union["Stub", "PythonTask", "NotebookTask", "ShellTask", "Parallel", "Map"]
@@ -66,7 +70,9 @@ class Catalog(BaseModel):
 
     """
 
-    model_config = ConfigDict(extra="forbid")  # Need to be for command, would be validated later
+    model_config = ConfigDict(
+        extra="forbid"
+    )  # Need to be for command, would be validated later
     # Note: compute_data_folder was confusing to explain, might be introduced later.
     # compute_data_folder: str = Field(default="", alias="compute_data_folder")
     get: List[str] = Field(default_factory=list, alias="get")
@@ -95,14 +101,18 @@ class BaseTraversal(ABC, BaseModel):
 
     def __rshift__(self, other: StepType) -> StepType:
         if self.next_node:
-            raise Exception(f"The node {self} already has a next node: {self.next_node}")
+            raise Exception(
+                f"The node {self} already has a next node: {self.next_node}"
+            )
         self.next_node = other.name
 
         return other
 
     def __lshift__(self, other: TraversalNode) -> TraversalNode:
         if other.next_node:
-            raise Exception(f"The {other} node already has a next node: {other.next_node}")
+            raise Exception(
+                f"The {other} node already has a next node: {other.next_node}"
+            )
         other.next_node = self.name
 
         return other
@@ -112,7 +122,9 @@ class BaseTraversal(ABC, BaseModel):
         assert not isinstance(node, Fail)
 
         if node.next_node:
-            raise Exception(f"The {node} node already has a next node: {node.next_node}")
+            raise Exception(
+                f"The {node} node already has a next node: {node.next_node}"
+            )
 
         node.next_node = self.name
         return self
@@ -124,7 +136,9 @@ class BaseTraversal(ABC, BaseModel):
 
         if self.terminate_with_failure or self.terminate_with_success:
             if self.next_node and self.next_node not in ["success", "fail"]:
-                raise AssertionError("A node being terminated cannot have a user defined next node")
+                raise AssertionError(
+                    "A node being terminated cannot have a user defined next node"
+                )
 
         if self.terminate_with_failure:
             self.next_node = "fail"
@@ -135,8 +149,7 @@ class BaseTraversal(ABC, BaseModel):
         return self
 
     @abstractmethod
-    def create_node(self) -> TraversalNode:
-        ...
+    def create_node(self) -> TraversalNode: ...
 
 
 class BaseTask(BaseTraversal):
@@ -146,12 +159,16 @@ class BaseTask(BaseTraversal):
 
     catalog: Optional[Catalog] = Field(default=None, alias="catalog")
     overrides: Dict[str, Any] = Field(default_factory=dict, alias="overrides")
-    returns: List[Union[str, TaskReturns]] = Field(default_factory=list, alias="returns")
+    returns: List[Union[str, TaskReturns]] = Field(
+        default_factory=list, alias="returns"
+    )
     secrets: List[str] = Field(default_factory=list)
 
     @field_validator("returns", mode="before")
     @classmethod
-    def serialize_returns(cls, returns: List[Union[str, TaskReturns]]) -> List[TaskReturns]:
+    def serialize_returns(
+        cls, returns: List[Union[str, TaskReturns]]
+    ) -> List[TaskReturns]:
         task_returns = []
 
         for x in returns:
@@ -167,9 +184,18 @@ class BaseTask(BaseTraversal):
     def create_node(self) -> TaskNode:
         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")
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
 
-        return TaskNode.parse_from_config(self.model_dump(exclude_none=True, by_alias=True))
+        return TaskNode.parse_from_config(
+            self.model_dump(exclude_none=True, by_alias=True)
+        )
+
+    def create_job(self) -> RunnableTask:
+        raise NotImplementedError(
+            "This method should be implemented in the child class"
+        )
 
 
 class PythonTask(BaseTask):
@@ -254,6 +280,11 @@ class PythonTask(BaseTask):
 
         return f"{module}.{name}"
 
+    def create_job(self) -> RunnableTask:
+        self.terminate_with_success = True
+        node = self.create_node()
+        return node.executable
+
 
 class NotebookTask(BaseTask):
     """
@@ -326,12 +357,19 @@ class NotebookTask(BaseTask):
     """
 
     notebook: str = Field(serialization_alias="command")
-    optional_ploomber_args: Optional[Dict[str, Any]] = Field(default=None, alias="optional_ploomber_args")
+    optional_ploomber_args: Optional[Dict[str, Any]] = Field(
+        default=None, alias="optional_ploomber_args"
+    )
 
     @computed_field
     def command_type(self) -> str:
         return "notebook"
 
+    def create_job(self) -> RunnableTask:
+        self.terminate_with_success = True
+        node = self.create_node()
+        return node.executable
+
 
 class ShellTask(BaseTask):
     """
@@ -416,7 +454,9 @@ class Stub(BaseTraversal):
     def create_node(self) -> StubNode:
         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")
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
 
         return StubNode.parse_from_config(self.model_dump(exclude_none=True))
 
@@ -439,14 +479,23 @@ class Parallel(BaseTraversal):
     @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()}
+        return {
+            name: pipeline._dag.model_copy() for name, pipeline in self.branches.items()
+        }
 
     def create_node(self) -> ParallelNode:
         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")
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
 
-        node = ParallelNode(name=self.name, branches=self.graph_branches, internal_name="", next_node=self.next_node)
+        node = ParallelNode(
+            name=self.name,
+            branches=self.graph_branches,
+            internal_name="",
+            next_node=self.next_node,
+        )
         return node
 
 
@@ -483,7 +532,9 @@ class Map(BaseTraversal):
     def create_node(self) -> MapNode:
         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")
+                raise AssertionError(
+                    "A node not being terminated must have a user defined next node"
+                )
 
         node = MapNode(
             name=self.name,
@@ -587,6 +638,7 @@ class Pipeline(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     def _validate_path(self, path: List[StepType], failure_path: bool = False) -> None:
+        # TODO: Drastically simplify this
         # Check if one and only one step terminates with success
         # Check no more than one step terminates with failure
 
@@ -596,16 +648,22 @@ class Pipeline(BaseModel):
         for step in path:
             if step.terminate_with_success:
                 if reached_success:
-                    raise Exception("A pipeline cannot have more than one step that terminates with success")
+                    raise Exception(
+                        "A pipeline cannot have more than one step that terminates with success"
+                    )
                 reached_success = True
                 continue
             if step.terminate_with_failure:
                 if reached_failure:
-                    raise Exception("A pipeline cannot have more than one step that terminates with failure")
+                    raise Exception(
+                        "A pipeline cannot have more than one step that terminates with failure"
+                    )
                 reached_failure = True
 
         if not reached_success and not reached_failure:
-            raise Exception("A pipeline must have at least one step that terminates with success")
+            raise Exception(
+                "A pipeline must have at least one step that terminates with success"
+            )
 
     def _construct_path(self, path: List[StepType]) -> None:
         prev_step = path[0]
@@ -615,7 +673,9 @@ class Pipeline(BaseModel):
                 continue
 
             if prev_step.terminate_with_success or prev_step.terminate_with_failure:
-                raise Exception(f"A step that terminates with success/failure cannot have a next step: {prev_step}")
+                raise Exception(
+                    f"A step that terminates with success/failure cannot have a next step: {prev_step}"
+                )
 
             if prev_step.next_node and prev_step.next_node not in ["success", "fail"]:
                 raise Exception(f"Step already has a next node: {prev_step} ")
@@ -646,7 +706,9 @@ class Pipeline(BaseModel):
         on_failure_paths: List[List[StepType]] = []
 
         for step in self.steps:
-            if isinstance(step, (Stub, PythonTask, NotebookTask, ShellTask, Parallel, Map)):
+            if isinstance(
+                step, (Stub, PythonTask, NotebookTask, ShellTask, Parallel, Map)
+            ):
                 success_path.append(step)
                 continue
             # on_failure_paths.append(step)
@@ -690,6 +752,16 @@ class Pipeline(BaseModel):
         dag_definition = self._dag.model_dump(by_alias=True, exclude_none=True)
         return graph.create_graph(dag_definition)
 
+    def _is_called_for_definition(self) -> bool:
+        """
+        If the run context is set, we are coming in only to get the pipeline definition.
+        """
+        from runnable.context import run_context
+
+        if run_context is None:
+            return False
+        return True
+
     def execute(
         self,
         configuration_file: str = "",
@@ -699,39 +771,20 @@ class Pipeline(BaseModel):
         log_level: str = defaults.LOG_LEVEL,
     ):
         """
-        *Execute* the Pipeline.
-
-        Execution of pipeline could either be:
-
-        Traverse and execute all the steps of the pipeline, eg. [local execution](configurations/executors/local.md).
-
-        Or create the representation of the pipeline for other executors.
-
-        Please refer to [concepts](concepts/executor.md) for more information.
-
-        Args:
-            configuration_file (str, optional): The path to the configuration file. Defaults to "".
-                The configuration file can be overridden by the environment variable RUNNABLE_CONFIGURATION_FILE.
-
-            run_id (str, optional): The ID of the run. Defaults to "".
-            tag (str, optional): The tag of the run. Defaults to "".
-                Use to group multiple runs.
-
-            parameters_file (str, optional): The path to the parameters file. Defaults to "".
-
-            log_level (str, optional): The log level. Defaults to defaults.LOG_LEVEL.
+        Overloaded method:
+        - Could be called by the user when executing the pipeline via SDK
+        - Could be called by the system itself when getting the pipeline definition
         """
-
-        # py_to_yaml is used by non local executors to generate the yaml representation of the pipeline.
-        py_to_yaml = os.environ.get("RUNNABLE_PY_TO_YAML", "false")
-
-        if py_to_yaml == "true":
+        if self._is_called_for_definition():
+            # 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)
-        configuration_file = os.environ.get("RUNNABLE_CONFIGURATION_FILE", configuration_file)
+        configuration_file = os.environ.get(
+            "RUNNABLE_CONFIGURATION_FILE", configuration_file
+        )
         run_context = entrypoints.prepare_configurations(
             configuration_file=configuration_file,
             run_id=run_id,
@@ -739,19 +792,22 @@ class Pipeline(BaseModel):
             parameters_file=parameters_file,
         )
 
-        run_context.execution_plan = defaults.EXECUTION_PLAN.CHAINED.value
-        utils.set_runnable_environment_variables(run_id=run_id, configuration_file=configuration_file, tag=tag)
+        assert isinstance(run_context.executor, BasePipelineExecutor)
 
-        dag_definition = self._dag.model_dump(by_alias=True, exclude_none=True)
+        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._local:
-            # We are not working with non local executor
+        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]
@@ -761,37 +817,145 @@ class Pipeline(BaseModel):
             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.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)),
+            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
-                pipeline_execution_task = progress.add_task("[dark_orange] Starting execution .. ", total=1)
+
                 run_context.executor.execute_graph(dag=run_context.dag)
 
-                if not run_context.executor._local:
+                if not run_context.executor._is_local:
+                    # non local executors just traverse the graph and do nothing
                     return {}
 
-                run_log = run_context.run_log_store.get_run_log_by_id(run_id=run_context.run_id, full=False)
+                run_log = run_context.run_log_store.get_run_log_by_id(
+                    run_id=run_context.run_id, full=False
+                )
 
                 if run_log.status == defaults.SUCCESS:
-                    progress.update(pipeline_execution_task, description="[green] Success", completed=True)
+                    progress.update(
+                        pipeline_execution_task,
+                        description="[green] Success",
+                        completed=True,
+                    )
                 else:
-                    progress.update(pipeline_execution_task, description="[red] Failed", completed=True)
+                    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)
+                progress.update(
+                    pipeline_execution_task,
+                    description="[red] Errored execution",
+                    completed=True,
+                )
                 raise
 
-        if run_context.executor._local:
-            return run_context.run_log_store.get_run_log_by_id(run_id=run_context.run_id)
+        if run_context.executor._is_local:
+            return run_context.run_log_store.get_run_log_by_id(
+                run_id=run_context.run_id
+            )
+
+
+class Job(BaseModel):
+    name: str
+    task: BaseTask
+
+    def return_task(self) -> RunnableTask:
+        return self.task.create_job()
+
+    def return_catalog_settings(self) -> Optional[List[str]]:
+        if self.task.catalog is None:
+            return []
+        return self.task.catalog.put
+
+    def _is_called_for_definition(self) -> bool:
+        """
+        If the run context is set, we are coming in only to get the pipeline definition.
+        """
+        from runnable.context import run_context
+
+        if run_context is None:
+            return False
+        return True
+
+    def execute(
+        self,
+        configuration_file: str = "",
+        job_id: str = "",
+        tag: str = "",
+        parameters_file: str = "",
+        log_level: str = defaults.LOG_LEVEL,
+    ):
+        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)
+        configuration_file = os.environ.get(
+            "RUNNABLE_CONFIGURATION_FILE", configuration_file
+        )
+        run_context = entrypoints.prepare_configurations(
+            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
+        )
+
+        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.task.create_job()
+        catalog_settings = self.return_catalog_settings()
+
+        run_context.executor.submit_job(job, catalog_settings=catalog_settings)
+
+        logger.info(
+            "Executing the job from the user. We are still in the caller's compute environment"
+        )
+
+        if run_context.executor._is_local:
+            return run_context.run_log_store.get_run_log_by_id(
+                run_id=run_context.run_id
+            )

runnable/secrets.py

@@ -92,4 +92,6 @@ class EnvSecretsManager(BaseSecrets):
         try:
             return os.environ[name]
         except KeyError:
-            raise exceptions.SecretNotFoundError(secret_name=name, secret_setting="environment variables")
+            raise exceptions.SecretNotFoundError(
+                secret_name=name, secret_setting="environment variables"
+            )