def test_when(self):
def random_code():
import random
res = "heads" if random.randint(0, 1) == 0 else "tails"
print(res)
def heads():
return couler.run_container(
image="alpine:3.6",
command=["sh", "-c", 'echo "it was heads"'])
def tails():
return couler.run_container(
image="alpine:3.6",
command=["sh", "-c", 'echo "it was tails"'])
result = couler.run_script(image="python:alpine3.6",
source=random_code)
couler.when(couler.equal(result, "heads"), lambda: heads())
couler.when(couler.equal(result, "tails"), lambda: tails())
proto_wf = get_default_proto_workflow()
step_heads = proto_wf.steps[1].steps[0]
# condition is like: "{{steps.test-when-550.outputs.result}} == heads"
self.assertTrue(step_heads.when.startswith("{{steps.test-when-"))
self.assertTrue(step_heads.when.endswith(".outputs.result}} == heads"))
def test_when_order_two(self):
couler.steps = OrderedDict()
couler.update_steps = True
couler.when(couler.equal(flip_coin(), "heads"), lambda: heads())
couler.when(couler.equal(flip_coin(), "tails"), lambda: tails())
steps = couler.workflow.get_steps_dict()
hope_steps = [
[OrderedDict({"name": "flip-coin-67", "template": "flip-coin"})],
[
OrderedDict(
{
"name": "heads-67",
"template": "heads",
"when": "{{steps.flip-coin-67.outputs.result}} == heads", # noqa: E501
}
)
],
[OrderedDict({"name": "flip-coin-68", "template": "flip-coin"})],
[
OrderedDict(
{
"name": "tails-68",
"template": "tails",
"when": "{{steps.flip-coin-68.outputs.result}} == tails", # noqa: E501
}
)
],
]
self.assertEqual(hope_steps, steps)
def test_when_order_three(self):
couler.steps = OrderedDict()
couler.update_steps = True
output_1 = flip_coin()
couler.when(couler.equal(output_1, "heads"), lambda: heads())
output_2 = flip_coin()
couler.when(couler.equal(output_1, "tails"), lambda: tails())
couler.when(couler.equal(output_2, "heads"), lambda: heads())
couler.when(couler.equal(output_2, "tails"), lambda: tails())
steps = couler.workflow.get_steps_dict()
hope_steps = [
[
OrderedDict([("name", "flip-coin-116"),
("template", "flip-coin")])
],
[
OrderedDict([
("name", "heads-117"),
("template", "heads"),
(
"when",
"{{steps.flip-coin-116.outputs.result}} == heads",
),
]),
OrderedDict([
("name", "tails-120"),
("template", "tails"),
(
"when",
"{{steps.flip-coin-116.outputs.result}} == tails",
),
]),
],
[
OrderedDict([("name", "flip-coin-119"),
("template", "flip-coin")])
],
[
OrderedDict([
("name", "heads-121"),
("template", "heads"),
(
"when",
"{{steps.flip-coin-119.outputs.result}} == heads",
),
]),
OrderedDict([
("name", "tails-122"),
("template", "tails"),
(
"when",
"{{steps.flip-coin-119.outputs.result}} == tails",
),
]),
],
]
self.assertEqual(hope_steps, steps)
def test_while_callable(self):
cls = self.create_callable_cls(lambda: flip_coin())
instance = cls()
func_names = ["a", "b", "c", "self"]
for func_name in func_names:
if func_name == "self":
couler.exec_while(couler.equal("tails"), instance)
else:
couler.exec_while(couler.equal("tails"),
getattr(instance, func_name))
self.check_argo_yaml("while_golden.yaml")
couler._cleanup()
def test_when_callable(self):
couler.steps = OrderedDict()
couler.update_steps = True
cls = self.create_callable_cls(lambda: heads())
instance = cls()
func_names = ["a", "b", "c", "self"]
for func_name in func_names:
if func_name == "self":
couler.when(couler.equal(flip_coin(), "heads"), instance)
else:
couler.when(
couler.equal(flip_coin(), "heads"),
getattr(instance, func_name),
)
couler.when(couler.equal(flip_coin(), "tails"), lambda: tails())
couler._cleanup()
def test_set_dependencies_with_conditional(self):
couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
couler.set_dependencies(lambda: job_b(message="B"), dependencies=["A"])
couler.set_dependencies(
lambda: couler.when(couler.equal(flip_coin(), "heads"), lambda:
heads()),
dependencies=["B"],
)
self.check_argo_yaml("dag_golden_2.yaml")
def flip_coin():
return couler.run_script(image="python:alpine3.6", source=random_code)
def heads():
return couler.run_container(image="alpine:3.6",
command=["sh", "-c", 'echo "it was heads"'])
def tails():
return couler.run_container(image="alpine:3.6",
command=["sh", "-c", 'echo "it was tails"'])
if __name__ == "__main__":
for impl_type in [_SubmitterImplTypes.GO, _SubmitterImplTypes.PYTHON]:
os.environ[_SUBMITTER_IMPL_ENV_VAR_KEY] = impl_type
print("Submitting flip coin example workflow via %s implementation" %
impl_type)
couler.config_workflow(
name="flip-coin-%s" % impl_type.lower(),
timeout=3600,
time_to_clean=3600 * 1.5,
)
result = flip_coin()
couler.when(couler.equal(result, "heads"), lambda: heads())
couler.when(couler.equal(result, "tails"), lambda: tails())
submitter = ArgoSubmitter(namespace="argo")
couler.run(submitter=submitter)
def test_while(self):
couler.exec_while(couler.equal("tails"), lambda: flip_coin())
self.check_argo_yaml("while_golden.yaml")
[lambda: job_a(message="A"), lambda: job_c(message="C")], # A -> C
[lambda: job_b(message="B"), lambda: job_d(message="D")], # B -> D
[lambda: job_c(message="C"), lambda: job_d(message="D")], # C -> D
])
if __name__ == "__main__":
couler.config_workflow(timeout=3600, time_to_clean=3600 * 1.5)
# 1) Add a linear DAG.
linear()
# 2) Add another step that depends on D and flips a coin.
# 3) If the result is "heads", another child step is also
# added to the entire workflow.
couler.set_dependencies(
lambda: couler.when(
couler.equal(conditional_parent(), "heads"),
lambda: conditional_child(),
),
dependencies=["D"],
)
# 4) Add an exit handler that runs when the workflow succeeds.
couler.set_exit_handler(couler.WFStatus.Succeeded, exit_handler_succeeded)
# 5) Add an exit handler that runs when the workflow failed.
couler.set_exit_handler(couler.WFStatus.Failed, exit_handler_failed)
submitter = ArgoSubmitter(namespace="argo")
wf = couler.run(submitter=submitter)
wf_name = wf["metadata"]["name"]
print("Workflow %s has been submitted for DAG example" % wf_name)
