def test_generate_task_definition_uses_run_config_task_definition( self, use_path, monkeypatch): task_definition = { "tags": [{ "key": "mykey", "value": "myvalue" }], "cpu": "2048", "memory": "4096", } if use_path: data = yaml.safe_dump(task_definition) run_config = ECSRun(task_definition_path="s3://test/path.yaml") monkeypatch.setattr( "prefect.agent.ecs.agent.read_bytes_from_path", MagicMock(wraps=read_bytes_from_path, return_value=data), ) else: run_config = ECSRun(task_definition=task_definition) res = self.generate_task_definition(run_config) assert any(e == { "key": "mykey", "value": "myvalue" } for e in res["tags"]) assert res["memory"] == "4096" assert res["cpu"] == "2048"
def test_task_definition_arn(): config = ECSRun(task_definition_arn="my-task-definition") assert config.task_definition_arn == "my-task-definition" assert config.task_definition is None assert config.task_definition_path is None # Can't mix `image` and `task_definition_arn` with pytest.raises(ValueError, match="task_definition_arn"): ECSRun(task_definition_arn="my-task-definition", image="my-image")
def test_cpu_and_memory_acceptable_types(): config = ECSRun() assert config.cpu is None assert config.memory is None config = ECSRun(cpu="1 vcpu", memory="1 GB") assert config.cpu == "1 vcpu" assert config.memory == "1 GB" config = ECSRun(cpu=1024, memory=2048) assert config.cpu == "1024" assert config.memory == "2048"
def test_generate_task_definition_command(self): taskdef = self.generate_task_definition(ECSRun()) assert taskdef["containerDefinitions"][0]["command"] == [ "/bin/sh", "-c", "prefect execute flow-run", ]
def test_local_task_definition_path(tmpdir, scheme): task_definition = { "containerDefinitions": [{ "name": "flow", "environment": [{ "name": "TEST", "value": "VALUE" }] }] } path = str(tmpdir.join("test.yaml")) if scheme is None: task_definition_path = path else: if sys.platform == "win32": pytest.skip("Schemes are not supported on win32") task_definition_path = f"{scheme}://" + path with open(path, "w") as f: yaml.safe_dump(task_definition, f) config = ECSRun(task_definition_path=task_definition_path) assert config.task_definition_path is None assert config.task_definition_arn is None assert config.task_definition == task_definition
def test_generate_task_definition_family_and_tags(self): taskdef = self.generate_task_definition(ECSRun()) assert taskdef["family"] == "prefect-test-flow" assert sorted(taskdef["tags"], key=lambda x: x["key"]) == [ {"key": "prefect:flow-id", "value": "flow-id"}, {"key": "prefect:flow-version", "value": "1"}, ]
def test_get_task_definition_arn_provided_task_definition_arn(): run_config = ECSRun(task_definition_arn="my-taskdef-arn") flow_run = GraphQLResult({"flow": GraphQLResult({"id": "flow-id", "version": 1})}) agent = ECSAgent() res = agent.get_task_definition_arn(flow_run, run_config) assert res == "my-taskdef-arn"
def test_get_task_definition_arn(aws, kind): if kind == "exists": aws.resourcegroupstaggingapi.get_resources.return_value = { "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}] } expected = "my-taskdef-arn" elif kind == "missing": aws.resourcegroupstaggingapi.get_resources.return_value = { "ResourceTagMappingList": [] } expected = None else: from botocore.exceptions import ClientError aws.resourcegroupstaggingapi.get_resources.side_effect = ClientError( {}, "GetResources" ) expected = None run_config = ECSRun() flow_run = GraphQLResult({"flow": GraphQLResult({"id": "flow-id", "version": 1})}) agent = ECSAgent() res = agent.get_task_definition_arn(flow_run, run_config) assert res == expected kwargs = aws.resourcegroupstaggingapi.get_resources.call_args[1] assert sorted(kwargs["TagFilters"], key=lambda x: x["Key"]) == [ {"Key": "prefect:flow-id", "Values": ["flow-id"]}, {"Key": "prefect:flow-version", "Values": ["1"]}, ] assert kwargs["ResourceTypeFilters"] == ["ecs:task-definition"]
def test_local_task_definition_path(tmpdir, scheme): task_definition = { "containerDefinitions": [{ "name": "flow", "environment": [{ "name": "TEST", "value": "VALUE" }] }] } path = str(tmpdir.join("test.yaml")) if scheme is None: task_definition_path = path else: # With a scheme, unix-style slashes are required task_definition_path = f"{scheme}://" + os.path.splitdrive( path)[1].replace("\\", "/") with open(path, "w") as f: yaml.safe_dump(task_definition, f) config = ECSRun(task_definition_path=task_definition_path) assert config.task_definition_path is None assert config.task_definition_arn is None assert config.task_definition == task_definition
def test_environment_has_agent_token_from_config(self): with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): env_list = self.get_run_task_kwargs( ECSRun())["overrides"]["containerOverrides"][0]["environment"] env = {item["name"]: item["value"] for item in env_list} assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_generate_task_definition_environment(self): run_config = ECSRun( image="test-image", task_definition={ "containerDefinitions": [{ "name": "flow", "environment": [ { "name": "CUSTOM1", "value": "VALUE1" }, { "name": "CUSTOM2", "value": "VALUE2" }, ], }] }, env={"CUSTOM4": "VALUE4"}, ) taskdef = self.generate_task_definition(run_config, env_vars={"CUSTOM3": "VALUE3"}) env_list = taskdef["containerDefinitions"][0]["environment"] env = {item["name"]: item["value"] for item in env_list} # Agent and run-config level envs are only set at runtime assert env == { "PREFECT__CONTEXT__IMAGE": "test-image", "CUSTOM1": "VALUE1", "CUSTOM2": "VALUE2", }
def test_get_run_task_kwargs_command(self): kwargs = self.get_run_task_kwargs(ECSRun()) assert kwargs["overrides"]["containerOverrides"][0]["command"] == [ "/bin/sh", "-c", "prefect execute flow-run", ]
def test_get_task_run_kwargs_execution_role_arn( self, on_run_config, on_agent, expected ): kwargs = self.get_run_task_kwargs( ECSRun(execution_role_arn=on_run_config), execution_role_arn=on_agent ) assert kwargs["overrides"].get("executionRoleArn") == expected
def test_generate_task_definition_execution_role_arn( self, on_run_config, on_agent, expected ): taskdef = self.generate_task_definition( ECSRun(execution_role_arn=on_run_config), execution_role_arn=on_agent ) assert taskdef.get("executionRoleArn") == expected
def test_generate_task_definition_environment(self): run_config = ECSRun( image="test-image", task_definition={ "containerDefinitions": [ { "name": "flow", "environment": [ {"name": "CUSTOM1", "value": "VALUE1"}, {"name": "CUSTOM2", "value": "VALUE2"}, ], } ] }, env={"CUSTOM4": "VALUE4"}, ) taskdef = self.generate_task_definition( run_config, env_vars={"CUSTOM3": "VALUE3"} ) env_list = taskdef["containerDefinitions"][0]["environment"] env = {item["name"]: item["value"] for item in env_list} # Agent and run-config level envs are only set at runtime assert env == { "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", "PREFECT__CONTEXT__IMAGE": "test-image", "CUSTOM1": "VALUE1", "CUSTOM2": "VALUE2", }
def test_get_run_task_kwargs_environment(self, tmpdir, backend): path = str(tmpdir.join("kwargs.yaml")) with open(path, "w") as f: yaml.safe_dump( { "overrides": { "containerOverrides": [{ "name": "flow", "environment": [ { "name": "CUSTOM1", "value": "VALUE1" }, { "name": "CUSTOM2", "value": "VALUE2" }, ], }] } }, f, ) kwargs = self.get_run_task_kwargs( ECSRun(env={ "CUSTOM3": "OVERRIDE3", "CUSTOM4": "VALUE4" }), env_vars={ "CUSTOM2": "OVERRIDE2", "CUSTOM3": "VALUE3" }, run_task_kwargs_path=path, ) env_list = kwargs["overrides"]["containerOverrides"][0]["environment"] env = {item["name"]: item["value"] for item in env_list} assert env == { "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", "PREFECT__BACKEND": backend, "PREFECT__CLOUD__API": prefect.config.cloud.api, "PREFECT__CLOUD__AUTH_TOKEN": "TEST_TOKEN", "PREFECT__CLOUD__AGENT__LABELS": "[]", "PREFECT__CONTEXT__FLOW_RUN_ID": "flow-run-id", "PREFECT__CONTEXT__FLOW_ID": "flow-id", "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS": "true", "PREFECT__LOGGING__LOG_TO_CLOUD": "true", "PREFECT__LOGGING__LEVEL": prefect.config.logging.level, "CUSTOM1": "VALUE1", "CUSTOM2": "OVERRIDE2", # agent envs override agent run-task-kwargs "CUSTOM3": "OVERRIDE3", # run-config envs override agent "CUSTOM4": "VALUE4", }
def test_get_run_task_kwargs_common(self, launch_type, cluster): kwargs = self.get_run_task_kwargs(ECSRun(), launch_type=launch_type, cluster=cluster) assert kwargs["launchType"] == launch_type assert kwargs.get("cluster") == cluster assert ("networkConfiguration" in kwargs) == (launch_type == "FARGATE") assert kwargs["overrides"]["containerOverrides"][0]["name"] == "flow"
def test_get_task_run_kwargs_capacity_provider_run_config( self, on_run_config, on_agent, expected_run_config, expected_agent ): kwargs = self.get_run_task_kwargs( ECSRun(run_task_kwargs=on_run_config), launch_type=on_agent ) assert kwargs.get("capacityProviderStrategy") == expected_run_config assert kwargs.get("launchType") == expected_agent
def test_deploy_flow_uses_provided_task_definition_arn(self, aws): aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]} res = self.deploy_flow(ECSRun(task_definition_arn="my-taskdef-arn")) assert not aws.ecs.register_task_definition.called assert aws.ecs.run_task.called assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn" assert "my-task-arn" in res
def test_deploy_flow_errors_if_mix_task_definition_arn_and_docker_storage(self): with pytest.raises( ValueError, match="Cannot provide `task_definition_arn` when using `Docker` storage", ): self.deploy_flow( ECSRun(task_definition_arn="my-taskdef-arn"), storage=Docker(registry_url="test", image_name="name", image_tag="tag"), )
def test_environment_has_api_key_from_config(self, config_with_api_key): env_list = self.get_run_task_kwargs(ECSRun())["overrides"][ "containerOverrides" ][0]["environment"] env = {item["name"]: item["value"] for item in env_list} assert env["PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key assert env["PREFECT__CLOUD__AUTH_TOKEN"] == config_with_api_key.cloud.api_key assert env["PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def configure_run_config(cluster: Cluster, recipe_bakery: RecipeBakery, recipe_name: str, secrets: Dict): if cluster.type == FARGATE_CLUSTER: definition = { "networkMode": "awsvpc", "cpu": 2048, "memory": 16384, "containerDefinitions": [{ "name": "flow" }], "executionRoleArn": cluster.cluster_options.execution_role_arn, } run_config = ECSRun( image=cluster.worker_image, labels=[recipe_bakery.id], task_definition=definition, run_task_kwargs={ "tags": [ { "key": "Project", "value": "pangeo-forge" }, { "key": "Recipe", "value": recipe_name }, ] }, ) return run_config elif cluster.type == AKS_CLUSTER: job_template = yaml.safe_load(""" apiVersion: batch/v1 kind: Job metadata: annotations: "cluster-autoscaler.kubernetes.io/safe-to-evict": "false" spec: template: spec: containers: - name: flow """) run_config = KubernetesRun( job_template=job_template, image=cluster.worker_image, labels=[recipe_bakery.id], memory_request="10000Mi", cpu_request="2048m", env={ "AZURE_STORAGE_CONNECTION_STRING": secrets[cluster.flow_storage_options.secret] }, ) return run_config else: raise UnsupportedClusterType
def test_prefect_logging_level_override_logic(self, config, agent_env_vars, run_config_env_vars, expected_logging_level): with set_temporary_config(config): kwargs = self.get_run_task_kwargs(ECSRun(env=run_config_env_vars), env_vars=agent_env_vars) env_list = kwargs["overrides"]["containerOverrides"][0][ "environment"] env = {item["name"]: item["value"] for item in env_list} assert env["PREFECT__LOGGING__LEVEL"] == expected_logging_level
def test_deploy_flow_run_task_fails(self, aws): aws.ecs.run_task.return_value = { "tasks": [], "failures": [{"reason": "my-reason"}], } with pytest.raises(ValueError) as exc: self.deploy_flow(ECSRun()) assert aws.ecs.run_task.called assert aws.ecs.deregister_task_definition.called assert "my-reason" in str(exc.value)
def test_generate_task_definition_requires_compatibilities_capacity_provider( self, tmpdir ): path = str(tmpdir.join("kwargs.yaml")) with open(path, "w") as f: yaml.safe_dump( {"capacityProviderStrategy": [{"capacityProvider", "FARGATE_SPOT"}]}, f ) taskdef = self.generate_task_definition(ECSRun(), run_task_kwargs_path=path) assert taskdef.get("requiresCompatibilities") == None
def test_deploy_flow_forwards_run_task_kwargs(self, aws): aws.ecs.register_task_definition.return_value = { "taskDefinition": {"taskDefinitionArn": "my-taskdef-arn"} } aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]} res = self.deploy_flow(ECSRun(run_task_kwargs={"enableECSManagedTags": True})) assert aws.ecs.run_task.called assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn" assert aws.ecs.run_task.call_args[1]["enableECSManagedTags"] is True assert "my-task-arn" in res
def test_no_args(): config = ECSRun() assert config.task_definition is None assert config.task_definition_path is None assert config.image is None assert config.env is None assert config.cpu is None assert config.memory is None assert config.task_role_arn is None assert config.run_task_kwargs is None assert config.labels == set()
def test_deploy_flow_forwards_run_config_settings(self, aws): aws.ecs.register_task_definition.return_value = { "taskDefinition": {"taskDefinitionArn": "my-taskdef-arn"} } aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]} self.deploy_flow(ECSRun(cpu=8, memory=1024)) aws.ecs.run_task.assert_called_once() assert aws.ecs.run_task.call_args[1]["overrides"]["cpu"] == "8" assert aws.ecs.run_task.call_args[1]["overrides"]["memory"] == "1024"
def test_environment_has_api_key_from_config(self, tenant_id): with set_temporary_config({ "cloud.api_key": "TEST_KEY", "cloud.tenant_id": tenant_id }): env_list = self.get_run_task_kwargs( ECSRun())["overrides"]["containerOverrides"][0]["environment"] env = {item["name"]: item["value"] for item in env_list} assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY" assert env.get("PREFECT__CLOUD__TENANT_ID") == tenant_id
def test_task_definition(): task_definition = { "containerDefinitions": [ {"name": "flow", "environment": [{"name": "TEST", "value": "VALUE"}]} ] } config = ECSRun(task_definition=task_definition) assert config.task_definition_path is None assert config.task_definition_arn is None assert config.task_definition == task_definition