<commit_before># Copyright 2021 Google LLC
#
# 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 unittest
from typing import Callable

import kfp
from kfp import LocalClient, run_pipeline_func_locally

InputPath = kfp.components.InputPath()
OutputPath = kfp.components.OutputPath()

BASE_IMAGE = "python:3.7"


def light_component(
    base_image: str = BASE_IMAGE,
):
    """Decorator of kfp light component with customized parameters

    Usage:
    ```python
    @light_component(base_image="python:3.7")
    def a_component(src: kfp.components.InputPath(), ...):
        ...
    ```
    """

    def wrapper(func: Callable):
        return kfp.components.create_component_from_func(
            func=func,
            base_image=base_image,
        )

    return wrapper


@light_component()
def hello(name: str):
    print(f"hello {name}")


@light_component()
def local_loader(src: str, dst: kfp.components.OutputPath()):
    import os
    import shutil

    if os.path.exists(src):
        shutil.copyfile(src, dst)


@light_component()
def flip_coin(dst: kfp.components.OutputPath()):
    import random

    result = "head" if random.randint(0, 1) == 0 else "tail"
    with open(dst, "w") as f:
        f.write(result)


@light_component()
def list(dst: kfp.components.OutputPath()):
    import json

    with open(dst, "w") as f:
        json.dump(["hello", "world", "kfp"], f)


@light_component()
def component_connect_demo(
    src: kfp.components.InputPath(), dst: kfp.components.OutputPath()
):
    with open(src, "r") as f:
        line = f.readline()
        print(f"read first line: {line}")
        with open(dst, "w") as fw:
            fw.write(f"{line} copied")


class LocalRunnerTest(unittest.TestCase):
    def setUp(self):
        import tempfile

        self.temp_file = tempfile.NamedTemporaryFile()
        with open(self.temp_file.name, "w") as f:
            f.write("hello world")

    def test_run_local(self):
        def _pipeline(name: str):
            hello(name)

        run_pipeline_func_locally(
            _pipeline,
            {"name": "world"},
            execution_mode=LocalClient.ExecutionMode("local"),
        )

    def test_local_file(self):
        def _pipeline(file_path: str):
            local_loader(file_path)

        run_result = run_pipeline_func_locally(
            _pipeline,
            {"file_path": self.temp_file.name},
            execution_mode=LocalClient.ExecutionMode("local"),
        )
        output_file_path = run_result.get_output_file("local-loader")

        with open(output_file_path, "r") as f:
            line = f.readline()
            assert "hello" in line

    def test_condition(self):
        def _pipeline():
            _flip = flip_coin()
            with kfp.dsl.Condition(_flip.output == "head"):
                hello("head")

            with kfp.dsl.Condition(_flip.output == "tail"):
                hello("tail")

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    def test_for(self):
        @light_component()
        def cat(item, dst: OutputPath):
            with open(dst, "w") as f:
                f.write(item)

        def _pipeline():
            with kfp.dsl.ParallelFor(list().output) as item:
                cat(item)

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    def test_connect(self):
        def _pipeline():
            _local_loader = local_loader(self.temp_file.name)
            component_connect_demo(_local_loader.output)

        run_result = run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )
        output_file_path = run_result.get_output_file("component-connect-demo")

        with open(output_file_path, "r") as f:
            line = f.readline()
            assert "copied" in line

    def test_command_argument_in_any_format(self):
        def echo():
            return kfp.dsl.ContainerOp(
                name="echo",
                image=BASE_IMAGE,
                command=["echo", "hello world", ">", "/tmp/outputs/output_file"],
                arguments=[],
                file_outputs={"output": "/tmp/outputs/output_file"},
            )

        def _pipeline():
            _echo = echo()
            component_connect_demo(_echo.output)

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    def test_execution_mode_exclude_op(self):
        @light_component(base_image="image_not_exist")
        def cat_on_image_not_exist(name: str, dst: OutputPath):
            with open(dst, "w") as f:
                f.write(name)

        def _pipeline():
            cat_on_image_not_exist("exclude ops")

        run_result = run_pipeline_func_locally(
            _pipeline,
            {},
            execution_mode=LocalClient.ExecutionMode(mode="docker"),
        )
        output_file_path = run_result.get_output_file("cat-on-image-not-exist")
        import os
