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
