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
assert not os.path.exists(output_file_path)
run_result = run_pipeline_func_locally(
_pipeline,
{},
execution_mode=LocalClient.ExecutionMode(
mode="docker", ops_to_exclude=["cat-on-image-not-exist"]),
)
output_file_path = run_result.get_output_file("cat-on-image-not-exist")
with open(output_file_path, "r") as f:
line = f.readline()
assert "exclude ops" in line
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_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_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_connect(self):
def _pipeline():
_local_loader = local_loader(self.temp_file_path)
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_local_file(self):
def _pipeline(file_path: str):
local_loader(file_path)
run_result = run_pipeline_func_locally(
_pipeline,
{"file_path": self.temp_file_path},
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_docker_options(self):
@light_component()
def check_option(dst: OutputPath):
import os
with open(dst, "w") as f:
f.write(os.environ["foo"])
def _pipeline():
check_option()
run_result = run_pipeline_func_locally(
_pipeline, {},
execution_mode=LocalClient.ExecutionMode(
mode="docker", docker_options=["-e", "foo=bar"]))
assert run_result.success
output_file_path = run_result.get_output_file("check-option")
with open(output_file_path, "r") as f:
line = f.readline()
assert "bar" in line
