def test_set_dependencies_none(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: job_c(message="C"), dependencies=None) pyaml.dump(couler.workflow_yaml()) wf_tree = couler.workflow_yaml() tasks = wf_tree["spec"]["templates"][0]["dag"]["tasks"] assert_deps = {"A": None, "B": ["A"], "C": None} self.assertEqual(len(tasks), 3) for task in tasks: assert_dep = assert_deps[task["name"]] self.assertEqual(task.get("dependencies"), assert_dep)
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_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 test_run_container_env(self): self.run_a_container() workflow = couler.workflow_yaml() checked = False for template in workflow["spec"]["templates"]: if template["name"] == "run-a-container": checked = True container = template["container"] self.assertEqual(len(container), 3) self.assertEqual(container["image"], "python:3.6") self.assertEqual(container["command"], ["bash", "-c", "echo ${MESSAGE}"]) self.assertEqual(len(container["env"]), len(self.envs) + 2) for env in container["env"]: if env["name"] == "NVIDIA_VISIBLE_DEVICES": self.assertEqual(env["value"], "") elif env["name"] == "NVIDIA_DRIVER_CAPABILITIES": self.assertEqual(env["value"], "") elif env["name"] == "str_env": self.assertEqual(env["value"], self.envs["str_env"]) elif env["name"] == "bool_env": self.assertEqual(env["value"], "'%s'" % self.envs["bool_env"]) elif env["name"] == "num_env": self.assertEqual(env["value"], "%s" % self.envs["num_env"]) self.assertTrue(checked)
def test_run_script_env(self): self.run_a_script() workflow = couler.workflow_yaml() checked = False for template in workflow["spec"]["templates"]: if template["name"] == "run-a-script": checked = True script = template["script"] self.assertEqual(len(script), 4) self.assertEqual(script["image"], "python:3.6") self.assertEqual(script["command"], ["python"]) self.assertEqual(len(script["env"]), len(self.envs) + 2) for env in script["env"]: if env["name"] == "NVIDIA_VISIBLE_DEVICES": self.assertEqual(env["value"], "") elif env["name"] == "NVIDIA_DRIVER_CAPABILITIES": self.assertEqual(env["value"], "") elif env["name"] == "str_env": self.assertEqual(env["value"], self.envs["str_env"]) elif env["name"] == "bool_env": self.assertEqual(env["value"], "'%s'" % self.envs["bool_env"]) elif env["name"] == "num_env": self.assertEqual(env["value"], "%s" % self.envs["num_env"]) self.assertTrue(checked)
def test_run_gpu_script_env(self): self.run_a_gpu_script() workflow = couler.workflow_yaml() checked = False for template in workflow["spec"]["templates"]: if template["name"] == "run-a-gpu-script": checked = True script = template["script"] self.assertEqual(len(script), 5) self.assertEqual(script["image"], "python:3.6") self.assertEqual(script["command"], ["python"]) self.assertEqual(len(script["env"]), len(self.envs)) for env in script["env"]: if env["name"] not in self.envs: raise AssertionError("Unrecognized env variable %s" % env["name"]) elif env["name"] == "str_env": self.assertEqual(env["value"], self.envs["str_env"]) elif env["name"] == "bool_env": self.assertEqual(env["value"], "'%s'" % self.envs["bool_env"]) elif env["name"] == "num_env": self.assertEqual(env["value"], "%s" % self.envs["num_env"]) self.assertTrue(checked)
def test_workflow_name_salt_applies(self): couler.config_defaults(name_salter=lambda name: "%s-salted" % name) couler.config_workflow(name="test_name_salter_applies") ret = couler.workflow_yaml() self.assertEqual( ret["metadata"]["name"], "test_name_salter_applies-salted" )
def test_run_gpu_container_env(self): self.run_a_gpu_container() workflow = couler.workflow_yaml() checked = False for template in workflow["spec"]["templates"]: if template["name"] == "run-a-gpu-container": checked = True container = template["container"] self.assertEqual(len(container), 4) self.assertEqual(container["image"], "python:3.6") self.assertEqual(container["command"], ["bash", "-c", "echo ${MESSAGE}"]) self.assertEqual(len(container["env"]), len(self.envs)) for env in container["env"]: if env["name"] not in self.envs: raise AssertionError("Unrecognized env variable %s" % env["name"]) elif env["name"] == "str_env": self.assertEqual(env["value"], self.envs["str_env"]) elif env["name"] == "bool_env": self.assertEqual(env["value"], "'%s'" % self.envs["bool_env"]) elif env["name"] == "num_env": self.assertEqual(env["value"], "%s" % self.envs["num_env"]) self.assertTrue(checked)
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_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_workflow_config(self): flip_coin() tails() couler.config_workflow( name="test-workflow", user_id="88888888", time_to_clean=100 ) wf = couler.workflow_yaml() expected_meta = { "name": "test-workflow", # "labels": {"couler_job_user": "******"}, } self.assertEqual(wf["metadata"], expected_meta)
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 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_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_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_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_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 check_argo_yaml(self, expected_fn): test_data_dir = os.path.join(os.path.dirname(__file__), _test_data_dir) with open(os.path.join(test_data_dir, expected_fn), "r") as f: expected = yaml.safe_load(f) output = yaml.safe_load( pyaml.dump(couler.workflow_yaml(), string_val_style="plain")) def dump(x): x = re.sub(r"-[0-9]*", "-***", json.dumps(self.mock_dict(x), indent=2)) return x output_j, expected_j = dump(output), dump(expected) self.maxDiff = None self.assertEqual(output_j, expected_j)
def test_input_args_as_othertypes(self): def whalesay(para_integer, para_boolean, para_float): return couler.run_container( image="docker/whalesay", command=["cowsay"], args=[para_integer, para_boolean, para_float], ) whalesay(1, True, 1.1) wf = couler.workflow_yaml() parameters = wf["spec"]["templates"][0]["steps"][0][0]["arguments"][ "parameters"] self.assertEqual(int(parameters[0]["value"].strip(" " "' ")), 1) self.assertEqual(bool(parameters[1]["value"].strip(" " "' ")), True) self.assertEqual(float(parameters[2]["value"].strip(" " "' ")), 1.1)
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 verify_concurrent_step(self, yaml_file_name): _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, yaml_file_name), "r") as f: expected = yaml.safe_load(f) output = yaml.safe_load( pyaml.dump(couler.workflow_yaml(), string_val_style="plain")) # Because test environment between local and CI is different, # we can not compare the YAML directly. steps = output["spec"]["templates"][0]["steps"][0] expected_steps = expected["spec"]["templates"][0]["steps"][0] self.assertEqual(len(steps), len(expected_steps)) for index in range(len(steps)): _step = steps[index] _expected_step = expected_steps[index] self.assertEqual(_step["template"], _expected_step["template"])
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_artifact_passing_script(self): def producer(): output_artifact = couler.create_local_artifact(path="/mnt/t1.txt") outputs = couler.run_script( image="docker/whalesay:latest", args=["echo -n hello world > %s" % output_artifact.path], command=["bash", "-c"], output=output_artifact, source="sadfa\nasdf", step_name="producer", ) return outputs def consumer(inputs): # read the content from local artifact couler.run_script( image="docker/whalesay:latest", args=inputs, command=[("cat %s" % inputs[0].path)], source="sadfa\nasdf", ) outputs = couler.set_dependencies(lambda: producer(), dependencies=None) couler.set_dependencies(lambda: consumer(outputs), dependencies=["producer"]) wf = couler.workflow_yaml() dag = wf["spec"]["templates"][0]["dag"] self.assertEqual(len(dag["tasks"][1]["arguments"]["artifacts"]), 1) self.assertIsInstance( dag["tasks"][1]["arguments"]["artifacts"][0]["from"], str) self.assertIsInstance( dag["tasks"][1]["arguments"]["artifacts"][0]["name"], str) template = wf["spec"]["templates"][1] self.assertEqual(len(template["outputs"]["artifacts"]), 1) template = wf["spec"]["templates"][2] self.assertEqual(len(template["inputs"]["artifacts"]), 1) self.assertNotIn("args", template)
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_katib_with_xgboost_training(self): katib.run( tuning_params=[ { "name": "max_depth", "type": "int", "range": [2, 10] }, { "name": "num_round", "type": "int", "range": [50, 100] }, ], objective={ "type": "maximize", "goal": 1.01, "metric_name": "accuracy", }, success_condition="status.trialsSucceeded > 4", failure_condition="status.trialsFailed > 3", algorithm="random", raw_template=xgboost_mainifest_template, parallel_trial_count=4, max_trial_count=16, max_failed_trial_count=3, ) wf = couler.workflow_yaml() self.assertEqual(len(wf["spec"]["templates"]), 2) # Check steps template template0 = wf["spec"]["templates"][0] self.assertEqual(len(template0["steps"]), 1) self.assertEqual(len(template0["steps"][0]), 1) # Check train template template1 = wf["spec"]["templates"][1] self.assertTrue(template1["name"] in ["run", "test-run-python-script"])