<commit_msg>fix(tests): Disable the SDK unittest that requires docker binary in local environment. (#5211)<commit_after># Copyright 2021 Google LLC
#
# 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 unittest
from typing import Callable

import kfp
from kfp import LocalClient, run_pipeline_func_locally

InputPath = kfp.components.InputPath()
OutputPath = kfp.components.OutputPath()

BASE_IMAGE = "python:3.7"


def light_component(
    base_image: str = BASE_IMAGE,
):
    """Decorator of kfp light component with customized parameters

    Usage:
    ```python
    @light_component(base_image="python:3.7")
    def a_component(src: kfp.components.InputPath(), ...):
        ...
    ```
    """

    def wrapper(func: Callable):
        return kfp.components.create_component_from_func(
            func=func,
            base_image=base_image,
        )

    return wrapper


@light_component()
def hello(name: str):
    print(f"hello {name}")


@light_component()
def local_loader(src: str, dst: kfp.components.OutputPath()):
    import os
    import shutil

    if os.path.exists(src):
        shutil.copyfile(src, dst)


@light_component()
def flip_coin(dst: kfp.components.OutputPath()):
    import random

    result = "head" if random.randint(0, 1) == 0 else "tail"
    with open(dst, "w") as f:
        f.write(result)


@light_component()
def list(dst: kfp.components.OutputPath()):
    import json

    with open(dst, "w") as f:
        json.dump(["hello", "world", "kfp"], f)


@light_component()
def component_connect_demo(
    src: kfp.components.InputPath(), dst: kfp.components.OutputPath()
):
    with open(src, "r") as f:
        line = f.readline()
        print(f"read first line: {line}")
        with open(dst, "w") as fw:
            fw.write(f"{line} copied")


class LocalRunnerTest(unittest.TestCase):
    def setUp(self):
        import tempfile

        self.temp_file = tempfile.NamedTemporaryFile()
        with open(self.temp_file.name, "w") as f:
            f.write("hello world")

    def test_run_local(self):
        def _pipeline(name: str):
            hello(name)

        run_pipeline_func_locally(
            _pipeline,
            {"name": "world"},
            execution_mode=LocalClient.ExecutionMode("local"),
        )

    def test_local_file(self):
        def _pipeline(file_path: str):
            local_loader(file_path)

        run_result = run_pipeline_func_locally(
            _pipeline,
            {"file_path": self.temp_file.name},
            execution_mode=LocalClient.ExecutionMode("local"),
        )
        output_file_path = run_result.get_output_file("local-loader")

        with open(output_file_path, "r") as f:
            line = f.readline()
            assert "hello" in line

    def test_condition(self):
        def _pipeline():
            _flip = flip_coin()
            with kfp.dsl.Condition(_flip.output == "head"):
                hello("head")

            with kfp.dsl.Condition(_flip.output == "tail"):
                hello("tail")

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    def test_for(self):
        @light_component()
        def cat(item, dst: OutputPath):
            with open(dst, "w") as f:
                f.write(item)

        def _pipeline():
            with kfp.dsl.ParallelFor(list().output) as item:
                cat(item)

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    def test_connect(self):
        def _pipeline():
            _local_loader = local_loader(self.temp_file.name)
            component_connect_demo(_local_loader.output)

        run_result = run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )
        output_file_path = run_result.get_output_file("component-connect-demo")

        with open(output_file_path, "r") as f:
            line = f.readline()
            assert "copied" in line

    def test_command_argument_in_any_format(self):
        def echo():
            return kfp.dsl.ContainerOp(
                name="echo",
                image=BASE_IMAGE,
                command=["echo", "hello world", ">", "/tmp/outputs/output_file"],
                arguments=[],
                file_outputs={"output": "/tmp/outputs/output_file"},
            )

        def _pipeline():
            _echo = echo()
            component_connect_demo(_echo.output)

        run_pipeline_func_locally(
            _pipeline, {}, execution_mode=LocalClient.ExecutionMode("local")
        )

    @unittest.skip('docker is not installed in CI environment.')
    def test_execution_mode_exclude_op(self):
        @light_component(base_image="image_not_exist")
        def cat_on_image_not_exist(name: str, dst: OutputPath):
            with open(dst, "w") as f:
                f.write(name)

        def _pipeline():
            cat_on_image_not_exist("exclude ops")

        run_result = run_pipeline_func_locally(
            _pipeline,
            {},
            execution_mode=LocalClient.ExecutionMode(mode="docker"),
        )
        output_file_path = run_result.get_output_file("cat-on-image-not-exist")
        import os

{'fraction_non_alphanumeric': 0.11037024039305142, 'fraction_numerical': 0.0029829794700824705, 'mean_word_length': 2.45948406676783, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4316710', 'n_tokens_mistral': 3432, 'n_tokens_neox': 3347, 'n_words': 767}