argo-kedro 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

argo_kedro/framework/cli/cli.py

@@ -0,0 +1,242 @@
+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/pipeline/__init__.py

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

argo_kedro/pipeline/fused_pipeline.py

@@ -0,0 +1,58 @@
+from typing import Iterable, List
+from kedro.pipeline import Pipeline, Node
+from functools import cached_property
+
+class FusedNode(Node):
+    """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):
+        self._nodes = nodes
+        self._name = name
+        self._namespace = None
+        self._inputs = []
+        self._outputs = []
+        self._confirms = []
+        self._func = lambda: None
+        self._tags = []
+
+        for node in nodes:
+            self._inputs.extend(node.inputs)
+            self._outputs.extend(node.outputs)
+            self._tags.extend(node._tags)
+
+        # NOTE: Exclude ouputs made as part of the intermediate nodes
+        for node in self._outputs:
+            if node in self._inputs:
+                self._inputs.remove(node)
+
+        self._tags = list(set(self._tags))
+
+    @cached_property
+    def inputs(self) -> list[str]:
+        return self._inputs # TODO: Remove transcoding?
+
+
+class FusedPipeline(Pipeline):
+    """Fused pipeline allows for wrapping nodes for execution by the underlying
+    pipeline execution framework.
+
+    This is needed, as Kedro immediately translates a pipeline to a list of nodes
+    to execute, where any pipeline structure is flatmapped. The FusedPipeline produces
+    a _single_ FusedNode that contains the wrapped nodes."""
+
+    def __init__(
+        self,
+        nodes: Iterable[Node | Pipeline],
+        name: str,
+        *,
+        tags: str | Iterable[str] | None = None,
+    ):
+        self._name = name
+        super().__init__(nodes, tags=tags)
+
+    @property
+    def nodes(self) -> list[Node]:
+        return [FusedNode(self._nodes, name=self._name)]

argo_kedro/runners/__init__.py

@@ -0,0 +1,3 @@
+from .fuse_runner import FusedRunner
+
+__all__ = ["FusedRunner"]

argo_kedro/runners/fuse_runner.py

@@ -0,0 +1,92 @@
+from kedro.io import DataCatalog
+from kedro.framework.project import pipelines
+from kedro.pipeline import Pipeline
+from kedro.io.memory_dataset import MemoryDataset
+from kedro.runner.sequential_runner import SequentialRunner
+from pluggy import PluginManager
+
+from argo_kedro.pipeline.fused_pipeline import FusedNode
+
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, List
+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.project import pipelines as kedro_pipelines
+from kedro.framework.cli.utils import find_run_command
+from kedro.pipeline import Pipeline
+from kedro.pipeline.node import Node
+
+LOGGER = getLogger(__name__)
+ARGO_TEMPLATES_DIR_PATH = Path(__file__).parent.parent.parent / "templates"
+
+
+class FusedRunner(SequentialRunner):
+    """Fused runner is an extension of the SequentialRunner that
+    essentially unpacks the FusedNode back to the contained nodes for
+    execution."""
+
+    def __init__(
+        self,
+        is_async: bool = False,
+        pipeline_name: str | None = None,
+    ):
+        """Instantiates the runner class.
+
+        The runner requires access to the pipeline name under execution to correctly handle
+        node fusing, as each node during parallell execution is wrapped as a single unit. To
+        properly fuse, the runner needs to know the pipeline execution boundary.
+
+        Args:
+            is_async: If True, the node inputs and outputs are loaded and saved
+                asynchronously with threads. Defaults to False.
+            pipeline_name: Name of the pipeline to run.
+        """
+        self._is_async = is_async
+        self._pipeline_name = pipeline_name
+
+    def _run(
+        self,
+        pipeline: Pipeline,
+        catalog: DataCatalog,
+        hook_manager: PluginManager,
+        session_id: str | None = None,
+    ) -> None:
+        nodes = pipeline.nodes
+
+        LOGGER.warning(f"Running pipeline: {self._pipeline_name}")
+
+        for node in nodes:
+            if isinstance(node, FusedNode):
+                pipeline = Pipeline(node._nodes)
+
+                outputs = pipeline.outputs()
+                for dataset in pipeline.datasets():
+
+                    found = False
+                    for pipeline_node in pipelines[self._pipeline_name].nodes:
+                        if node.name != pipeline_node.name:
+                            if dataset in pipeline_node.inputs:
+                                found = True
+                                break
+
+                    if found:
+                        print(f"{dataset} found as input to other pipeline node")
+                        outputs.append(dataset)
+
+                for dataset in pipeline.datasets().difference(pipeline.inputs().union(outputs)):
+                    catalog._datasets[dataset] = MemoryDataset()
+
+        # Invoke super runner
+        super()._run(
+            Pipeline([Pipeline(node._nodes) if isinstance(node, FusedNode) else node for node in nodes]),
+            catalog,
+            hook_manager,
+            session_id,
+        )

argo_kedro/templates/argo_wf_spec.tmpl

@@ -0,0 +1,53 @@
+{# <project_root>/templates/argo_spec.tmpl #}
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  generateName: workflow-
+  namespace: {{ namespace }}
+spec:
+  workflowMetadata:
+    labels:
+      plugin: argo-kedro 
+  entrypoint: "pipeline"
+  templates:
+  - name: kedro
+    metadata:
+      labels:
+        app: argo-kedro
+    inputs:
+      parameters:
+      - name: pipeline
+      - name: kedro_nodes
+    container:
+      image: {{ image }}
+      command: ["kedro"]
+      imagePullPolicy: Always
+      args:
+      - "run"
+      - "--pipeline"
+      - "{{ '{{inputs.parameters.pipeline}}' }}"
+      - "--nodes"
+      - "{{ '{{inputs.parameters.kedro_nodes}}' }}"
+      - "--env"
+      - "{{ environment }}"
+
+  - name: pipeline
+    dag:
+      tasks:
+      {% for task in pipeline_tasks %}
+      - name: {{ task.name }}
+        template: {{ task.get('template', 'kedro') }}
+        {% if task.deps %}
+        dependencies:
+        {% for dep in task.deps %}
+          - {{ dep }}
+        {% endfor %}
+        {% endif %}
+        arguments:
+          parameters:
+          - name: pipeline
+            value: {{ pipeline_name }}
+          - name: kedro_nodes
+            value: {{ task.nodes }}
+
+      {% endfor %}

{argo_kedro-0.1.2.dist-info → argo_kedro-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: argo-kedro
-Version: 0.1.2
+Version: 0.1.3
 Summary: Kedro plugin for running pipelines on Argo Workflows
 Author-email: Laurens Vijnck <laurens@everycure.org>, Nelson Alfonso <nelson@everycure.org>
 License: MIT

argo_kedro-0.1.3.dist-info/RECORD

@@ -0,0 +1,14 @@
+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.3.dist-info/licenses/LICENSE,sha256=deLLtAUKpK9aD3f3sgr6-FP6M1K5Mh9ai0-EBBUZMqA,1080
+argo_kedro-0.1.3.dist-info/METADATA,sha256=y4UcEXsyETGe8MYdboIyiOd2tkMU2hhAj5S9vKMh3Po,3272
+argo_kedro-0.1.3.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+argo_kedro-0.1.3.dist-info/entry_points.txt,sha256=czlBy9HPiG00bIpokpPOaE2EFr9YlWRSDhWPbLDcEYU,134
+argo_kedro-0.1.3.dist-info/top_level.txt,sha256=bkDBnht8zOdNxOcy4MwQU2MoRz5-eOww8MVzW2CLEdE,11
+argo_kedro-0.1.3.dist-info/RECORD,,

{argo_kedro-0.1.2.dist-info → argo_kedro-0.1.3.dist-info}/WHEEL

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

argo_kedro-0.1.3.dist-info/entry_points.txt

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

argo_kedro-0.1.3.dist-info/top_level.txt

@@ -0,0 +1 @@
+argo_kedro

argo_kedro-0.1.2.dist-info/RECORD

@@ -1,6 +0,0 @@
-argo_kedro-0.1.2.dist-info/licenses/LICENSE,sha256=deLLtAUKpK9aD3f3sgr6-FP6M1K5Mh9ai0-EBBUZMqA,1080
-argo_kedro-0.1.2.dist-info/METADATA,sha256=AxKrCOSc0__esa0Fp7xiVWDih1QAzslbBEXJeHtn8aI,3272
-argo_kedro-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-argo_kedro-0.1.2.dist-info/entry_points.txt,sha256=_v1LY-D5TaFe5teVUEGpSGSYElj2FoV592cxE_ewwv4,134
-argo_kedro-0.1.2.dist-info/top_level.txt,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-argo_kedro-0.1.2.dist-info/RECORD,,

argo_kedro-0.1.2.dist-info/entry_points.txt

@@ -1,5 +0,0 @@
-[kedro.global_commands]
-run = kedro_argo.framework.cli.cli:cli
-
-[kedro.project_commands]
-argo = kedro_argo.framework.cli.cli:commands

argo_kedro-0.1.2.dist-info/top_level.txt

@@ -1 +0,0 @@
-