_local_client.py

# Copyright 2021 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
import json
import logging
import os
import re
import subprocess
import tempfile
import warnings
from collections import deque
from typing import Any, Callable, Dict, List, Mapping, Optional, Union, cast

from . import dsl
from .compiler.compiler import sanitize_k8s_name


class _Dag:
    """DAG stands for Direct Acyclic Graph.

    DAG here is used to decide the order to execute pipeline ops.

    For more information on DAG, please refer to `wiki <https://en.wikipedia.org/wiki/Directed_acyclic_graph>`_.
    """

    def __init__(self, nodes: List[str]) -> None:
        """

        Args::
          nodes: List of DAG nodes, each node is identified by an unique name.
        """
        self._graph = {node: [] for node in nodes}
        self._reverse_graph = {node: [] for node in nodes}

    @property
    def graph(self):
        return self._graph

    @property
    def reverse_graph(self):
        return self._reverse_graph

    def add_edge(self, edge_source: str, edge_target: str) -> None:
        """Add an edge between DAG nodes.

        Args::
          edge_source: the source node of the edge
          edge_target: the target node of the edge
        """
        self._graph[edge_source].append(edge_target)
        self._reverse_graph[edge_target].append(edge_source)

    def get_follows(self, source_node: str) -> List[str]:
        """Get all target nodes start from the specified source node.

        Args::
          source_node: the source node
        """
        return self._graph.get(source_node, [])

    def get_dependencies(self, target_node: str) -> List[str]:
        """Get all source nodes end with the specified target node.

        Args::
          target_node: the target node
        """
        return self._reverse_graph.get(target_node, [])

    def topological_sort(self) -> List[str]:
        """List DAG nodes in topological order."""

        in_degree = {node: 0 for node in self._graph.keys()}

        for i in self._graph:
            for j in self._graph[i]:
                in_degree[j] += 1

        queue = deque()
        for node, degree in in_degree.items():
            if degree == 0:
                queue.append(node)

        sorted_nodes = []

        while queue:
            u = queue.popleft()
            sorted_nodes.append(u)

            for node in self._graph[u]:
                in_degree[node] -= 1

                if in_degree[node] == 0:
                    queue.append(node)

        return sorted_nodes


def _extract_pipeline_param(param: str) -> dsl.PipelineParam:
    """Extract PipelineParam from string."""
    matches = re.findall(r"{{pipelineparam:op=([\w\s_-]*);name=([\w\s_-]+)}}",
                         param)
    op_dependency_name = matches[0][0]
    output_file_name = matches[0][1]
    return dsl.PipelineParam(output_file_name, op_dependency_name)


def _get_op(ops: List[dsl.ContainerOp],
            op_name: str) -> Union[dsl.ContainerOp, None]:
    """Get the first op with specified op name."""
    return next(filter(lambda op: op.name == op_name, ops), None)


def _get_subgroup(groups: List[dsl.OpsGroup],
                  group_name: str) -> Union[dsl.OpsGroup, None]:
    """Get the first OpsGroup with specified group name."""
    return next(filter(lambda g: g.name == group_name, groups), None)


