argo-kedro 0.1.5__tar.gz → 0.1.7__tar.gz

{argo_kedro-0.1.5 → argo_kedro-0.1.7}/PKG-INFO

@@ -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/argo_kedro/framework/cli/cli.py

@@ -0,0 +1,400 @@
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Iterable, Union
+from logging import getLogger
+
+import click
+import yaml
+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,
+    FROM_INPUTS_HELP,
+    FROM_NODES_HELP,
+    LOAD_VERSION_HELP,
+    NODE_ARG_HELP,
+    PARAMS_ARG_HELP,
+    PIPELINE_ARG_HELP,
+    RUNNER_ARG_HELP,
+    TAG_ARG_HELP,
+    TO_NODES_HELP,
+    TO_OUTPUTS_HELP,
+    project_group,
+)
+from kedro.framework.project import pipelines as kedro_pipelines
+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
+
+@cli.command(name="run")
+@click.option("--pipeline", "-p", type=str, default="__default__", help="Name of the pipeline to execute")
+@click.option("--env", "-e", type=str, default=None, help="Kedro environment to run the pipeline in")
+@click.option("--config", "-c", type=str, multiple=True, help="Extra config to pass to KedroContext")
+@click.option("--params", type=str, multiple=True, help="Override parameters")
+@click.option("--tags", "-t", type=str, multiple=True, help=TAG_ARG_HELP)
+@click.option("--nodes", "-n", type=str, multiple=True, help="Run only nodes with specified names")
+@click.option("--to-nodes", type=str, multiple=True, help="Run a sub-pipeline up to certain nodes")
+@click.option("--from-nodes", type=str, multiple=True, help="Run a sub-pipeline starting from certain nodes")
+@click.option("--from-inputs", type=str, multiple=True, help="Run a sub-pipeline starting from nodes that produce these inputs")
+@click.option("--to-outputs", type=str, multiple=True, help="Run a sub-pipeline up to nodes that produce these outputs")
+@click.option("--load-version", type=str, multiple=True, help="Specify a particular dataset version")
+@click.option("--namespaces", type=str, multiple=True, help="Namespaces of the pipeline")
+@click.pass_obj
+def _run_command_impl(
+    ctx,
+    pipeline: str,
+    env: str,
+    config: tuple,
+    params: tuple,
+    tags: list[str],
+    nodes: tuple,
+    to_nodes: tuple,
+    from_nodes: tuple,
+    from_inputs: tuple,
+    to_outputs: tuple,
+    load_version: tuple,
+    namespaces: Iterable[str],
+):    
+    """Run the pipeline with the FusedRunner."""
+
+    LOGGER.warning(f"Using plugin entrypoint")
+    
+    load_versions = None
+    if load_version:
+        load_versions = {}
+        for version_spec in load_version:
+            if ":" in version_spec:
+                dataset, version = version_spec.split(":", 1)
+                load_versions[dataset] = version
+
+    conf_source = getattr(ctx, "conf_source", None)
+    env_value = env or getattr(ctx, "env", None)
+
+    with KedroSession.create(
+        env=env_value,
+        conf_source=conf_source,
+    ) as session:
+
+        session.run(
+            pipeline_name=pipeline,
+            tags=tags,
+            runner=FusedRunner(pipeline_name=pipeline),
+            node_names=list(nodes) if nodes else None,
+            from_nodes=list(from_nodes) if from_nodes else None,
+            to_nodes=list(to_nodes) if to_nodes else None,
+            from_inputs=list(from_inputs) if from_inputs else None,
+            to_outputs=list(to_outputs) if to_outputs else None,
+            load_versions=load_versions,
+            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="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.pass_obj
+def submit(
+    ctx,
+    pipeline: str,
+    image: str,
+    environment: str
+):
+    """Submit the pipeline to Argo."""
+    LOGGER.info("Loading spec template..")
+
+    loader = FileSystemLoader(searchpath=ARGO_TEMPLATES_DIR_PATH)
+    template_env = Environment(loader=loader, trim_blocks=True, lstrip_blocks=True)
+    template = template_env.get_template("argo_wf_spec.tmpl")
+
+    LOGGER.info("Rendering Argo spec...")
+
+    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)
+
+
+class ArgoTask:
+    """Class to model an Argo task.
+
+    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, machine_type: MachineType):
+        self._node = node
+        self._parents = []
+        self._machine_type = machine_type
+
+    @property
+    def node(self):
+        return self._node
+
+    def add_parents(self, nodes: List[Node]):
+        self._parents.extend(nodes)
+
+    def to_dict(self):
+        return {
+            "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, 
+    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
+    each node.
+
+    NOTE: This function is now agnostic to the fact that nodes might be fused. The nodes
+    returned as part of the pipeline may optionally contain FusedNodes, which have correct
+    inputs and outputs for the perspective of the Argo Task.
+    """
+    tasks = {}
+
+    # The `grouped_nodes` property returns the nodes list, in a toplogical order,
+    # allowing us to easily translate the Kedro DAG to an Argo WF.
+    for group in pipeline.grouped_nodes:
+        for target_node in group:
+            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
+                    for parent in tasks.values()
+                    if set(clean_dependencies(target_node.inputs)) & set(clean_dependencies(parent.node.outputs))
+                ]
+            )
+
+            tasks[target_node.name] = task
+
+    return tasks
+
+
+def clean_name(name: str) -> str:
+    """Function to clean the node name.
+
+    Args:
+        name: name of the node
+    Returns:
+        Clean node name, according to Argo's requirements
+    """
+    return re.sub(r"[\W_]+", "-", name).strip("-")
+
+
+def clean_dependencies(elements) -> List[str]:
+    """Function to clean node dependencies.
+
+    Operates by removing `params:` from the list and dismissing
+    the transcoding operator.
+    """
+    return [el.split("@")[0] for el in elements if not el.startswith("params:")]

