def test_set_workflow_exit_handler(self):
couler._cleanup()
flip_coin()
couler.set_exit_handler(couler.WFStatus.Succeeded, heads)
couler.set_exit_handler(couler.WFStatus.Failed, tails)
self.check_argo_yaml("workflow_basic_golden.yaml")
couler._cleanup()
def test_exit_handler(self):
def send_mail():
return couler.run_container(image="alpine:3.6",
command=["echo", "send mail"])
couler.run_container(image="alpine:3.6", command=["exit", "1"])
couler.set_exit_handler(couler.WFStatus.Failed, send_mail)
proto_wf = get_default_proto_workflow()
self.assertEqual(len(proto_wf.exit_handler_steps), 1)
s = proto_wf.exit_handler_steps[0]
self.assertEqual(s.when, "{{workflow.status}} == Failed")
def test_raise(self):
with self.assertRaises(ValueError):
def some_raise():
raise ValueError("test this")
couler.set_exit_handler(couler.WFStatus.Failed, some_raise)
couler._cleanup()
couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
content = pyaml.dump(couler.workflow_yaml())
self.assertNotIn("onExit", content)
def test_set_dependencies_with_exit_handler(self):
def producer():
return couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
)
def consumer():
return couler.run_container(
image="docker/whalesay:latest",
command=["cowsay"],
step_name="B",
)
def exit_handler_succeeded():
return couler.run_container(
image="docker/whalesay:latest",
command=["cowsay"],
step_name="success-exit",
)
def exit_handler_failed():
return couler.run_container(
image="docker/whalesay:latest",
command=["cowsay"],
step_name="failure-exit",
)
couler.set_dependencies(lambda: producer(), dependencies=None)
couler.set_dependencies(lambda: consumer(), dependencies=["A"])
couler.set_exit_handler(couler.WFStatus.Succeeded,
exit_handler_succeeded)
couler.set_exit_handler(couler.WFStatus.Failed, exit_handler_failed)
wf = couler.workflow_yaml()
self.assertEqual(wf["spec"]["onExit"], "exit-handler")
expected_container_spec = (
"container",
OrderedDict([("image", "docker/whalesay:latest"),
("command", ["cowsay"])]),
)
self.assertEqual(
wf["spec"]["templates"][3],
OrderedDict([("name", "success-exit"), expected_container_spec]),
)
self.assertEqual(
wf["spec"]["templates"][4],
OrderedDict([("name", "failure-exit"), expected_container_spec]),
)
couler._cleanup()
def test_set_dependencies_none_with_exit_handler(self):
couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
couler.set_dependencies(lambda: job_b(message="B"), dependencies=["A"])
def job_exit():
return couler.run_container(
image="docker/whalesay:latest",
command=["cowsay"],
step_name="C",
)
couler.set_exit_handler(couler.WFStatus.Failed, job_exit)
content = pyaml.dump(couler.workflow_yaml())
self.assertIn("{{workflow.status}} == Failed", content)
def test_exit_callable(self):
cls = self.create_callable_cls(lambda: heads())
instance = cls()
func_names = ["a", "b", "c", "self"]
for func_name in func_names:
flip_coin()
if func_name == "self":
couler.set_exit_handler(couler.WFStatus.Succeeded, instance)
else:
couler.set_exit_handler(couler.WFStatus.Succeeded,
getattr(instance, func_name))
couler.set_exit_handler(couler.WFStatus.Failed, tails)
self.check_argo_yaml("workflow_basic_golden.yaml")
couler._cleanup()
[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)
