runnable 0.21.0__tar.gz → 0.23.0__tar.gz

{runnable-0.21.0 → runnable-0.23.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: runnable
-Version: 0.21.0
+Version: 0.23.0
 Summary: Add your description here
 Author-email: "Vammi, Vijay" <vijay.vammi@astrazeneca.com>
 License-File: LICENSE

{runnable-0.21.0 → runnable-0.23.0}/extensions/job_executor/__init__.py

@@ -2,7 +2,7 @@ import logging
 import os
 from typing import Dict, List, Optional
 
-from runnable import context, defaults, exceptions, parameters, utils
+from runnable import context, defaults, exceptions, parameters, task_console, utils
 from runnable.datastore import DataCatalog, JobLog, JsonParameter
 from runnable.executor import BaseJobExecutor
 
@@ -158,3 +158,17 @@ class GenericJobExecutor(BaseJobExecutor):
             data_catalogs.extend(data_catalog)
 
         return data_catalogs
+
+    def add_task_log_to_catalog(
+        self, name: str, map_variable: Dict[str, str | int | float] | None = None
+    ):
+        log_file_name = utils.make_log_file_name(
+            name=name,
+            map_variable=map_variable,
+        )
+        task_console.save_text(log_file_name)
+        # Put the log file in the catalog
+        self._context.catalog_handler.put(
+            name=log_file_name, run_id=self._context.run_id
+        )
+        os.remove(log_file_name)

{runnable-0.21.0 → runnable-0.23.0}/extensions/pipeline_executor/__init__.py

@@ -210,6 +210,20 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         """
         return int(os.environ.get(defaults.ATTEMPT_NUMBER, 1))
 
+    def add_task_log_to_catalog(
+        self, name: str, map_variable: Dict[str, str | int | float] | None = None
+    ):
+        log_file_name = utils.make_log_file_name(
+            name=name,
+            map_variable=map_variable,
+        )
+        task_console.save_text(log_file_name)
+        # Put the log file in the catalog
+        self._context.catalog_handler.put(
+            name=log_file_name, run_id=self._context.run_id
+        )
+        os.remove(log_file_name)
+
     def _execute_node(
         self,
         node: BaseNode,
@@ -336,14 +350,6 @@ class GenericPipelineExecutor(BasePipelineExecutor):
         )
         self.trigger_node_execution(node=node, map_variable=map_variable, **kwargs)
 
-        log_file_name = utils.make_log_file_name(node=node, map_variable=map_variable)
-        task_console.save_text(log_file_name, clear=True)
-
-        self._context.catalog_handler.put(
-            name=log_file_name, run_id=self._context.run_id
-        )
-        os.remove(log_file_name)
-
     def trigger_node_execution(
         self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs
     ):
@@ -482,6 +488,11 @@ 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]")
 

{runnable-0.21.0 → runnable-0.23.0}/extensions/pipeline_executor/argo.py

@@ -6,7 +6,7 @@ import string
 from collections import namedtuple
 from enum import Enum
 from functools import cached_property
-from typing import Annotated, Literal, Optional
+from typing import Annotated, Any, Literal, Optional, cast
 
 from pydantic import (
     BaseModel,
@@ -20,7 +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
+from extensions.nodes.nodes import (
+    MapNode,
+    ParallelNode,
+    StubNode,
+    SuccessNode,
+    TaskNode,
+)
 from extensions.pipeline_executor import GenericPipelineExecutor
 from runnable import defaults, utils
 from runnable.defaults import TypeMapVariable
@@ -229,9 +235,8 @@ class Resources(BaseModel):
     requests: Request = Field(default=Request(), serialization_alias="requests")
 
 
-# This is what the user can override per template
-# Some are specific to container and some are specific to dag
-class TemplateDefaults(BaseModelWIthConfig):
+# Lets construct this from UserDefaults
+class ArgoTemplateDefaults(BaseModelWIthConfig):
     active_deadline_seconds: Optional[int] = Field(default=86400)  # 1 day
     fail_fast: bool = Field(default=True)
     node_selector: dict[str, str] = Field(default_factory=dict)
@@ -240,10 +245,32 @@ class TemplateDefaults(BaseModelWIthConfig):
     timeout: Optional[str] = Field(default=None)
     tolerations: Optional[list[Toleration]] = Field(default=None)
 
-    # These are in addition to what argo spec provides
-    image: str
-    image_pull_policy: Optional[ImagePullPolicy] = Field(default=ImagePullPolicy.Always)
+    model_config = ConfigDict(
+        extra="ignore",
+    )
+
+
+class CommonDefaults(BaseModelWIthConfig):
+    active_deadline_seconds: Optional[int] = Field(default=86400)  # 1 day
+    fail_fast: bool = Field(default=True)
+    node_selector: dict[str, str] = Field(default_factory=dict)
+    parallelism: Optional[int] = Field(default=None)
+    retry_strategy: Optional[RetryStrategy] = Field(default=None)
+    timeout: Optional[str] = Field(default=None)
+    tolerations: Optional[list[Toleration]] = Field(default=None)
+    image_pull_policy: ImagePullPolicy = Field(default=ImagePullPolicy.Always)
     resources: Resources = Field(default_factory=Resources)
+    env: list[EnvVar | SecretEnvVar] = Field(default_factory=list, exclude=True)
+
+
+# The user provided defaults at the top level
+class UserDefaults(CommonDefaults):
+    image: str
+
+
+# Overrides need not have image
+class Overrides(CommonDefaults):
+    image: Optional[str] = Field(default=None)
 
 
 # User provides this as part of the argoSpec
@@ -257,8 +284,8 @@ class ArgoWorkflowSpec(BaseModelWIthConfig):
     pod_gc: Optional[PodGC] = Field(default=None, serialization_alias="podGC")
     retry_strategy: Optional[RetryStrategy] = Field(default=None)
     service_account_name: Optional[str] = Field(default=None)
-    template_defaults: TemplateDefaults
     tolerations: Optional[list[Toleration]] = Field(default=None)
+    template_defaults: Optional[ArgoTemplateDefaults] = Field(default=None)
 
 
 class ArgoMetadata(BaseModelWIthConfig):
@@ -320,7 +347,6 @@ class ContainerTemplate((BaseModelWIthConfig)):
     inputs: Optional[Inputs] = Field(default=None)
     outputs: Optional[Outputs] = Field(default=None)
 
-    # The remaining can be from template defaults or node overrides
     active_deadline_seconds: Optional[int] = Field(default=86400)  # 1 day
     metadata: Optional[PodMetaData] = Field(default=None)
     node_selector: dict[str, str] = Field(default_factory=dict)
@@ -355,13 +381,7 @@ class ArgoExecutor(GenericPipelineExecutor):
         from_attributes=True,
         use_enum_values=True,
     )
-
-    argo_workflow: ArgoWorkflow
-
-    # Lets use a generic one
     pvc_for_runnable: Optional[str] = Field(default=None)
-    # pvc_for_catalog: Optional[str] = Field(default=None)
-    # pvc_for_run_log: Optional[str] = Field(default=None)
     custom_volumes: Optional[list[CustomVolume]] = Field(
         default_factory=list[CustomVolume]
     )
@@ -373,6 +393,11 @@ class ArgoExecutor(GenericPipelineExecutor):
         default="INFO"
     )
 
+    defaults: UserDefaults
+    argo_workflow: ArgoWorkflow
+
+    overrides: dict[str, Overrides] = Field(default_factory=dict)
+
     # This should be used when we refer to run_id or log_level in the containers
     _run_id_as_parameter: str = PrivateAttr(default="{{workflow.parameters.run_id}}")
     _log_level_as_parameter: str = PrivateAttr(
@@ -386,6 +411,11 @@ class ArgoExecutor(GenericPipelineExecutor):
     _container_catalog_location: str = PrivateAttr(default="/tmp/catalog/")
     _added_initial_container: bool = PrivateAttr(default=False)
 
+    def model_post_init(self, __context: Any) -> None:
+        self.argo_workflow.spec.template_defaults = ArgoTemplateDefaults(
+            **self.defaults.model_dump()
+        )
+
     def sanitize_name(self, name: str) -> str:
         formatted_name = name.replace(" ", "-").replace(".", "-").replace("_", "-")
         tag = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
@@ -423,8 +453,6 @@ class ArgoExecutor(GenericPipelineExecutor):
         parameters: Optional[list[Parameter]],
         task_name: str,
     ):
-        template_defaults = self.argo_workflow.spec.template_defaults.model_dump()
-
         map_variable: TypeMapVariable = {}
         for parameter in parameters or []:
             map_variable[parameter.name] = (  # type: ignore
@@ -440,8 +468,8 @@ class ArgoExecutor(GenericPipelineExecutor):
 
         core_container_template = CoreContainerTemplate(
             command=shlex.split(fan_command),
-            image=template_defaults["image"],
-            image_pull_policy=template_defaults["image_pull_policy"],
+            image=self.defaults.image,
+            image_pull_policy=self.defaults.image_pull_policy,
             volume_mounts=[
                 volume_pair.volume_mount for volume_pair in self.volume_pairs
             ],
@@ -457,12 +485,17 @@ class ArgoExecutor(GenericPipelineExecutor):
             outputs = Outputs(parameters=[OutputParameter(name="iterate-on")])
 
         container_template = ContainerTemplate(
-            container=core_container_template,
             name=task_name,
-            volumes=[volume_pair.volume for volume_pair in self.volume_pairs],
+            container=core_container_template,
             inputs=Inputs(parameters=parameters),
             outputs=outputs,
-            **template_defaults,
+            active_deadline_seconds=self.defaults.active_deadline_seconds,
+            node_selector=self.defaults.node_selector,
+            parallelism=self.defaults.parallelism,
+            retry_strategy=self.defaults.retry_strategy,
+            timeout=self.defaults.timeout,
+            tolerations=self.defaults.tolerations,
+            volumes=[volume_pair.volume for volume_pair in self.volume_pairs],
         )
 
         self._templates.append(container_template)
@@ -473,14 +506,23 @@ class ArgoExecutor(GenericPipelineExecutor):
         task_name: str,
         inputs: Optional[Inputs] = None,
     ) -> ContainerTemplate:
-        template_defaults = self.argo_workflow.spec.template_defaults.model_dump()
+        assert (
+            isinstance(node, TaskNode)
+            or isinstance(node, StubNode)
+            or isinstance(node, SuccessNode)
+        )
 
-        node_overide = {}
-        if hasattr(node, "overides"):
-            node_overide = node.overides
+        node_override = None
+        if hasattr(node, "overrides"):
+            override_key = node.overrides.get(self.service_name, "")
+            try:
+                node_override = self.overrides.get(override_key)
+            except:  # noqa
+                raise Exception("Override not found for: ", override_key)
 
-        # update template defaults with node overrides
-        template_defaults.update(node_overide)
+        effective_settings = self.defaults.model_dump()
+        if node_override:
+            effective_settings.update(node_override.model_dump())
 
         inputs = inputs or Inputs(parameters=[])
 
@@ -500,35 +542,66 @@ class ArgoExecutor(GenericPipelineExecutor):
 
         core_container_template = CoreContainerTemplate(
             command=shlex.split(command),
-            image=template_defaults["image"],
-            image_pull_policy=template_defaults["image_pull_policy"],
+            image=effective_settings["image"],
+            image_pull_policy=effective_settings["image_pull_policy"],
+            resources=effective_settings["resources"],
             volume_mounts=[
                 volume_pair.volume_mount for volume_pair in self.volume_pairs
             ],
         )
 
         self._set_up_initial_container(container_template=core_container_template)
+        self._expose_secrets_to_task(
+            working_on=node, container_template=core_container_template
+        )
+        self._set_env_vars_to_task(node, core_container_template)
 
         container_template = ContainerTemplate(
-            container=core_container_template,
             name=task_name,
+            container=core_container_template,
             inputs=Inputs(
                 parameters=[
                     Parameter(name=param.name) for param in inputs.parameters or []
                 ]
             ),
             volumes=[volume_pair.volume for volume_pair in self.volume_pairs],
-            **template_defaults,
+            **node_override.model_dump() if node_override else {},
         )
 
         return container_template
 
+    def _set_env_vars_to_task(
+        self, working_on: BaseNode, container_template: CoreContainerTemplate
+    ):
+        if not isinstance(working_on, TaskNode):
+            return
+
+        global_envs: dict[str, str] = {}
+
+        for env_var in self.defaults.env:
+            env_var = cast(EnvVar, env_var)
+            global_envs[env_var.name] = env_var.value
+
+        override_key = working_on.overrides.get(self.service_name, "")
+        node_override = self.overrides.get(override_key, None)
+
+        # Update the global envs with the node overrides
+        if node_override:
+            for env_var in node_override.env:
+                env_var = cast(EnvVar, env_var)
+                global_envs[env_var.name] = env_var.value
+
+        for key, value in global_envs.items():
+            env_var_to_add = EnvVar(name=key, value=value)
+            container_template.env.append(env_var_to_add)
+
     def _expose_secrets_to_task(
         self,
         working_on: BaseNode,
         container_template: CoreContainerTemplate,
     ):
-        assert isinstance(working_on, TaskNode)
+        if not isinstance(working_on, TaskNode):
+            return
         secrets = working_on.executable.secrets
         for secret in secrets:
             assert self.secret_from_k8s is not None
@@ -770,11 +843,6 @@ class ArgoExecutor(GenericPipelineExecutor):
 
         argo_workflow_dump = self.argo_workflow.model_dump(
             by_alias=True,
-            exclude={
-                "spec": {
-                    "template_defaults": {"image_pull_policy", "image", "resources"}
-                }
-            },
             exclude_none=True,
             round_trip=False,
         )

{runnable-0.21.0 → runnable-0.23.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "runnable"
-version = "0.21.0"
+version = "0.23.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [

{runnable-0.21.0 → runnable-0.23.0}/runnable/entrypoints.py

@@ -345,15 +345,9 @@ def execute_single_node(
     try:
         executor.execute_node(node=node_to_execute, map_variable=map_variable_dict)
     finally:
-        log_file_name = utils.make_log_file_name(
-            node=node_to_execute,
-            map_variable=map_variable_dict,
+        run_context.executor.add_task_log_to_catalog(
+            name=node_to_execute.internal_name, map_variable=map_variable_dict
         )
-        task_console.save_text(log_file_name)
-
-        # Put the log file in the catalog
-        run_context.catalog_handler.put(name=log_file_name, run_id=run_context.run_id)
-        os.remove(log_file_name)
 
     executor.send_return_code()
 
@@ -408,7 +402,10 @@ def execute_job_yaml_spec(
     )
 
     assert isinstance(executor, BaseJobExecutor)
-    executor.submit_job(job, catalog_settings=catalog_config)
+    try:
+        executor.submit_job(job, catalog_settings=catalog_config)
+    finally:
+        run_context.executor.add_task_log_to_catalog("job")
 
     executor.send_return_code()
 
@@ -483,9 +480,12 @@ def execute_job_non_local(
         "Executing the job from the user. We are still in the caller's compute environment"
     )
 
-    run_context.executor.execute_job(
-        run_context.job, catalog_settings=run_context.job_catalog_settings
-    )
+    try:
+        run_context.executor.execute_job(
+            run_context.job, catalog_settings=run_context.job_catalog_settings
+        )
+    finally:
+        run_context.executor.add_task_log_to_catalog("job")
 
     run_context.executor.send_return_code()
 

{runnable-0.21.0 → runnable-0.23.0}/runnable/executor.py

@@ -84,6 +84,11 @@ class BaseExecutor(ABC, BaseModel):
         """
         ...
 
+    @abstractmethod
+    def add_task_log_to_catalog(
+        self, name: str, map_variable: Optional[TypeMapVariable] = None
+    ): ...
+
 
 class BaseJobExecutor(BaseExecutor):
     service_type: str = "job_executor"

{runnable-0.21.0 → runnable-0.23.0}/runnable/sdk.py

@@ -899,7 +899,10 @@ class BaseJob(BaseModel):
         job = self.get_task()
         catalog_settings = self.return_catalog_settings()
 
-        run_context.executor.submit_job(job, catalog_settings=catalog_settings)
+        try:
+            run_context.executor.submit_job(job, catalog_settings=catalog_settings)
+        finally:
+            run_context.executor.add_task_log_to_catalog("job")
 
         logger.info(
             "Executing the job from the user. We are still in the caller's compute environment"

{runnable-0.21.0 → runnable-0.23.0}/runnable/utils.py

@@ -109,8 +109,16 @@ def apply_variables(
         raise Exception("Argument Variables should be dict")
 
     json_d = json.dumps(apply_to)
-    transformed = str_template(json_d).substitute(**variables)
-    return json.loads(transformed)
+    string_template = str_template(json_d)
+
+    template = string_template.safe_substitute(variables)
+
+    if "$" in template:
+        logger.warning(
+            "Not all variables found in the config are found in the variables"
+        )
+
+    return json.loads(template)
 
 
 def get_module_and_attr_names(command: str) -> Tuple[str, str]:
@@ -620,9 +628,9 @@ def gather_variables() -> Dict[str, str]:
     return variables
 
 
-def make_log_file_name(node: BaseNode, map_variable: TypeMapVariable) -> str:
+def make_log_file_name(name: str, map_variable: TypeMapVariable) -> str:
     random_tag = "".join(random.choices(string.ascii_uppercase + string.digits, k=3))
-    log_file_name = node.name
+    log_file_name = name
 
     if map_variable:
         for _, value in map_variable.items():