def test_map_function_callable(self):
test_paras = ["t1", "t2", "t3"]
callable_cls = self.create_callable_cls(lambda x: consume(x))
instance = callable_cls()
func_names = ["a", "b", "c", "self"]
for func_name in func_names:
if func_name == "self":
couler.map(instance, test_paras)
else:
couler.map(getattr(instance, func_name), test_paras)
expected_with_items = [
{
"para-consume-0": "t1"
},
{
"para-consume-0": "t2"
},
{
"para-consume-0": "t3"
},
]
wf = couler.workflow_yaml()
templates = wf["spec"]["templates"]
steps_template = templates[0]
map_step = steps_template["steps"][0][0]
self.assertListEqual(map_step["withItems"], expected_with_items)
couler._cleanup()
def test_resource_setup(self):
couler.run_container(
image="docker/whalesay",
command=["cowsay"],
args=["resource test"],
resources={"cpu": "1", "memory": "100Mi"},
)
# Because test environment between local and CI is different,
# we can not compare the YAML directly.
_test_data_dir = "test_data"
test_data_dir = os.path.join(os.path.dirname(__file__), _test_data_dir)
with open(
os.path.join(test_data_dir, "resource_config_golden.yaml"), "r"
) as f:
expected = yaml.safe_load(f)
output = yaml.safe_load(
pyaml.dump(couler.workflow_yaml(), string_val_style="plain")
)
_resources = output["spec"]["templates"][1]["container"]["resources"]
_expected_resources = expected["spec"]["templates"][1]["container"][
"resources"
]
self.assertEqual(_resources, _expected_resources)
couler._cleanup()
def test_run_container_with_workflow_volume(self):
pvc = VolumeClaimTemplate("workdir")
volume_mount = VolumeMount("workdir", "/mnt/vol")
couler.create_workflow_volume(pvc)
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
volume_mounts=[volume_mount],
)
volume_mount = VolumeMount("workdir", "/mnt/vol")
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
volume_mounts=[volume_mount],
)
wf = couler.workflow_yaml()
self.assertEqual(len(wf["spec"]["volumeClaimTemplates"]), 1)
self.assertEqual(wf["spec"]["volumeClaimTemplates"][0], pvc.to_dict())
self.assertEqual(
wf["spec"]["templates"][1]["container"]["volumeMounts"][0],
volume_mount.to_dict(),
)
couler._cleanup()
def pipeline(data: str, workflow_name: str, image_registry: str,
work_path: str):
'''
This is the main entry point
Called from __init__ ArgoExecutor
'''
couler._cleanup()
sb_step.global_steps = []
couler.workflow.name = workflow_name if workflow_name else "workflow"
wfl = workflow()
wfl.image_registry = image_registry
wfl.work_path = work_path
data = wfl.parse(data)
couler.workflow.dag_mode = True
builder_as_deps = []
for c in wfl.containers:
couler.set_dependencies(lambda: wfl.builder_phase(c),
dependencies=None)
# From cargo documntation: https://github.com/argoproj/argo-workflows/blob/master/docs/enhanced-depends-logic.md
# Create a string representing the enhanced depends logic that specifies
# dependencies based on their statuses.
builder_as_deps.append(f"builder{c.name}.Succeeded"
) # += " builder"+c.name+".Succeeded" + " &&"
builder_as_deps = ' && '.join(
builder_as_deps) # builder_as_deps[:len(builder_as_deps) - 2]
wfl.dag_phase(data, builder_as_deps)
ret = yaml()
#couler._cleanup()
return ret
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_run_container_with_dependency_implicit_params_passing(self):
output_path = "/tmp/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(params["value"] in [
# Note that the "output-id-92" case is needed for
# Python 3.8.
'"{{workflow.outputs.parameters.output-id-113}}"',
'"{{workflow.outputs.parameters.output-id-114}}"',
])
# Check input parameters for step B
template = wf["spec"]["templates"][2]
self.assertEqual(template["inputs"]["parameters"], [{
"name": "para-B-0"
}])
couler._cleanup()
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_map_function(self):
test_paras = ["t1", "t2", "t3"]
couler.map(lambda x: consume(x), test_paras)
wf = couler.workflow_yaml()
templates = wf["spec"]["templates"]
self.assertEqual(len(templates), 2)
# We should have a 'consume' template
consume_template = templates[1]
self.assertEqual(consume_template["name"], "consume")
# Check input parameters
expected_paras = [{"name": "para-consume-0"}]
self.assertListEqual(consume_template["inputs"]["parameters"],
expected_paras)
# Check container
expected_container = {
"image": "docker/whalesay:latest",
"command": ["cowsay"],
"args": ['"{{inputs.parameters.para-consume-0}}"'],
}
self.assertDictEqual(consume_template["container"], expected_container)
# Check the steps template
steps_template = templates[0]
self.assertTrue(steps_template["name"] in ["pytest", "runpy"])
self.assertEqual(len(steps_template["steps"]), 1)
self.assertEqual(len(steps_template["steps"][0]), 1)
map_step = steps_template["steps"][0][0]
self.assertIn("consume", map_step["name"])
self.assertEqual(map_step["template"], "consume")
# Check arguments
expected_paras = [{
"name": "para-consume-0",
"value": '"{{item.para-consume-0}}"'
}]
self.assertListEqual(map_step["arguments"]["parameters"],
expected_paras)
# Check withItems
expected_with_items = [
{
"para-consume-0": "t1"
},
{
"para-consume-0": "t2"
},
{
"para-consume-0": "t3"
},
]
self.assertListEqual(map_step["withItems"], expected_with_items)
couler._cleanup()
def setUp(self):
self.state_keys = dir(states)
self.tc_dump = {}
for key in self.state_keys:
if key.startswith("__"):
continue
if key.startswith("_"):
value = getattr(states, key)
if callable(value):
continue
self.tc_dump[key] = value
couler._cleanup()
def test_run_concurrent_callable(self):
callable_cls = self.create_callable_cls(lambda: whalesay("hello1"))
func_names = ["a", "b", "c", "self"]
instance = callable_cls()
for func_name in func_names:
if func_name == "self":
func = instance
else:
func = getattr(instance, func_name)
couler.concurrent([func, lambda: heads(), lambda: tails()])
self.verify_concurrent_step("run_concurrent_golden.yaml")
couler._cleanup()
def tearDown(self):
couler._cleanup()
for key in self.state_keys:
if key.startswith("__"):
continue
if key.startswith("_"):
value = getattr(states, key)
if callable(value):
continue
self.assertEqual(value,
self.tc_dump[key],
msg="state not cleanup:%s" % key)
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_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_run_container_with_toleration(self):
toleration = Toleration("example-toleration", "Exists", "NoSchedule")
couler.add_toleration(toleration)
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
)
wf = couler.workflow_yaml()
self.assertEqual(wf["spec"]["tolerations"][0], toleration.to_dict())
couler._cleanup()
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()
def test_run_container_with_node_selector(self):
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
node_selector={"beta.kubernetes.io/arch": "amd64"},
)
wf = couler.workflow_yaml()
self.assertEqual(
wf["spec"]["templates"][1]["nodeSelector"],
{"beta.kubernetes.io/arch": "amd64"},
)
couler._cleanup()
def test_cluster_config(self):
couler.config_workflow(cluster_config_file=os.path.join(
os.path.dirname(__file__), "test_data/dummy_cluster_config.py"))
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
)
wf = couler.workflow_yaml()
self.assertTrue(wf["spec"]["hostNetwork"])
self.assertEqual(wf["spec"]["templates"][1]["tolerations"], [])
couler._cleanup()
def test_workflow_service_account(self):
self.assertIsNone(couler.workflow.service_account)
flip_coin()
self.assertNotIn("serviceAccountName", couler.workflow_yaml()["spec"])
couler.config_workflow(service_account="test-serviceaccount")
self.assertEqual(couler.workflow.service_account,
"test-serviceaccount")
self.assertEqual(
couler.workflow_yaml()["spec"]["serviceAccountName"],
"test-serviceaccount",
)
couler._cleanup()
self.assertIsNone(couler.workflow.service_account)
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_passing_artifact_implicitly(self):
default_path = "/mnt/t1.txt"
def producer(step_name):
output_artifact = couler.create_oss_artifact(
default_path,
bucket="test-bucket/",
accesskey_id="abcde",
accesskey_secret="abc12345",
key="osspath/t1",
endpoint="xyz.com",
)
outputs = couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world > %s" % output_artifact.path],
command=["bash", "-c"],
output=output_artifact,
step_name=step_name,
)
return outputs
def consumer(step_name):
# read the content from an OSS bucket
# inputs = couler.get_step_output(step_name="A")
couler.run_container(
image="docker/whalesay:latest",
args=["--test 1"],
command=[("cat %s" % default_path)],
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]
artifact = template["outputs"]["artifacts"][0]
self._oss_check_helper(artifact)
template = wf["spec"]["templates"][2]
artifact = template["inputs"]["artifacts"][0]
self._oss_check_helper(artifact)
couler._cleanup()
def test_create_job(self):
success_condition = "status.succeeded > 0"
failure_condition = "status.failed > 3"
# Null manifest
with self.assertRaises(ValueError):
couler.run_job(
manifest=None,
success_condition=success_condition,
failure_condition=failure_condition,
)
# Have a manifest
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"]
"""
for set_owner in (True, False):
couler.run_job(
manifest=manifest,
success_condition=success_condition,
failure_condition=failure_condition,
set_owner_reference=set_owner,
)
self.assertEqual(len(couler.workflow.templates), 1)
template = couler.workflow.get_template(
"test-create-job"
).to_dict()
resource = template["resource"]
self.assertEqual(template["name"], "test-create-job")
self.assertEqual(resource["action"], "create")
self.assertEqual(
resource["setOwnerReference"], "true" if set_owner else "false"
)
self.assertEqual(resource["successCondition"], success_condition)
self.assertEqual(resource["failureCondition"], failure_condition)
self.assertEqual(resource["manifest"], manifest)
couler._cleanup()
def test_run_container_with_image_pull_secret(self):
secret = ImagePullSecret("test-secret")
couler.add_image_pull_secret(secret)
secret1 = ImagePullSecret("test-secret1")
couler.add_image_pull_secret(secret1)
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
working_dir="/mnt/src",
)
wf = couler.workflow_yaml()
self.assertEqual(wf["spec"]["imagePullSecrets"][0], secret.to_dict())
self.assertEqual(wf["spec"]["imagePullSecrets"][1], secret1.to_dict())
couler._cleanup()
def test_run_container_with_volume(self):
volume = Volume("workdir", "my-existing-volume")
volume_mount = VolumeMount("workdir", "/mnt/vol")
couler.add_volume(volume)
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world"],
command=["bash", "-c"],
step_name="A",
volume_mounts=[volume_mount],
)
wf = couler.workflow_yaml()
self.assertEqual(wf["spec"]["volumes"][0], volume.to_dict())
self.assertEqual(
wf["spec"]["templates"][1]["container"]["volumeMounts"][0],
volume_mount.to_dict(),
)
couler._cleanup()
def test_run_container_with_dns(self):
dns_config = DnsConfig(
["10.1.0.1", "0.0.0.0"],
[DnsConfigOption("timeout", "1")],
["domain.google.com"],
)
couler.set_dns("None", dns_config)
wf = couler.workflow_yaml()
self.assertEqual(wf["spec"]["dnsPolicy"], "None")
self.assertEqual(len(wf["spec"]["dnsConfig"]["nameservers"]), 2)
self.assertEqual(wf["spec"]["dnsConfig"]["nameservers"][0], "10.1.0.1")
self.assertEqual(wf["spec"]["dnsConfig"]["nameservers"][1], "0.0.0.0")
self.assertEqual(len(wf["spec"]["dnsConfig"]["options"]), 1)
self.assertEqual(wf["spec"]["dnsConfig"]["options"][0]["name"],
"timeout")
self.assertEqual(wf["spec"]["dnsConfig"]["options"][0]["value"], "1")
self.assertEqual(len(wf["spec"]["dnsConfig"]["searches"]), 1)
self.assertEqual(wf["spec"]["dnsConfig"]["searches"][0],
"domain.google.com")
couler._cleanup()
def test_artifact_passing(self):
def producer():
output_artifact = couler.create_oss_artifact(
path="/mnt/t1.txt",
bucket="test-bucket/",
accesskey_id="abcde",
accesskey_secret="abc12345",
key="osspath/t1",
endpoint="xyz.com",
)
outputs = couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world > %s" % output_artifact.path],
command=["bash", "-c"],
output=output_artifact,
)
return outputs
def consumer(inputs):
# read the content from an OSS bucket
couler.run_container(
image="docker/whalesay:latest",
args=inputs,
command=[("cat %s" % inputs[0].path)],
)
outputs = producer()
consumer(outputs)
wf = couler.workflow_yaml()
template = wf["spec"]["templates"][1]
artifact = template["outputs"]["artifacts"][0]
self.assertEqual(len(template["outputs"]["artifacts"]), 1)
self._oss_check_helper(artifact)
template = wf["spec"]["templates"][2]
artifact = template["inputs"]["artifacts"][0]
self.assertEqual(len(template["inputs"]["artifacts"]), 1)
self._oss_check_helper(artifact)
couler._cleanup()
def test_output_s3_artifact(self):
# the content of local file would be uploaded to OSS
output_artifact = couler.create_s3_artifact(
path="/mnt/t1.txt",
bucket="test-bucket/",
accesskey_id="abcde",
accesskey_secret="abc12345",
key="s3path/t1",
endpoint="xyz.com",
)
couler.run_container(
image="docker/whalesay:latest",
args=["echo -n hello world > %s" % output_artifact.path],
command=["bash", "-c"],
output=output_artifact,
)
wf = couler.workflow_yaml()
template = wf["spec"]["templates"][1]
artifact = template["outputs"]["artifacts"][0]
self._s3_check_helper(artifact)
couler._cleanup()
def test_input_oss_artifact(self):
input_artifact = couler.create_oss_artifact(
path="/mnt/t1.txt",
bucket="test-bucket/",
accesskey_id="abcde",
accesskey_secret="abc12345",
key="osspath/t1",
endpoint="xyz.com",
)
# read the content from an OSS bucket
couler.run_container(
image="docker/whalesay:latest",
args=["cat %s" % input_artifact.path],
command=["bash", "-c"],
input=input_artifact,
)
wf = couler.workflow_yaml()
template = wf["spec"]["templates"][1]
artifact = template["inputs"]["artifacts"][0]
self._oss_check_helper(artifact)
couler._cleanup()
def tearDown(self):
couler._cleanup()
def setUp(self):
couler._cleanup()
self.envs = {"str_env": "abc", "bool_env": False, "num_env": 1234}
def setUp(self) -> None:
couler._cleanup()