argo_kedro-0.1.7/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-0.1.7/argo_kedro/pipeline/__init__.py

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

{argo_kedro-0.1.5 → argo_kedro-0.1.7}/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-0.1.7/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-0.1.7/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-0.1.5 → argo_kedro-0.1.7}/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 → argo_kedro-0.1.7}/argo_kedro.egg-info/PKG-INFO

@@ -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.5 → argo_kedro-0.1.7}/argo_kedro.egg-info/SOURCES.txt

@@ -15,8 +15,16 @@ argo_kedro.egg-info/top_level.txt
 argo_kedro/framework/__init__.py
 argo_kedro/framework/cli/__init__.py
 argo_kedro/framework/cli/cli.py
+argo_kedro/framework/hooks/argo_hook.py
 argo_kedro/pipeline/__init__.py
 argo_kedro/pipeline/fused_pipeline.py
+argo_kedro/pipeline/node.py
 argo_kedro/runners/__init__.py
 argo_kedro/runners/fuse_runner.py
-argo_kedro/templates/argo_wf_spec.tmpl
+argo_kedro/templates/argo.yml
+argo_kedro/templates/argo_wf_spec.tmpl
+tests/__init__.py
+tests/cli/__init__.py
+tests/cli/test_cli.py
+tests/pipeline/__init__.py
+tests/pipeline/test_fused_pipeline.py

{argo_kedro-0.1.5 → argo_kedro-0.1.7}/argo_kedro.egg-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 → argo_kedro-0.1.7}/argo_kedro.egg-info/requires.txt

@@ -2,3 +2,4 @@ kedro
 pyyaml>=6.0.2
 jinja2>=3.0.0
 kubernetes>=35.0.0
+pydantic>=2.0.0

{argo_kedro-0.1.5 → argo_kedro-0.1.7}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "argo-kedro"
-version = "0.1.5"
+version = "0.1.7"
 description = "Kedro plugin for running pipelines on Argo Workflows"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -33,6 +33,7 @@ dependencies = [
     "pyyaml>=6.0.2",
     "jinja2>=3.0.0",
     "kubernetes>=35.0.0",
+    "pydantic>=2.0.0",
 ]
 
 [project.urls]
