def test_wf_container_task_multiple():
square = ContainerTask(
name="square",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(val=int),
outputs=kwtypes(out=int),
image="alpine",
command=["sh", "-c", "echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"],
)
sum = ContainerTask(
name="sum",
input_data_dir="/var/flyte/inputs",
output_data_dir="/var/flyte/outputs",
inputs=kwtypes(x=int, y=int),
outputs=kwtypes(out=int),
image="alpine",
command=["sh", "-c", "echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out"],
)
@workflow
def raw_container_wf(val1: int, val2: int) -> int:
return sum(x=square(val=val1), y=square(val=val2))
with task_mock(square) as square_mock, task_mock(sum) as sum_mock:
square_mock.side_effect = lambda val: val * val
assert square(val=10) == 100
sum_mock.side_effect = lambda x, y: x + y
assert sum(x=10, y=10) == 20
assert raw_container_wf(val1=10, val2=10) == 200
def test_serialization():
square = ContainerTask(
name="square",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(val=int),
outputs=kwtypes(out=int),
image="alpine",
command=[
"sh", "-c",
"echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"
],
)
sum = ContainerTask(
name="sum",
input_data_dir="/var/flyte/inputs",
output_data_dir="/var/flyte/outputs",
inputs=kwtypes(x=int, y=int),
outputs=kwtypes(out=int),
image="alpine",
command=[
"sh", "-c",
"echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out"
],
)
@workflow
def raw_container_wf(val1: int, val2: int) -> int:
return sum(x=square(val=val1), y=square(val=val2))
default_img = Image(name="default", fqn="test", tag="tag")
serialization_settings = context_manager.SerializationSettings(
project="project",
domain="domain",
version="version",
env=None,
image_config=ImageConfig(default_image=default_img,
images=[default_img]),
)
wf_spec = get_serializable(OrderedDict(), serialization_settings,
raw_container_wf)
assert wf_spec is not None
assert wf_spec.template is not None
assert len(wf_spec.template.nodes) == 3
sqn_spec = get_serializable(OrderedDict(), serialization_settings, square)
assert sqn_spec.template.container.image == "alpine"
sumn_spec = get_serializable(OrderedDict(), serialization_settings, sum)
assert sumn_spec.template.container.image == "alpine"
def test_wf_container_task():
@task
def t1(a: int) -> (int, str):
return a + 2, str(a) + "-HELLO"
t2 = ContainerTask(
"raw",
image="alpine",
inputs=kwtypes(a=int, b=str),
input_data_dir="/tmp",
output_data_dir="/tmp",
command=["cat"],
arguments=["/tmp/a"],
)
@workflow
def wf(a: int):
x, y = t1(a=a)
t2(a=x, b=y)
with task_mock(t2) as mock:
mock.side_effect = lambda a, b: None
assert t2(a=10, b="hello") is None
wf(a=10)
def test_resolver_load_task():
# any task is fine, just copied one
square = ContainerTask(
name="square",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(val=int),
outputs=kwtypes(out=int),
image="alpine",
command=[
"sh", "-c",
"echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"
],
)
resolver = TaskTemplateResolver()
ts = get_serializable(OrderedDict(), serialization_settings, square)
file = tempfile.NamedTemporaryFile().name
# load_task should create an instance of the path to the object given, doesn't need to be a real executor
write_proto_to_file(ts.template.to_flyte_idl(), file)
shim_task = resolver.load_task(
[file, f"{Placeholder.__module__}.Placeholder"])
assert isinstance(shim_task.executor, Placeholder)
assert shim_task.task_template.id.name == "square"
assert shim_task.task_template.interface.inputs["val"] is not None
assert shim_task.task_template.interface.outputs["out"] is not None
def test_container():
@task
def t1(a: int) -> (int, str):
return a + 2, str(a) + "-HELLO"
t2 = ContainerTask(
"raw",
image="alpine",
inputs=kwtypes(a=int, b=str),
input_data_dir="/tmp",
output_data_dir="/tmp",
command=["cat"],
arguments=["/tmp/a"],
)
sdk_task = get_serializable(serialization_settings, t2, fast=True)
assert "pyflyte" not in sdk_task.container.args
def test_container():
@task
def t1(a: int) -> (int, str):
return a + 2, str(a) + "-HELLO"
t2 = ContainerTask(
"raw",
image="alpine",
inputs=kwtypes(a=int, b=str),
input_data_dir="/tmp",
output_data_dir="/tmp",
command=["cat"],
arguments=["/tmp/a"],
requests=Resources(mem="400Mi", cpu="1"),
)
task_spec = get_serializable(OrderedDict(),
serialization_settings,
t2,
fast=True)
assert "pyflyte" not in task_spec.template.container.args
#
# A :py:class:`flytekit.ContainerTask` denotes an arbitrary container. In the following example, the name of the task
# is ``calculate_ellipse_area_shell``. This name has to be unique in the entire project. Users can specify:
#
# - ``input_data_dir`` -> where inputs will be written to
# - ``output_data_dir`` -> where Flyte will expect the outputs to exist.
#
# inputs and outputs specify the interface for the task, thus it should be an ordered dictionary of typed input and
# output variables.
calculate_ellipse_area_shell = ContainerTask(
name="ellipse-area-metadata-shell",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(a=float, b=float),
outputs=kwtypes(area=float, metadata=str),
image="rawcontainers-shell:v1",
command=[
"./calculate-ellipse-area.sh",
"/var/inputs",
"/var/outputs",
],
)
calculate_ellipse_area_python = ContainerTask(
name="ellipse-area-metadata-python",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(a=float, b=float),
outputs=kwtypes(area=float, metadata=str),
image="rawcontainers-python:v1",
command=[
# %%
# A :py:class:`flytekit.ContainerTask` denotes an arbitrary container. In the following example, the name of the task
# is square. This name has to be unique in the entire project. Users can specify some
#
# - ``input_data_dir`` -> where inputs will be written to
# - ``output_data_dir`` -> where Flyte will expect the outputs to exist.
#
# inputs and outputs specify the interface for the task, thus it should be an ordered dictionary of typed input and
# output variables
square = ContainerTask(
name="square",
input_data_dir="/var/inputs",
output_data_dir="/var/outputs",
inputs=kwtypes(val=int),
outputs=kwtypes(out=int),
image="alpine",
command=[
"sh",
"-c",
"echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out",
],
)
sum = ContainerTask(
name="sum",
input_data_dir="/var/flyte/inputs",
output_data_dir="/var/flyte/outputs",
inputs=kwtypes(x=int, y=int),
outputs=kwtypes(out=int),
image="alpine",