argo-kedro 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

argo_kedro/framework/cli/cli.py

@@ -1,6 +1,6 @@
 import re
 from pathlib import Path
-from typing import Any, Dict, List, Iterable
+from typing import Any, Dict, List, Iterable, Union
 from logging import getLogger
 
 import click
@@ -9,7 +9,10 @@ from kubernetes import config
 from kubernetes.dynamic import DynamicClient
 from jinja2 import Environment, FileSystemLoader
 from kedro.framework.cli.utils import CONTEXT_SETTINGS, KedroCliError
+from kedro.framework.project import pipelines, settings
 from kedro.framework.session import KedroSession
+from kedro.framework.startup import bootstrap_project
+from kedro.utils import find_kedro_project, is_kedro_project
 from kedro.framework.cli.project import (
     ASYNC_ARG_HELP,
     CONF_SOURCE_HELP,
@@ -30,11 +33,48 @@ from kedro.pipeline import Pipeline
 from kedro.pipeline.node import Node
 from kedro.runner.sequential_runner import SequentialRunner
 from argo_kedro.runners.fuse_runner import FusedRunner
+from argo_kedro.framework.hooks.argo_hook import MachineType
+from argo_kedro.pipeline.node import ArgoNode
 
 LOGGER = getLogger(__name__)
 ARGO_TEMPLATES_DIR_PATH = Path(__file__).parent.parent.parent / "templates"
 
 
+def render_jinja_template(
+    src: Union[str, Path], **kwargs
+) -> str:
+    """This functions enable to copy a file and render the
+    tags (identified by {{ my_tag }}) with the values provided in kwargs.
+
+    Arguments:
+        src {Union[str, Path]} -- The path to the template which should be rendered
+
+    Returns:
+        str -- A string that contains all the files with replaced tags.
+    """
+    src = Path(src)
+    template_loader = FileSystemLoader(searchpath=src.parent.as_posix())
+    template_env = Environment(loader=template_loader, keep_trailing_newline=True)
+    template = template_env.get_template(src.name)
+    return template.render(**kwargs)
+
+
+def write_jinja_template(
+    src: Union[str, Path], dst: Union[str, Path], **kwargs
+) -> None:
+    """Write a template file and replace tis jinja's tags
+     (identified by {{ my_tag }}) with the values provided in kwargs.
+
+    Arguments:
+        src {Union[str, Path]} -- Path to the template which should be rendered
+        dst {Union[str, Path]} -- Path where the rendered template should be saved
+    """
+    dst = Path(dst)
+    parsed_template = render_jinja_template(src, **kwargs)
+    with open(dst, "w") as file_handler:
+        file_handler.write(parsed_template)
+
+
 @click.group(context_settings=CONTEXT_SETTINGS)
 def cli():
     pass
@@ -101,21 +141,108 @@ def _run_command_impl(
             namespaces=namespaces,
         )
 
+class KedroClickGroup(click.Group):
+    def reset_commands(self):
+        self.commands = {}
+
+        # add commands on the fly based on conditions
+        if is_kedro_project(find_kedro_project(Path.cwd())):
+            self.add_command(init)
+            self.add_command(submit)
+
+    def list_commands(self, ctx):
+        self.reset_commands()
+        commands_list = sorted(self.commands)
+        return commands_list
+
+    def get_command(self, ctx, cmd_name):
+        self.reset_commands()
+        return self.commands.get(cmd_name)
+
 @click.group(name="argo")
 def commands():
     pass
 
-@commands.command(name="submit")
+@commands.command(name="argo", cls=KedroClickGroup)
+def argo_commands():
+    """Use mlflow-specific commands inside kedro project."""
+    pass  # pragma: no cover
+
+@argo_commands.command()
+@click.option(
+    "--env",
+    "-e",
+    default="base",
+    help="The name of the kedro environment where the 'argo.yml' should be created. Default to 'base'",
+)
+@click.option(
+    "--force",
+    "-f",
+    is_flag=True,
+    default=False,
+    help="Update the template without any checks.",
+)
+@click.option(
+    "--silent",
+    "-s",
+    is_flag=True,
+    default=False,
+    help="Should message be logged when files are modified?",
+)
+def init(env: str, force: bool, silent: bool):
+    """Updates the template of a kedro project.
+    Running this command is mandatory to use argo-kedro.
+    This adds "conf/base/argo.yml": This is a configuration file
+    used for run parametrization when calling "kedro run" command.
+    """
+
+    # get constants
+    argo_yml = "argo.yml"
+    project_path = find_kedro_project(Path.cwd()) or Path.cwd()
+    project_metadata = bootstrap_project(project_path)
+    argo_yml_path = project_path / settings.CONF_SOURCE / env / argo_yml
+
+    # mlflow.yml is just a static file,
+    # but the name of the experiment is set to be the same as the project
+    if argo_yml_path.is_file() and not force:
+        click.secho(
+            click.style(
+                f"A 'argo.yml' already exists at '{argo_yml_path}' You can use the ``--force`` option to override it.",
+                fg="red",
+            )
+        )
+    else:
+        try:
+            write_jinja_template(
+                src=ARGO_TEMPLATES_DIR_PATH / argo_yml,
+                is_cookiecutter=False,
+                dst=argo_yml_path,
+                python_package=project_metadata.package_name,
+            )
+            if not silent:
+                click.secho(
+                    click.style(
+                        f"'{settings.CONF_SOURCE}/{env}/{argo_yml}' successfully updated.",
+                        fg="green",
+                    )
+                )
+        except FileNotFoundError:
+            click.secho(
+                click.style(
+                    f"No env '{env}' found. Please check this folder exists inside '{settings.CONF_SOURCE}' folder.",
+                    fg="red",
+                )
+            )
+
+@argo_commands.command(name="submit")
 @click.option("--pipeline", "-p", type=str, default="__default__", help="Specify which pipeline to execute")
 @click.option("--environment", "-e", type=str, default="base", help="Kedro environment to execute in")
 @click.option("--image", type=str, required=True, help="Image to execute")
-@click.option("--namespace", "-n", type=str, required=True, help="Namespace to execute in")
 @click.pass_obj
 def submit(
     ctx,
     pipeline: str,
     image: str,
-    namespace: str,
     environment: str
 ):
     """Submit the pipeline to Argo."""
@@ -125,43 +252,61 @@ def submit(
     template_env = Environment(loader=loader, trim_blocks=True, lstrip_blocks=True)
     template = template_env.get_template("argo_wf_spec.tmpl")
 
-    pipeline_tasks = get_argo_dag(kedro_pipelines[pipeline])
-
     LOGGER.info("Rendering Argo spec...")
 
-    # Render the template
-    rendered_template = template.render(
-        pipeline_tasks=[task.to_dict() for task in pipeline_tasks.values()],
-        pipeline_name=pipeline,
-        image=image,
-        namespace=namespace,
-        environment=environment
-    )
-
-    # Load as yaml
-    yaml_data = yaml.safe_load(rendered_template)
-    yaml_without_anchors = yaml.dump(yaml_data, sort_keys=False, default_flow_style=False)
-    save_argo_template(
-        yaml_without_anchors,
-    )
-
-    # Use kubeconfig to submit to kubernetes
-    config.load_kube_config()
-    client = DynamicClient(config.new_client_from_config())
-
-    resource = client.resources.get(
-        api_version=yaml_data["apiVersion"],
-        kind=yaml_data["kind"],
-    )
-
-    resource.create(
-        body=yaml_data,
-        namespace=namespace
-    )
+    project_path = find_kedro_project(Path.cwd()) or Path.cwd()
+    bootstrap_project(project_path)
+    with KedroSession.create(
+        project_path=project_path,
+        env=environment,
+    ) as session:
+        context = session.load_context()
+        pipeline_tasks = get_argo_dag(
+            kedro_pipelines[pipeline], 
+            machine_types=context.argo.machine_types,
+            default_machine_type=context.argo.default_machine_type
+        )
+
+        # Render the template
+        rendered_template = template.render(
+            pipeline_tasks=[task.to_dict() for task in pipeline_tasks.values()],
+            pipeline_name=pipeline,
+            image=image,
+            namespace=context.argo.namespace,
+            environment=environment
+        )
+
+        # Load as yaml
+        yaml_data = yaml.safe_load(rendered_template)
+        yaml_without_anchors = yaml.dump(yaml_data, sort_keys=False, default_flow_style=False)
+        save_argo_template(
+            yaml_without_anchors,
+        )
+
+        # Use kubeconfig to submit to kubernetes
+        config.load_kube_config()
+        client = DynamicClient(config.new_client_from_config())
+
+        resource = client.resources.get(
+            api_version=yaml_data["apiVersion"],
+            kind=yaml_data["kind"],
+        )
+
+        response = resource.create(
+            body=yaml_data,
+            namespace=context.argo.namespace
+        )
+        
+        workflow_name = response.metadata.name
+        LOGGER.info(f"Workflow submitted successfully: {workflow_name}")
+        LOGGER.info(f"View workflow at: https://argo.ai-platform.dev.everycure.org/workflows/{context.argo.namespace}/{workflow_name}")
+        
+        return workflow_name
 
 
 def save_argo_template(argo_template: str) -> str:
     file_path = Path("templates") / "argo-workflow-template.yml"
+    file_path.parent.mkdir(parents=True, exist_ok=True)
     with open(file_path, "w") as f:
         f.write(argo_template)
     return str(file_path)
@@ -173,9 +318,10 @@ class ArgoTask:
     Argo's operating model slightly differs from Kedro's, i.e., while Kedro uses dataset
     dependecies to model relationships, Argo uses task dependencies."""
 
-    def __init__(self, node: Node):
+    def __init__(self, node: Node, machine_type: MachineType):
         self._node = node
         self._parents = []
+        self._machine_type = machine_type
 
     @property
     def node(self):
@@ -189,10 +335,17 @@ class ArgoTask:
             "name": clean_name(self._node.name),
             "nodes": self._node.name,
             "deps": [clean_name(parent.name) for parent in sorted(self._parents)],
+            "mem": self._machine_type.mem,
+            "cpu": self._machine_type.cpu,
+            "num_gpu": self._machine_type.num_gpu,
         }
 
 
-def get_argo_dag(pipeline: Pipeline) -> List[Dict[str, Any]]:
+def get_argo_dag(
+    pipeline: Pipeline, 
+    machine_types: dict[str, MachineType],
+    default_machine_type: str,
+) -> List[Dict[str, Any]]:
     """Function to convert the Kedro pipeline into Argo Tasks. The function
     iterates the nodes of the pipeline and generates Argo tasks with dependencies.
     These dependencies are inferred based on the input and output datasets for
@@ -208,7 +361,12 @@ def get_argo_dag(pipeline: Pipeline) -> List[Dict[str, Any]]:
     # allowing us to easily translate the Kedro DAG to an Argo WF.
     for group in pipeline.grouped_nodes:
         for target_node in group:
-            task = ArgoTask(target_node)
+            try:
+                task = ArgoTask(target_node, machine_types[target_node.machine_type] if isinstance(target_node, ArgoNode) and target_node.machine_type is not None else machine_types[default_machine_type])
+            except KeyError as e:
+                LOGGER.error(f"Machine type not found for node `{target_node.name}`")
+                raise KeyError(f"Machine type `{target_node.machine_type}` not found for node `{target_node.name}`")
+            
             task.add_parents(
                 [
                     parent.node

argo_kedro/framework/hooks/argo_hook.py

@@ -0,0 +1,66 @@
+
+import os
+import re
+from logging import Logger, getLogger
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Union
+
+from kedro.config import MissingConfigException
+from kedro.framework.context import KedroContext
+from kedro.framework.hooks import hook_impl
+from kedro.framework.startup import _get_project_metadata
+from kedro.io import CatalogProtocol, DataCatalog
+from kedro.pipeline import Pipeline
+from kedro.pipeline.node import Node
+from omegaconf import OmegaConf
+
+
+from pydantic import BaseModel
+
+class MachineType(BaseModel):
+    mem: int
+    cpu: int
+    num_gpu: int
+
+class ArgoConfig(BaseModel):
+    namespace: str
+    machine_types: dict[str, MachineType]
+    default_machine_type: str
+
+
+class ArgoHook:
+    @property
+    def _logger(self) -> Logger:
+        return getLogger(__name__)
+
+    @hook_impl
+    def after_context_created(
+        self,
+        context: KedroContext,
+    ) -> None:
+        """Hooks to be invoked after a `KedroContext` is created. This is the earliest
+        hook triggered within a Kedro run. The `KedroContext` stores useful information
+        such as `credentials`, `config_loader` and `env`.
+        Args:
+            context: The context that was created.
+        """
+        try:
+            if "argo" not in context.config_loader.config_patterns.keys():
+                context.config_loader.config_patterns.update(
+                    {"argo": ["argo*", "argo*/**", "**/argo*"]}
+                )
+            conf_argo_yml = context.config_loader["argo"]
+        except MissingConfigException:
+            self._logger.warning(
+                "No 'argo.yml' config file found in environment. Default configuration will be used. Use ``kedro argo init`` command in CLI to customize the configuration."
+            )
+            # we create an empty dict to have the same behaviour when the argo.yml
+            # is commented out. In this situation there is no MissingConfigException
+            # but we got an empty dict
+            conf_argo_yml = {}
+
+        conf_argo_yml = ArgoConfig.model_validate(conf_argo_yml)
+        context.__setattr__("argo", conf_argo_yml)
+
+argo_hook = ArgoHook()

argo_kedro/pipeline/__init__.py

@@ -1,3 +1,4 @@
 from .fused_pipeline import FusedPipeline
+from .node import ArgoNode
 
-__all__ = ["FusedPipeline", ]
+__all__ = ["FusedPipeline", "ArgoNode"]

argo_kedro/pipeline/fused_pipeline.py

@@ -1,14 +1,16 @@
 from typing import Iterable, List
-from kedro.pipeline import Pipeline, Node
+from kedro.pipeline import Pipeline
 from functools import cached_property
+from argo_kedro.pipeline.node import ArgoNode
+from kedro.pipeline.node import Node
 
-class FusedNode(Node):
+class FusedNode(ArgoNode):
     """FusedNode is an extension of Kedro's internal node. The FusedNode
     wraps a set of nodes, and correctly sets it's `inputs` and `outputs`
     allowing it to act as a single unit for execution.
     """
 
-    def __init__(self, nodes: List[Node], name: str):
+    def __init__(self, nodes: List[Node], name: str, machine_type: str | None = None):
         self._nodes = nodes
         self._name = name
         self._namespace = None
@@ -17,6 +19,7 @@ class FusedNode(Node):
         self._confirms = []
         self._func = lambda: None
         self._tags = []
+        self._machine_type = machine_type
 
         for node in nodes:
             self._inputs.extend(node.inputs)
@@ -49,10 +52,12 @@ class FusedPipeline(Pipeline):
         name: str,
         *,
         tags: str | Iterable[str] | None = None,
+        machine_type: str | None = None,
     ):
         self._name = name
+        self._machine_type = machine_type
         super().__init__(nodes, tags=tags)
 
     @property
     def nodes(self) -> list[Node]:
-        return [FusedNode(self._nodes, name=self._name)]
+        return [FusedNode(self._nodes, name=self._name, machine_type=self._machine_type)]

argo_kedro/pipeline/node.py

@@ -0,0 +1,26 @@
+from kedro.pipeline import Node
+from typing import Callable, Iterable
+
+class ArgoNode(Node):
+    """ArgoNode is an extension of the Kedro node class, aimed at allowing
+    the node to be allocated to a specific machine type.
+    """
+    def __init__(
+        self,
+        func: Callable,
+        inputs: str | list[str] | dict[str, str] | None,
+        outputs: str | list[str] | dict[str, str] | None,
+        *,
+        name: str | None = None,
+        machine_type: str | None = None,
+        tags: str | Iterable[str] | None = None,
+        confirms: str | list[str] | None = None,
+        namespace: str | None = None,
+    ):
+
+        super().__init__(func, inputs, outputs, name=name, tags=tags, confirms=confirms, namespace=namespace)
+        self._machine_type = machine_type
+
+    @property
+    def machine_type(self) -> str:
+        return self._machine_type

argo_kedro/templates/argo.yml

@@ -0,0 +1,9 @@
+namespace: argo-workflows
+
+machine_types:
+  default:
+    mem: 16
+    cpu: 4
+    num_gpu: 0
+
+default_machine_type: default

argo_kedro/templates/argo_wf_spec.tmpl

@@ -18,6 +18,25 @@ spec:
       parameters:
       - name: pipeline
       - name: kedro_nodes
+    podSpecPatch: |
+      containers:
+        - name: main
+        # Add tolerations for large memory nodes and GPU nodes
+          resources:
+            requests:
+              memory: {% raw %} "{{inputs.parameters.mem}}Gi"
+              {% endraw %}
+              cpu: {% raw %} "{{inputs.parameters.cpu}}"
+              {% endraw %}
+              nvidia.com/gpu: {% raw %} "{{inputs.parameters.num_gpu}}"
+              {% endraw %}
+            limits:
+              memory: {% raw %} "{{inputs.parameters.mem}}Gi"
+              {% endraw %}
+              cpu: {% raw %} "{{inputs.parameters.cpu}}"
+              {% endraw %}
+              nvidia.com/gpu: {% raw %} "{{inputs.parameters.num_gpu}}"
+              {% endraw %}
     container:
       image: {{ image }}
       command: ["kedro"]
@@ -49,5 +68,10 @@ spec:
             value: {{ pipeline_name }}
           - name: kedro_nodes
             value: {{ task.nodes }}
-
+          - name: num_gpu
+            value: {{ task.num_gpu }}
+          - name: mem
+            value: {{ task.mem }}
+          - name: cpu
+            value: {{ task.cpu }}
       {% endfor %}

{argo_kedro-0.1.5.dist-info → argo_kedro-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: argo-kedro
-Version: 0.1.5
+Version: 0.1.7
 Summary: Kedro plugin for running pipelines on Argo Workflows
 Author-email: Laurens Vijnck <laurens@everycure.org>, Nelson Alfonso <nelson@everycure.org>
 License: MIT
@@ -26,6 +26,7 @@ Requires-Dist: kedro
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: jinja2>=3.0.0
 Requires-Dist: kubernetes>=35.0.0
+Requires-Dist: pydantic>=2.0.0
 Dynamic: license-file
 
 # argo-kedro

argo_kedro-0.1.7.dist-info/RECORD

@@ -0,0 +1,17 @@
+argo_kedro/framework/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+argo_kedro/framework/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+argo_kedro/framework/cli/cli.py,sha256=3w-bdPbFtLobI3dtlHBKZECnG90nKwinDWc2vY84zxM,14071
+argo_kedro/framework/hooks/argo_hook.py,sha256=aBbq_u8CskwJDIZaV2Eb7DvR-yLR_1ywKl38PxPdtO0,2183
+argo_kedro/pipeline/__init__.py,sha256=OcViNwRwYyMCIYPPRHTGIiOH3UfkZzgXXn-g5vgRzgE,109
+argo_kedro/pipeline/fused_pipeline.py,sha256=Z28PI0ACincyKjqIuaBw4FaBwQx6r1MUc7_kDopYN2c,2116
+argo_kedro/pipeline/node.py,sha256=zhV_KNZIAg3xpRxoZCCbayK9tpD5crXllGE_9S-ovn8,868
+argo_kedro/runners/__init__.py,sha256=AfU9FbRebpfTYnliocXtdwALpfVlh19WXcrEBh4Wb78,63
+argo_kedro/runners/fuse_runner.py,sha256=K-OmciE8hLicMNiaLe5SjR4GwBH4Ud3mtNDknvzhTFA,3177
+argo_kedro/templates/argo.yml,sha256=YJ1tyEw7X8npRj0PPu-C8QGxCTsPHQMXvQnqvLdzKjI,121
+argo_kedro/templates/argo_wf_spec.tmpl,sha256=-2diQNjd4y3LqcQJ5br_fEU2ZG4aIT8M5CiEN1TseVw,2122
+argo_kedro-0.1.7.dist-info/licenses/LICENSE,sha256=deLLtAUKpK9aD3f3sgr6-FP6M1K5Mh9ai0-EBBUZMqA,1080
+argo_kedro-0.1.7.dist-info/METADATA,sha256=aV0fUpLwBxDV-yMwFdUPx7wp7jR6kzlqWdGXwV0j_Iw,3337
+argo_kedro-0.1.7.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+argo_kedro-0.1.7.dist-info/entry_points.txt,sha256=YAR67HB3IjHmg3cSXSx9WFpRct41kS2M9Y9JPbAZgSg,208
+argo_kedro-0.1.7.dist-info/top_level.txt,sha256=bkDBnht8zOdNxOcy4MwQU2MoRz5-eOww8MVzW2CLEdE,11
+argo_kedro-0.1.7.dist-info/RECORD,,

{argo_kedro-0.1.5.dist-info → argo_kedro-0.1.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{argo_kedro-0.1.5.dist-info → argo_kedro-0.1.7.dist-info}/entry_points.txt

@@ -1,5 +1,8 @@
 [kedro.global_commands]
 run = argo_kedro.framework.cli.cli:cli
 
+[kedro.hooks]
+argo_hook = argo_kedro.framework.hooks.argo_hook:argo_hook
+
 [kedro.project_commands]
 argo = argo_kedro.framework.cli.cli:commands

argo_kedro-0.1.5.dist-info/RECORD

@@ -1,14 +0,0 @@
-argo_kedro/framework/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-argo_kedro/framework/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-argo_kedro/framework/cli/cli.py,sha256=x1wZfDJ3-GyxhUDnmdSzcmpVLrQu1SowZH9Wnjl0mz0,8176
-argo_kedro/pipeline/__init__.py,sha256=eADPWCo5qxWr3nWXJJK7yugfJ37-zGYdY7frE-8dLcs,72
-argo_kedro/pipeline/fused_pipeline.py,sha256=FHK5ZMQi21dHYGGewU-4QN31JwtiWbhtPiGrnC9QtpE,1844
-argo_kedro/runners/__init__.py,sha256=AfU9FbRebpfTYnliocXtdwALpfVlh19WXcrEBh4Wb78,63
-argo_kedro/runners/fuse_runner.py,sha256=K-OmciE8hLicMNiaLe5SjR4GwBH4Ud3mtNDknvzhTFA,3177
-argo_kedro/templates/argo_wf_spec.tmpl,sha256=7gLBGnA_d8bwhwOAMGgU_IgWIjpD8atDEGidrG5BNKU,1218
-argo_kedro-0.1.5.dist-info/licenses/LICENSE,sha256=deLLtAUKpK9aD3f3sgr6-FP6M1K5Mh9ai0-EBBUZMqA,1080
-argo_kedro-0.1.5.dist-info/METADATA,sha256=VwFG8P4SJEwE0VVz81RfTvVyS88DUr_FF7RT-VsoeKg,3306
-argo_kedro-0.1.5.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-argo_kedro-0.1.5.dist-info/entry_points.txt,sha256=czlBy9HPiG00bIpokpPOaE2EFr9YlWRSDhWPbLDcEYU,134
-argo_kedro-0.1.5.dist-info/top_level.txt,sha256=bkDBnht8zOdNxOcy4MwQU2MoRz5-eOww8MVzW2CLEdE,11
-argo_kedro-0.1.5.dist-info/RECORD,,