@@ -46,6 +47,9 @@ run = "argo_kedro.framework.cli.cli:cli"
 [project.entry-points."kedro.project_commands"]
 argo = "argo_kedro.framework.cli.cli:commands"
 
+[project.entry-points."kedro.hooks"]
+argo_hook = "argo_kedro.framework.hooks.argo_hook:argo_hook"
+
 [tool.setuptools]
 include-package-data = true
 
@@ -53,4 +57,4 @@ include-package-data = true
 include = ["argo_kedro*"]
 
 [tool.setuptools.package-data]
-argo_kedro = ["**/*.tmpl"]
+argo_kedro = ["**/*.tmpl", "**/*.yml"]

argo_kedro-0.1.7/tests/cli/test_cli.py

@@ -0,0 +1,206 @@
+import pytest
+
+from kedro.pipeline import Pipeline, node
+from argo_kedro.pipeline import FusedPipeline, ArgoNode
+from argo_kedro.framework.cli.cli import get_argo_dag, MachineType
+
+@pytest.fixture
+def machine_types() -> dict[str, MachineType]:
+    return {
+        "default": MachineType(mem=16, cpu=2, num_gpu=0),
+        "n1-standard-4": MachineType(mem=16, cpu=4, num_gpu=0),
+        "n1-standard-8": MachineType(mem=16, cpu=8, num_gpu=0),
+    }
+
+@pytest.fixture
+def default_machine_type() -> str:
+    return "default"
+
+@pytest.fixture
+def pipeline() -> Pipeline:
+    return Pipeline(
+        [
+            ArgoNode(
+                func=lambda x: x,
+                inputs="raw_data",
+                outputs="data",
+                tags=["preprocessing"],
+                name="preprocess_fun",
+                machine_type="n1-standard-4",
+            ),
+            ArgoNode(
+                func=lambda x: x,
+                inputs="data",
+                outputs="model",
+                tags=["training"],
+                name="train_fun",
+                machine_type="n1-standard-8",
+            ),
+        ]
+    )
+
+@pytest.fixture
+def fused_pipeline() -> Pipeline:
+    return Pipeline(
+        [
+            node(
+                func=lambda x: x,
+                inputs="raw_data",
+                outputs="data",
+                tags=["preprocessing"],
+                name="preprocess_fun",
+            ),
+            FusedPipeline(
+                [
+                    node(
+                        func=lambda x: x,
+                        inputs="data",
+                        outputs="model",
+                        tags=["training"],
+                        name="train_fun",
+                    ),
+                    node(
+                        func=lambda x: x,
+                        inputs="model",
+                        outputs="predictions",
+                        tags=["predictions"],
+                        name="create_predictions",
+                    ),
+                ],
+                name="fused_modelling",
+                machine_type="n1-standard-8",
+            ),
+        ]
+    )
+
+
+@pytest.fixture
+def fused_pipeline_complex() -> Pipeline:
+    return Pipeline(
+        [
+            node(
+                func=lambda x: x,
+                inputs="raw_data",
+                outputs="data",
+                tags=["preprocessing"],
+                name="preprocess_fun",
+            ),
+            node(
+                func=lambda x: x,
+                inputs="raw_customers",
+                outputs="customers",
+                tags=["preprocessing"],
+                name="preprocess_customers",
+            ),
+            FusedPipeline(
+                [
+                    node(
+                        func=lambda x: x,
+                        inputs="data",
+                        outputs="model",
+                        tags=["training"],
+                        name="train_fun",
+                    ),
+                    node(
+                        func=lambda x, y: x,
+                        inputs=["model", "customers"],
+                        outputs="predictions",
+                        tags=["predictions"],
+                        name="create_predictions",
+                    ),
+                ],
+                name="fused_modelling",
+                machine_type="n1-standard-8",
+            ),
+        ]
+    )
+
+
+def test_get_argo_dag(pipeline: Pipeline, machine_types: dict[str, MachineType], default_machine_type: str):
+
+    # When generating the argo DAG
+    argo_dag = get_argo_dag(pipeline, machine_types, default_machine_type)
+    expected = {
+        "preprocess_fun": { 
+            "name": "preprocess-fun",
+            "nodes": "preprocess_fun",
+            "deps": [],
+            "mem": 16,
+            "cpu": 4,
+            "num_gpu": 0,
+        },
+        "train_fun": {
+            "name": "train-fun",
+            "nodes": "train_fun",
+            "deps": ["preprocess-fun"],
+            "mem": 16,
+            "cpu": 8,
+            "num_gpu": 0,
+        }
+    }
+
+    # Assert resulting argo dag is correct
+    assert {key: task.to_dict() for key,task in argo_dag.items()} == expected
+
+
+def test_get_argo_dag_fused(fused_pipeline: Pipeline, machine_types: dict[str, MachineType], default_machine_type: str):
+
+    # When generating the argo DAG
+    argo_dag = get_argo_dag(fused_pipeline, machine_types, default_machine_type)
+    expected = {
+        "preprocess_fun": { 
+            "name": "preprocess-fun",
+            "nodes": "preprocess_fun",
+            "deps": [],
+            "mem": 16,
+            "cpu": 2,
+            "num_gpu": 0,
+        },
+        "fused_modelling": {
+            "name": "fused-modelling",
+            "nodes": "fused_modelling",
+            "deps": ["preprocess-fun"],
+            "mem": 16,
+            "cpu": 8,
+            "num_gpu": 0,
+        }
+    }
+
+    # Assert resulting argo dag is correct
+    assert {key: task.to_dict() for key,task in argo_dag.items()} == expected
+    
+
+def test_get_argo_dag_fused_complex(fused_pipeline_complex: Pipeline, machine_types: dict[str, MachineType], default_machine_type: str):
+
+    # When generating the argo DAG
+    argo_dag = get_argo_dag(fused_pipeline_complex, machine_types, default_machine_type)
+    expected = {
+        "preprocess_fun": { 
+            "name": "preprocess-fun",
+            "nodes": "preprocess_fun",
+            "deps": [],
+            "mem": 16,
+            "cpu": 2,
+            "num_gpu": 0,
+        },
+        "preprocess_customers": { 
+            "name": "preprocess-customers",
+            "nodes": "preprocess_customers",
+            "deps": [],
+            "mem": 16,
+            "cpu": 2,
+            "num_gpu": 0,
+        },
+        "fused_modelling": {
+            "name": "fused-modelling",
+            "nodes": "fused_modelling",
+            "deps": ["preprocess-customers", "preprocess-fun"],
+            "mem": 16,
+            "cpu": 8,
+            "num_gpu": 0,
+        }
+    }
+
+    # Assert resulting argo dag is correct
+    assert {key: task.to_dict() for key,task in argo_dag.items()} == expected
+    

