def linear():
couler.set_dependencies(lambda: job_a(message="A"), dependencies=None)
couler.set_dependencies(lambda: job_b(message="B"), dependencies=["A"])
couler.set_dependencies(lambda: job_c(message="C"), dependencies=["A"])
couler.set_dependencies(lambda: job_d(message="D"), dependencies=["B"])
def test_run_container_with_dependency_implicit_params_passing(self):
output_path = "/mnt/hello_world.txt"
def producer(step_name):
output_place = couler.create_parameter_artifact(
path=output_path, is_global=True
)
return couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world > %s" % output_place.path],
command=["bash", "-c"],
output=output_place,
step_name=step_name,
)
def consumer(step_name):
couler.run_container(
image="docker/whalesay:latest",
command=["cowsay"],
step_name=step_name,
)
couler.set_dependencies(
lambda: producer(step_name="A"), dependencies=None
)
couler.set_dependencies(
lambda: consumer(step_name="B"), dependencies=["A"]
)
wf = couler.workflow_yaml()
template = wf["spec"]["templates"][1]
# Check input parameters for step A
self.assertEqual(
template["inputs"]["parameters"], [{"name": "para-A-0"}]
)
# Check output parameters for step A
self.assertEqual(
output_path,
template["outputs"]["parameters"][0]["valueFrom"]["path"],
)
self.assertEqual(
"global-" + template["outputs"]["parameters"][0]["name"],
template["outputs"]["parameters"][0]["globalName"],
)
params = wf["spec"]["templates"][0]["dag"]["tasks"][1]["arguments"][
"parameters"
][0]
self.assertEqual(params["name"], "para-B-0")
self.assertTrue(
'"{{workflow.outputs.parameters.output-id-' in params["value"]
)
# Check automatically created emptyDir volume and volume mount
self.assertEqual(
template["volumes"], [{"emptyDir": {}, "name": "couler-out-dir-0"}]
)
self.assertEqual(
template["container"]["volumeMounts"],
[
OrderedDict(
[("name", "couler-out-dir-0"), ("mountPath", "/mnt")]
)
],
)
# Check input parameters for step B
template = wf["spec"]["templates"][2]
self.assertEqual(
template["inputs"]["parameters"], [{"name": "para-B-0"}]
)
def linear():
couler.set_dependencies(lambda: job(name="A"), dependencies=None)
couler.set_dependencies(lambda: job(name="B"), dependencies=["A"])
couler.set_dependencies(lambda: job(name="C"), dependencies=["A"])
couler.set_dependencies(lambda: job(name="D"), dependencies=["B"])
[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)
from couler.argo_submitter import ArgoSubmitter
if __name__ == "__main__":
couler.config_workflow(timeout=3600, time_to_clean=3600 * 1.5)
def pass_step(name):
return couler.run_container(image="alpine:3.6",
command=["sh", "-c", "exit 0"],
step_name=name)
def fail_step(name):
return couler.run_container(image="alpine:3.6",
command=["sh", "-c", "exit 1"],
step_name=name)
couler.set_dependencies(lambda: pass_step("A"), dependencies=None)
couler.set_dependencies(lambda: pass_step("B"), dependencies="A")
couler.set_dependencies(lambda: fail_step("C"), dependencies="A")
couler.set_dependencies(
lambda: pass_step("should-execute-1"),
dependencies="A && (C.Succeeded || C.Failed)",
)
couler.set_dependencies(lambda: pass_step("should-execute-2"),
dependencies="B || C")
couler.set_dependencies(lambda: pass_step("should-not-execute"),
dependencies="B && C")
couler.set_dependencies(
lambda: pass_step("should-execute-3"),
dependencies="should-execute-2.Succeeded || should-not-execute",
)
def test_run_job_with_dependency_implicit_params_passing_from_job(self):
success_condition = "status.succeeded > 0"
failure_condition = "status.failed > 3"
manifest = """
apiVersion: batch/v1
kind: Job
metadata:
generateName: rand-num-
spec:
template:
spec:
containers:
- name: rand
image: python:3.6
command: ["python random_num.py"]
"""
def producer(step_name):
couler.run_job(
manifest=manifest,
success_condition=success_condition,
failure_condition=failure_condition,
step_name=step_name,
)
def consumer(step_name):
return couler.run_container(
image="docker/whalesay:latest",
command=[
"bash",
"-c",
"echo '{{inputs.parameters.para-B-0}}'",
],
step_name=step_name,
)
couler.set_dependencies(lambda: producer(step_name="A"),
dependencies=None)
couler.set_dependencies(lambda: consumer(step_name="B"),
dependencies=["A"])
self.assertEqual(len(couler.workflow.templates), 2)
wf = couler.workflow_yaml()
# Check task for step B in dag tasks
template = wf["spec"]["templates"][0]
self.assertEqual(
[
{
"value": '"{{tasks.A.outputs.parameters.job-id}}"',
"name": "para-B-0",
},
{
"value": '"{{tasks.A.outputs.parameters.job-name}}"',
"name": "para-B-1",
},
{
"value": '"{{tasks.A.outputs.parameters.job-obj}}"',
"name": "para-B-2",
},
],
template["dag"]["tasks"][1]["arguments"]["parameters"],
)
# Check output parameters for step A
template = wf["spec"]["templates"][1]
self.assertEqual(
[
{
"name": "job-name",
"valueFrom": {
"jsonPath": '"{.metadata.name}"'
},
},
{
"name": "job-id",
"valueFrom": {
"jsonPath": '"{.metadata.uid}"'
},
},
{
"name": "job-obj",
"valueFrom": {
"jqFilter": '"."'
}
},
],
template["outputs"]["parameters"],
)
# Check input parameters for step A
template = wf["spec"]["templates"][2]
self.assertEqual(
[{
"name": "para-B-0"
}, {
"name": "para-B-1"
}, {
"name": "para-B-2"
}],
template["inputs"]["parameters"],
)
couler._cleanup()
def test_run_job_with_dependency_implicit_params_passing_from_container(
self):
success_condition = "status.succeeded > 0"
failure_condition = "status.failed > 3"
manifest = """
apiVersion: batch/v1
kind: Job
metadata:
generateName: rand-num-
spec:
template:
spec:
containers:
- name: rand
image: python:3.6
command: ["python random_num.py"]
"""
output_path = "/tmp/hello_world.txt"
def producer(step_name):
output_place = couler.create_parameter_artifact(path=output_path)
return couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world > %s" % output_place.path],
command=["bash", "-c"],
output=output_place,
step_name=step_name,
)
def consumer(step_name):
couler.run_job(
manifest=manifest,
success_condition=success_condition,
failure_condition=failure_condition,
step_name=step_name,
env={"k1": "v1"},
)
couler.set_dependencies(lambda: producer(step_name="A"),
dependencies=None)
couler.set_dependencies(lambda: consumer(step_name="B"),
dependencies=["A"])
self.assertEqual(len(couler.workflow.templates), 2)
wf = couler.workflow_yaml()
# Check input and output parameters for step A
template = wf["spec"]["templates"][1]
self.assertEqual(template["inputs"]["parameters"], [{
"name": "para-A-0"
}])
self.assertEqual(
output_path,
template["outputs"]["parameters"][0]["valueFrom"]["path"],
)
# Check env for step B
manifest_dict = yaml.safe_load(
wf["spec"]["templates"][2]["resource"]["manifest"])
self.assertEqual(manifest_dict["spec"]["env"][0], {
"name": "k1",
"value": "v1"
})
envs = manifest_dict["spec"]["env"][1]
self.assertEqual(envs["name"], "couler.inferred_outputs.0")
self.assertTrue(
"{{tasks.A.outputs.parameters.output-id-" in envs["value"])
couler._cleanup()
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]),
)
self.assertEqual(
wf["spec"]["templates"][5],
{
"name":
"exit-handler",
"steps": [
[
OrderedDict([
("name", "success-exit"),
("template", "success-exit"),
("when", "{{workflow.status}} == Succeeded"),
])
],
[
OrderedDict([
("name", "failure-exit"),
("template", "failure-exit"),
("when", "{{workflow.status}} == Failed"),
])
],
],
},
)
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"]),
(
"env",
[
{
"name": "NVIDIA_VISIBLE_DEVICES",
"value": ""
},
{
"name": "NVIDIA_DRIVER_CAPABILITIES",
"value": "",
},
],
),
]),
)
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()