class LocalClient:

    class ExecutionMode:
        """Configuration to decide whether the client executes a component in
        docker or in local process."""

        DOCKER = "docker"
        LOCAL = "local"

        def __init__(
            self,
            mode: str = DOCKER,
            images_to_exclude: List[str] = [],
            ops_to_exclude: List[str] = [],
            docker_options: List[str] = [],
        ) -> None:
            """Constructor.

            Args:
                mode: Default execution mode, default 'docker'
                images_to_exclude: If the image of op is in images_to_exclude, the op is
                    executed in the mode different from default_mode.
                ops_to_exclude: If the name of op is in ops_to_exclude, the op is
                    executed in the mode different from default_mode.
                docker_options: Docker options used in docker mode,
                    e.g. docker_options=["-e", "foo=bar"].
            """
            if mode not in [self.DOCKER, self.LOCAL]:
                raise Exception(
                    "Invalid execution mode, must be docker of local")
            self._mode = mode
            self._images_to_exclude = images_to_exclude
            self._ops_to_exclude = ops_to_exclude
            self._docker_options = docker_options

        @property
        def mode(self) -> str:
            return self._mode

        @property
        def images_to_exclude(self) -> List[str]:
            return self._images_to_exclude

        @property
        def ops_to_exclude(self) -> List[str]:
            return self._ops_to_exclude

        @property
        def docker_options(self) -> List[str]:
            return self._docker_options

    def __init__(self, pipeline_root: Optional[str] = None) -> None:
        """Construct the instance of LocalClient.

        Args：
            pipeline_root: The root directory where the output artifact of component
              will be saved.
        """
        warnings.warn(
            'LocalClient is an Alpha[1] feature. It may be deprecated in the future.\n'
            '[1] https://github.com/kubeflow/pipelines/blob/master/docs/release/feature-stages.md#alpha',
            category=FutureWarning,
        )

        pipeline_root = pipeline_root or tempfile.tempdir
        self._pipeline_root = pipeline_root

    def _find_base_group(self, groups: List[dsl.OpsGroup],
                         op_name: str) -> Union[dsl.OpsGroup, None]:
        """Find the base group of op in candidate group list."""
        if groups is None or len(groups) == 0:
            return None
        for group in groups:
            if _get_op(group.ops, op_name):
                return group
            else:
                _parent_group = self._find_base_group(group.groups, op_name)
                if _parent_group:
                    return group

        return None

    def _create_group_dag(self, pipeline_dag: _Dag,
                          group: dsl.OpsGroup) -> _Dag:
        """Create DAG within current group, it's a DAG of direct ops and direct
        subgroups.

        Each node of the DAG is either an op or a subgroup. For each
        node in current group, if one of its DAG follows is also an op
        in current group, add an edge to this follow op, otherwise, if
        this follow belongs to subgroups, add an edge to its subgroup.
        If this node has dependency from subgroups, then add an edge
        from this subgroup to current node.
        """
        group_dag = _Dag([op.name for op in group.ops] +
                         [g.name for g in group.groups])

        for op in group.ops:
            for follow in pipeline_dag.get_follows(op.name):
                if _get_op(group.ops, follow) is not None:
                    # add edge between direct ops
                    group_dag.add_edge(op.name, follow)
                else:
                    _base_group = self._find_base_group(group.groups, follow)
                    if _base_group:
                        # add edge to direct subgroup
                        group_dag.add_edge(op.name, _base_group.name)

            for dependency in pipeline_dag.get_dependencies(op.name):
                if _get_op(group.ops, dependency) is None:
                    _base_group = self._find_base_group(group.groups,
                                                        dependency)
                    if _base_group:
                        # add edge from direct subgroup
                        group_dag.add_edge(_base_group.name, op.name)

        return group_dag

    def _create_op_dag(self, p: dsl.Pipeline) -> _Dag:
        """Create the DAG of the pipeline ops."""
        dag = _Dag(p.ops.keys())

        for op in p.ops.values():
            # dependencies defined by inputs
            for input_value in op.inputs:
                if isinstance(input_value, dsl.PipelineParam):
                    input_param = _extract_pipeline_param(input_value.pattern)
                    if input_param.op_name:
                        dag.add_edge(input_param.op_name, op.name)
                    else:
                        logging.debug("%s depend on pipeline param", op.name)

            # explicit dependencies of current op
            for dependent in op.dependent_names:
                dag.add_edge(dependent, op.name)
        return dag

    def _make_output_file_path_unique(self, run_name: str, op_name: str,
                                      output_file: str) -> str:
        """Alter the file path of output artifact to make sure it's unique in
        local runner.

        kfp compiler will bound a tmp file for each component output,
        which is unique in kfp runtime, but not unique in local runner.
        We alter the file path of the name of current run and op, to
        make it unique in local runner.
        """
        if not output_file.startswith("/tmp/"):
            return output_file
        return f'{self._pipeline_root}/{run_name}/{op_name.lower()}/{output_file[len("/tmp/"):]}'

    def _get_output_file_path(
        self,
        run_name: str,
        pipeline: dsl.Pipeline,
        op_name: str,
        output_name: str = None,
    ) -> str:
        """Get the file path of component output."""

        op_dependency = pipeline.ops[op_name]
        if output_name is None and len(op_dependency.file_outputs) == 1:
            output_name = next(iter(op_dependency.file_outputs.keys()))
        output_file = op_dependency.file_outputs[output_name]
        unique_output_file = self._make_output_file_path_unique(
            run_name, op_name, output_file)
        return unique_output_file

    def _generate_cmd_for_subprocess_execution(
        self,
        run_name: str,
        pipeline: dsl.Pipeline,
        op: dsl.ContainerOp,
        stack: Dict[str, Any],
    ) -> List[str]:
        """Generate shell command to run the op locally."""
        cmd = op.command + op.arguments

        # In debug mode, for `python -c cmd` format command, pydev will insert code before
        # `cmd`, but there is no newline at the end of the inserted code, which will cause
        # syntax error, so we add newline before `cmd`.
        for i in range(len(cmd)):
            if cmd[i] == "-c":
                cmd[i + 1] = "\n" + cmd[i + 1]

        for index, cmd_item in enumerate(cmd):
            if cmd_item in stack:  # Argument is LoopArguments item
                cmd[index] = str(stack[cmd_item])
            elif cmd_item in op.file_outputs.values(
            ):  # Argument is output file
                output_name = next(
                    filter(lambda item: item[1] == cmd_item,
                           op.file_outputs.items()))[0]
                output_param = op.outputs[output_name]
                output_file = cmd_item
                output_file = self._make_output_file_path_unique(
                    run_name, output_param.op_name, output_file)

                os.makedirs(os.path.dirname(output_file), exist_ok=True)
                cmd[index] = output_file
            elif (cmd_item in op.input_artifact_paths.values()
                 ):  # Argument is input artifact file
                input_name = next(
                    filter(
                        lambda item: item[1] == cmd_item,
                        op.input_artifact_paths.items(),
                    ))[0]
                input_param_pattern = op.artifact_arguments[input_name]
                pipeline_param = _extract_pipeline_param(input_param_pattern)
                input_file = self._get_output_file_path(run_name, pipeline,
                                                        pipeline_param.op_name,
                                                        pipeline_param.name)

                cmd[index] = input_file

        return cmd

    def _generate_cmd_for_docker_execution(
        self,
        run_name: str,
        pipeline: dsl.Pipeline,
        op: dsl.ContainerOp,
        stack: Dict[str, Any],
        docker_options: List[str] = []
    ) -> List[str]:
        """Generate the command to run the op in docker locally."""
        cmd = self._generate_cmd_for_subprocess_execution(
            run_name, pipeline, op, stack)

        docker_cmd = [
            "docker",
            "run",
            *docker_options,
            "-v",
            "{pipeline_root}:{pipeline_root}".format(
                pipeline_root=self._pipeline_root),
            op.image,
        ] + cmd
        return docker_cmd

    def _run_group_dag(
        self,
        run_name: str,
        pipeline: dsl.Pipeline,
        pipeline_dag: _Dag,
        current_group: dsl.OpsGroup,
        stack: Dict[str, Any],
        execution_mode: ExecutionMode,
    ) -> bool:
        """Run ops in current group in topological order.

        Args:
            pipeline: kfp.dsl.Pipeline
            pipeline_dag: DAG of pipeline ops
            current_group: current ops group
            stack: stack to trace `LoopArguments`
            execution_mode: Configuration to decide whether the client executes
                component in docker or in local process.
        Returns:
            True if succeed to run the group dag.
        """
        group_dag = self._create_group_dag(pipeline_dag, current_group)

        for node in group_dag.topological_sort():
            subgroup = _get_subgroup(current_group.groups, node)
            if subgroup is not None:  # Node of DAG is subgroup
                success = self._run_group(run_name, pipeline, pipeline_dag, subgroup,
                                stack, execution_mode)
                if not success:
                    return False
            else:  # Node of DAG is op
                op = _get_op(current_group.ops, node)

                execution_mode = (
                    execution_mode
                    if execution_mode else LocalClient.ExecutionMode())
                can_run_locally = execution_mode.mode == LocalClient.ExecutionMode.LOCAL
                exclude = (
                    op.image in execution_mode.images_to_exclude or
                    op.name in execution_mode.ops_to_exclude)
                if exclude:
                    can_run_locally = not can_run_locally

                if can_run_locally:
                    cmd = self._generate_cmd_for_subprocess_execution(
                        run_name, pipeline, op, stack)
                else:
                    cmd = self._generate_cmd_for_docker_execution(
                        run_name, pipeline, op, stack, execution_mode.docker_options)
                process = subprocess.Popen(
                    cmd,
                    shell=False,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    universal_newlines=True,
                )
                # TODO support async process
                logging.info("start task：%s", op.name)
                stdout, stderr = process.communicate()
                if stdout:
                    logging.info(stdout)
                if stderr:
                    logging.error(stderr)
                if process.returncode != 0:
                    logging.error(cmd)
                    return False

    def _run_group(
        self,
        run_name: str,
        pipeline: dsl.Pipeline,
        pipeline_dag: _Dag,
        current_group: dsl.OpsGroup,
        stack: Dict[str, Any],
        execution_mode: ExecutionMode,
    ) -> bool:
        """Run all ops in current group.

        Args:
            run_name: str, the name of this run, can be used to query the run result
            pipeline: kfp.dsl.Pipeline
            pipeline_dag: DAG of pipeline ops
            current_group: current ops group
            stack: stack to trace `LoopArguments`
            execution_mode: Configuration to decide whether the client executes
                component in docker or in local process.
        Returns:
            True if succeed to run the group.
        """
        if current_group.type == dsl.ParallelFor.TYPE_NAME:
            current_group = cast(dsl.ParallelFor, current_group)

            if current_group.items_is_pipeline_param:
                _loop_args = current_group.loop_args
                _param_name = _loop_args.name[:-len(_loop_args
                                                    .LOOP_ITEM_NAME_BASE) - 1]

                _op_dependency = pipeline.ops[_loop_args.op_name]
                _list_file = _op_dependency.file_outputs[_param_name]
                _altered_list_file = self._make_output_file_path_unique(
                    run_name, _loop_args.op_name, _list_file)
                with open(_altered_list_file, "r") as f:
                    _param_values = json.load(f)
                for index, _param_value in enumerate(_param_values):
                    if isinstance(_param_values, (dict, list)):
                        _param_value = json.dumps(_param_value)
                    stack[_loop_args.pattern] = _param_value
                    loop_run_name = "{run_name}/{loop_index}".format(
                        run_name=run_name, loop_index=index)
                    success = self._run_group_dag(
                        loop_run_name,
                        pipeline,
                        pipeline_dag,
                        current_group,
                        stack,
                        execution_mode,
                    )
                    del stack[_loop_args.pattern]
                    if not success:
                        return False
                return True
            else:
                raise Exception("Not implemented")
        else:
            return self._run_group_dag(run_name, pipeline, pipeline_dag, current_group,
                                stack, execution_mode)

    def create_run_from_pipeline_func(
            self,
            pipeline_func: Callable,
            arguments: Mapping[str, str],
            execution_mode: ExecutionMode = ExecutionMode(),
    ):
        """Runs a pipeline locally, either using Docker or in a local process.

        Parameters:
          pipeline_func: pipeline function
          arguments: Arguments to the pipeline function provided as a dict, reference
              to `kfp.client.create_run_from_pipeline_func`
          execution_mode: Configuration to decide whether the client executes component
              in docker or in local process.
        """

        class RunPipelineResult:

            def __init__(self, client: LocalClient, pipeline: dsl.Pipeline,
                         run_id: str, success: bool):
                self._client = client
                self._pipeline = pipeline
                self.run_id = run_id
                self._success = success

            def get_output_file(self, op_name: str, output: str = None):
                return self._client._get_output_file_path(
                    self.run_id, self._pipeline, op_name, output)

            def success(self) -> bool:
                return self._success

            def __repr__(self):
                return "RunPipelineResult(run_id={})".format(self.run_id)

        pipeline_name = sanitize_k8s_name(
            getattr(pipeline_func, "_component_human_name", None) or
            pipeline_func.__name__)
        with dsl.Pipeline(pipeline_name) as pipeline:
            pipeline_func(**arguments)

        run_version = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        run_name = pipeline.name.replace(" ", "_").lower() + "_" + run_version

        pipeline_dag = self._create_op_dag(pipeline)
        success = self._run_group(run_name, pipeline, pipeline_dag, pipeline.groups[0],
                        {}, execution_mode)

        return RunPipelineResult(self, pipeline, run_name, success=success)