argo_kedro-0.1.7/tests/pipeline/test_fused_pipeline.py

@@ -0,0 +1,45 @@
+import pytest
+
+from kedro.pipeline import Pipeline, node
+from argo_kedro.pipeline.fused_pipeline import FusedPipeline, FusedNode
+
+@pytest.fixture
+def pipeline() -> Pipeline:
+    return Pipeline(
+        [
+            node(
+                func=lambda x: x,
+                inputs="raw_data",
+                outputs="data",
+                tags=["preprocessing"],
+                name="preprocess_fun",
+            ),
+            node(
+                func=lambda x: x,
+                inputs="data",
+                outputs="model",
+                tags=["training"],
+                name="train_fun",
+            ),
+        ]
+    )
+
+
+def test_fused_node_inputs(pipeline: Pipeline):
+
+    # Wrap pipeline in FusedNode
+    fused_node = FusedNode(pipeline.nodes, name="fused_node")
+
+    # Assert that the fused node inputs are the pure inputs of the pipeline, i.e.,
+    # all inputs not produced as part of intermediate nodes.
+    assert set(fused_node.inputs) == set(["raw_data"])
+
+
+def test_fused_pipeline_nodes(pipeline: Pipeline):
+
+    # Wrap pipeline in FusedPipeline
+    fused_pipeline = FusedPipeline(pipeline.nodes, name="fused_pipeline")
+
+    # Assert that the fused pipeline nodes are the same as the pipeline nodes
+    assert len(fused_pipeline.nodes) == 1
+    assert isinstance(fused_pipeline.nodes[0], FusedNode)

