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()