{argo_kedro-0.1.5 → argo_kedro-0.1.7}/uv.lock

@@ -356,6 +356,26 @@ wheels = [
 ]
 
 [[package]]
+<<<<<<< HEAD:kedro-argo/uv.lock
+name = "kedro-argo"
+version = "0.1.0"
+source = { editable = "." }
+dependencies = [
+    { name = "jinja2" },
+    { name = "kedro" },
+    { name = "pyyaml" },
+]
+
+[package.metadata]
+requires-dist = [
+    { name = "jinja2", specifier = ">=3.0.0" },
+    { name = "kedro" },
+    { name = "pyyaml", specifier = ">=6.0.2" },
+]
+
+[[package]]
+=======
+>>>>>>> refs/remotes/origin/main:argo-kedro/uv.lock
 name = "kedro-telemetry"
 version = "0.6.5"
 source = { registry = "https://pypi.org/simple" }
@@ -820,6 +840,8 @@ wheels = [
 ]
 
 [[package]]
+<<<<<<< HEAD:kedro-argo/uv.lock
+=======
 name = "websocket-client"
 version = "1.9.0"
 source = { registry = "https://pypi.org/simple" }
@@ -829,6 +851,7 @@ wheels = [
 ]
 
 [[package]]
+>>>>>>> refs/remotes/origin/main:argo-kedro/uv.lock
 name = "zipp"
 version = "3.23.0"
 source = { registry = "https://pypi.org/simple" }

argo_kedro-0.1.5/argo_kedro/framework/cli/cli.py

@@ -1,242 +0,0 @@
-import re
-from pathlib import Path
-from typing import Any, Dict, List, Iterable
-from logging import getLogger
-
-import click
-import yaml
-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.session import KedroSession
-from kedro.framework.cli.project import (
-    ASYNC_ARG_HELP,
-    CONF_SOURCE_HELP,
-    FROM_INPUTS_HELP,
-    FROM_NODES_HELP,
-    LOAD_VERSION_HELP,
-    NODE_ARG_HELP,
-    PARAMS_ARG_HELP,
-    PIPELINE_ARG_HELP,
-    RUNNER_ARG_HELP,
-    TAG_ARG_HELP,
-    TO_NODES_HELP,
-    TO_OUTPUTS_HELP,
-    project_group,
-)
-from kedro.framework.project import pipelines as kedro_pipelines
-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
-
-LOGGER = getLogger(__name__)
-ARGO_TEMPLATES_DIR_PATH = Path(__file__).parent.parent.parent / "templates"
-
-
-@click.group(context_settings=CONTEXT_SETTINGS)
-def cli():
-    pass
-
-@cli.command(name="run")
-@click.option("--pipeline", "-p", type=str, default="__default__", help="Name of the pipeline to execute")
-@click.option("--env", "-e", type=str, default=None, help="Kedro environment to run the pipeline in")
-@click.option("--config", "-c", type=str, multiple=True, help="Extra config to pass to KedroContext")
-@click.option("--params", type=str, multiple=True, help="Override parameters")
-@click.option("--tags", "-t", type=str, multiple=True, help=TAG_ARG_HELP)
-@click.option("--nodes", "-n", type=str, multiple=True, help="Run only nodes with specified names")
-@click.option("--to-nodes", type=str, multiple=True, help="Run a sub-pipeline up to certain nodes")
-@click.option("--from-nodes", type=str, multiple=True, help="Run a sub-pipeline starting from certain nodes")
-@click.option("--from-inputs", type=str, multiple=True, help="Run a sub-pipeline starting from nodes that produce these inputs")
-@click.option("--to-outputs", type=str, multiple=True, help="Run a sub-pipeline up to nodes that produce these outputs")
-@click.option("--load-version", type=str, multiple=True, help="Specify a particular dataset version")
-@click.option("--namespaces", type=str, multiple=True, help="Namespaces of the pipeline")
-@click.pass_obj
-def _run_command_impl(
-    ctx,
-    pipeline: str,
-    env: str,
-    config: tuple,
-    params: tuple,
-    tags: list[str],
-    nodes: tuple,
-    to_nodes: tuple,
-    from_nodes: tuple,
-    from_inputs: tuple,
-    to_outputs: tuple,
-    load_version: tuple,
-    namespaces: Iterable[str],
-):    
-    """Run the pipeline with the FusedRunner."""
-
-    LOGGER.warning(f"Using plugin entrypoint")
-    
-    load_versions = None
-    if load_version:
-        load_versions = {}
-        for version_spec in load_version:
-            if ":" in version_spec:
-                dataset, version = version_spec.split(":", 1)
-                load_versions[dataset] = version
-
-    conf_source = getattr(ctx, "conf_source", None)
-    env_value = env or getattr(ctx, "env", None)
-
-    with KedroSession.create(
-        env=env_value,
-        conf_source=conf_source,
-    ) as session:
-
-        session.run(
-            pipeline_name=pipeline,
-            tags=tags,
-            runner=FusedRunner(pipeline_name=pipeline),
-            node_names=list(nodes) if nodes else None,
-            from_nodes=list(from_nodes) if from_nodes else None,
-            to_nodes=list(to_nodes) if to_nodes else None,
-            from_inputs=list(from_inputs) if from_inputs else None,
-            to_outputs=list(to_outputs) if to_outputs else None,
-            load_versions=load_versions,
-            namespaces=namespaces,
-        )
-
-@click.group(name="argo")
-def commands():
-    pass
-
-@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."""
-    LOGGER.info("Loading spec template..")
-
-    loader = FileSystemLoader(searchpath=ARGO_TEMPLATES_DIR_PATH)
-    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
-    )
-
-
-def save_argo_template(argo_template: str) -> str:
-    file_path = Path("templates") / "argo-workflow-template.yml"
-    with open(file_path, "w") as f:
-        f.write(argo_template)
-    return str(file_path)
-
-
-class ArgoTask:
-    """Class to model an Argo task.
-
-    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):
-        self._node = node
-        self._parents = []
-
-    @property
-    def node(self):
-        return self._node
-
-    def add_parents(self, nodes: List[Node]):
-        self._parents.extend(nodes)
-
-    def to_dict(self):
-        return {
-            "name": clean_name(self._node.name),
-            "nodes": self._node.name,
-            "deps": [clean_name(parent.name) for parent in sorted(self._parents)],
-        }
-
-
-def get_argo_dag(pipeline: Pipeline) -> 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
-    each node.
-
-    NOTE: This function is now agnostic to the fact that nodes might be fused. The nodes
-    returned as part of the pipeline may optionally contain FusedNodes, which have correct
-    inputs and outputs for the perspective of the Argo Task.
-    """
-    tasks = {}
-
-    # The `grouped_nodes` property returns the nodes list, in a toplogical order,
-    # 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)
-            task.add_parents(
-                [
-                    parent.node
-                    for parent in tasks.values()
-                    if set(clean_dependencies(target_node.inputs)) & set(clean_dependencies(parent.node.outputs))
-                ]
-            )
-
-            tasks[target_node.name] = task
-
-    return tasks
-
-
-def clean_name(name: str) -> str:
-    """Function to clean the node name.
-
-    Args:
-        name: name of the node
-    Returns:
-        Clean node name, according to Argo's requirements
-    """
-    return re.sub(r"[\W_]+", "-", name).strip("-")
-
-
-def clean_dependencies(elements) -> List[str]:
-    """Function to clean node dependencies.
-
-    Operates by removing `params:` from the list and dismissing
-    the transcoding operator.
-    """
-    return [el.split("@")[0] for el in elements if not el.startswith("params:")]

argo_kedro-0.1.5/argo_kedro/pipeline/__init__.py

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