def test_parse_task_run_kwargs(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
kwarg_dict = {
"cluster": "test",
"count": "test",
"startedBy": "test",
"group": "test",
"placementConstraints": "test",
"placementStrategy": "test",
"platformVersion": "test",
"networkConfiguration": "test",
"tags": "test",
"enableECSManagedTags": "test",
"propagateTags": "test",
}
task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict)
assert task_run_kwargs == kwarg_dict
assert task_definition_kwargs == {
"placementConstraints": "test",
"tags": "test"
}
def test_fargate_agent_start_max_polls_zero(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
on_shutdown = MagicMock()
monkeypatch.setattr("prefect.agent.fargate.agent.FargateAgent.on_shutdown",
on_shutdown)
agent_process = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process",
agent_process)
agent_connect = MagicMock(return_value="id")
monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect",
agent_connect)
heartbeat = MagicMock()
monkeypatch.setattr("prefect.agent.fargate.agent.FargateAgent.heartbeat",
heartbeat)
agent = FargateAgent(max_polls=0)
agent.start()
assert on_shutdown.call_count == 1
assert agent_process.call_count == 0
assert heartbeat.call_count == 0
def test_deploy_flows_register_task_definition(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.run_task.return_value = {}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
})
])
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert (boto3_client.register_task_definition.call_args[1]["family"] ==
"prefect-task-id")
def test_deploy_flows_require_docker_storage(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {"tags": []}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
with pytest.raises(Exception) as excinfo:
agent = FargateAgent()
agent.deploy_flow(
GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"version": 2,
"name": "name",
}),
"id":
"id",
"name":
"name",
}))
assert boto3_client.describe_task_definition.not_called
assert boto3_client.run_task.not_called
def test_deploy_flow_raises(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
"name":
"name",
}))
assert boto3_client.describe_task_definition.called
assert boto3_client.run_task.called
def test_deploy_flows_no_security_group(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
})
])
assert boto3_client.describe_task_definition.called
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["cluster"] == "default"
assert boto3_client.run_task.call_args[1]["networkConfiguration"] == {
"awsvpcConfiguration": {
"subnets": [],
"assignPublicIp": "ENABLED"
}
}
def test_deploy_flow_register_task_definition_no_repo_credentials(
monkeypatch, runner_token
):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
with set_temporary_config({"logging.log_to_cloud": True}):
agent = FargateAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"id": "id",
}
),
"id": "id",
}
)
)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"
] == [
{
"name": "flow",
"image": "test/name:tag",
"command": ["/bin/sh", "-c", "prefect execute cloud-flow"],
"environment": [
{"name": "PREFECT__CLOUD__API", "value": "https://api.prefect.io"},
{"name": "PREFECT__CLOUD__AGENT__LABELS", "value": "[]"},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{"name": "PREFECT__LOGGING__LEVEL", "value": "DEBUG"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
],
"essential": True,
}
]
def test_parse_task_kwargs_invalid_value_removed(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
kwarg_dict = {"test": "not_real"}
task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict)
assert task_definition_kwargs == {}
assert task_run_kwargs == {}
def test_deploy_flows_enable_task_revisions_old_version_exists(
monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {
"tags": [
{
"key": "PrefectFlowId",
"value": "current_id"
},
{
"key": "PrefectFlowVersion",
"value": "5"
},
]
}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.get_resources.return_value = {
"ResourceTagMappingList": [{
"ResourceARN":
"arn:aws:ecs:us-east-1:12345:task-definition/flow:22"
}]
}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(enable_task_revisions=True)
agent.deploy_flow(
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
"version":
3,
"name":
"name",
}),
"id":
"id",
"name":
"name",
}))
assert boto3_client.describe_task_definition.called
assert boto3_client.get_resources.called
assert boto3_client.register_task_definition.not_called
assert boto3_client.run_task.called
assert (agent.task_definition_name ==
"arn:aws:ecs:us-east-1:12345:task-definition/flow:22")
def test_override_kwargs(monkeypatch, runner_token):
boto3_resource = MagicMock()
streaming_body = MagicMock()
streaming_body.read.return_value.decode.return_value = (
'{"cpu": "256", "networkConfiguration": "test"}')
boto3_resource.return_value.Object.return_value.get.return_value = {
"Body": streaming_body
}
monkeypatch.setattr("boto3.resource", boto3_resource)
agent = FargateAgent(
use_external_kwargs=True,
external_kwargs_s3_bucket="test-bucket",
external_kwargs_s3_key="prefect-artifacts/kwargs",
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
region_name="region",
labels=["aws", "staging"],
)
definition_kwargs = {}
run_kwargs = {}
agent._override_kwargs(
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
"version":
2,
"name":
"name",
}),
"id":
"id",
"name":
"name",
}),
definition_kwargs,
run_kwargs,
)
assert boto3_resource.called
assert streaming_body.read().decode.called
assert definition_kwargs == {"cpu": "256"}
assert run_kwargs == {"networkConfiguration": "test"}
def test_ecs_agent_config_env_vars(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "token")
monkeypatch.setenv("TASK_ROLE_ARN", "task_role_arn")
monkeypatch.setenv("EXECUTION_ROLE_ARN", "execution_role_arn")
monkeypatch.setenv("REGION_NAME", "region")
monkeypatch.setenv("CLUSTER", "cluster")
monkeypatch.setenv("REPOSITORY_CREDENTIALS", "repo")
monkeypatch.setenv("ASSIGN_PUBLIC_IP", "DISABLED")
monkeypatch.setenv("TASK_CPU", "1")
monkeypatch.setenv("TASK_MEMORY", "2")
agent = FargateAgent(subnets=["subnet"])
assert agent
assert agent.execution_role_arn == "execution_role_arn"
assert agent.task_role_arn == "task_role_arn"
assert agent.cluster == "cluster"
assert agent.repository_credentials == "repo"
assert agent.assign_public_ip == "DISABLED"
assert agent.task_cpu == "1"
assert agent.task_memory == "2"
boto3_client.assert_called_with(
"ecs",
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
region_name="region",
)
def test_fargate_agent_config_options(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
# Client args
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "")
monkeypatch.setenv("AWS_SESSION_TOKEN", "")
monkeypatch.setenv("REGION_NAME", "")
monkeypatch.delenv("AWS_ACCESS_KEY_ID")
monkeypatch.delenv("AWS_SECRET_ACCESS_KEY")
monkeypatch.delenv("AWS_SESSION_TOKEN")
monkeypatch.delenv("REGION_NAME")
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = FargateAgent(name="test", labels=["test"])
assert agent
assert agent.labels == ["test"]
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert agent.boto3_client
boto3_client.assert_called_with(
"ecs",
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
region_name=None,
)
def test_ecs_agent_init(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
assert agent
assert agent.boto3_client
def test_ecs_agent_config_options_init(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent(
aws_access_key_id="id",
aws_secret_access_key="secret",
region_name="region",
cluster="cluster",
subnets=["subnet"],
security_groups=["security_group"],
repository_credentials="repo",
assign_public_ip="DISABLED",
task_cpu="1",
task_memory="2",
)
assert agent
assert agent.cluster == "cluster"
assert agent.subnets == ["subnet"]
assert agent.security_groups == ["security_group"]
assert agent.repository_credentials == "repo"
assert agent.assign_public_ip == "DISABLED"
assert agent.task_cpu == "1"
assert agent.task_memory == "2"
boto3_client.assert_called_with(
"ecs",
aws_access_key_id="id",
aws_secret_access_key="secret",
region_name="region",
)
def test_parse_task_definition_kwargs_errors(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
kwarg_dict = {
"placementConstraints": "taskRoleArn='arn:aws:iam::543216789012:role/Dev",
}
task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict)
assert task_definition_kwargs == kwarg_dict
assert task_run_kwargs == {
"placementConstraints": "taskRoleArn='arn:aws:iam::543216789012:role/Dev"
}
def test_deploy_flows_enable_task_revisions_tags_passed_in(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {
"tags": [{"key": "PrefectFlowId", "value": "id"}]
}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(
enable_task_revisions=True,
tags=[
{"key": "PrefectFlowId", "value": "id"},
{"key": "PrefectFlowVersion", "value": "2"},
],
)
agent.deploy_flow(
GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"id": "id",
"version": 2,
"name": "name",
}
),
"id": "id",
"name": "name",
}
)
)
assert agent.task_definition_kwargs == {
"tags": [
{"key": "PrefectFlowId", "value": "id"},
{"key": "PrefectFlowVersion", "value": "2"},
]
}
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.not_called
assert boto3_client.run_task.called
assert boto3_client.run_task.called_with(taskDefinition="name")
def test_deploy_flow_register_task_definition_uses_user_env_vars(
monkeypatch, runner_token
):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar"))
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"id": "id",
}
),
"id": "id",
}
)
)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert (
boto3_client.register_task_definition.call_args[1]["family"]
== "prefect-task-id"
)
container_defs = boto3_client.register_task_definition.call_args[1][
"containerDefinitions"
]
user_vars = [
dict(name="AUTH_THING", value="foo"),
dict(name="PKG_SETTING", value="bar"),
]
assert container_defs[0]["environment"][-1] in user_vars
assert container_defs[0]["environment"][-2] in user_vars
def test_fargate_agent_config_options_default(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
assert agent
assert agent.labels == []
assert agent.name == "agent"
assert agent.task_definition_kwargs == {}
assert agent.task_run_kwargs == {}
assert agent.boto3_client
def test_default_subnets(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_subnets.return_value = {
"Subnets": [
{"MapPublicIpOnLaunch": False, "SubnetId": "id"},
{"MapPublicIpOnLaunch": True, "SubnetId": "id2"},
]
}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
assert agent.subnets == ["id"]
def test_k8s_agent_config_options(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = FargateAgent(name="test", labels=["test"])
assert agent
assert agent.labels == ["test"]
assert agent.name == "test"
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
assert agent.boto3_client
def test_deploy_flow_raises(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult(
{
"flow": GraphQLResult({
"storage": Local().serialize(),
"id": "id"
}),
"id": "id",
}))
assert not boto3_client.describe_task_definition.called
assert not boto3_client.run_task.called
def test_ecs_agent_config_options_default(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
assert agent
assert agent.cluster == "default"
assert not agent.subnets
assert not agent.security_groups
assert not agent.repository_credentials
assert agent.assign_public_ip == "ENABLED"
assert agent.task_cpu == "256"
assert agent.task_memory == "512"
assert agent.boto3_client
def test_parse_task_definition_kwargs(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
agent = FargateAgent()
kwarg_dict = {
"taskRoleArn": "test",
"executionRoleArn": "test",
"volumes": "test",
"placementConstraints": "test",
"cpu": "test",
"memory": "test",
"tags": "test",
"pidMode": "test",
"ipcMode": "test",
"proxyConfiguration": "test",
"inferenceAccelerators": "test",
}
task_definition_kwargs, task_run_kwargs = agent._parse_kwargs(kwarg_dict)
assert task_definition_kwargs == kwarg_dict
assert task_run_kwargs == {"placementConstraints": "test", "tags": "test"}
def test_fargate_agent_config_env_vars_lists_dicts(monkeypatch, runner_token):
boto3_client = MagicMock()
monkeypatch.setattr("boto3.client", boto3_client)
def_kwarg_dict = {
"placementConstraints": ["test"],
"proxyConfiguration": {
"test": "test"
},
}
run_kwarg_dict = {
"placementConstraints": ["test"],
"networkConfiguration": {
"test": "test"
},
}
# Client args
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "token")
monkeypatch.setenv("REGION_NAME", "region")
# Def / run args
monkeypatch.setenv("placementConstraints", "['test']")
monkeypatch.setenv("proxyConfiguration", "{'test': 'test'}")
monkeypatch.setenv("networkConfiguration", "{'test': 'test'}")
agent = FargateAgent(subnets=["subnet"])
assert agent
assert agent.task_definition_kwargs == def_kwarg_dict
assert agent.task_run_kwargs == run_kwarg_dict
boto3_client.assert_called_with(
"ecs",
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
region_name="region",
)
def get_fargate_agent_definition():
"""
Get the Prefect Agent definition for an environment that run workflows on AWS ECS Fargate
Returns:
[FargateAgent] -- Fargate Agent object from prefect.agent.fargate
"""
subnets_list = subnets.split("|")
return FargateAgent(
region_name=aws_region,
cpu=agent_cpu,
memory=agent_memory,
cluster=cluster_name,
taskRoleArn=task_role_arn,
executionRoleArn=execution_role_arn,
networkConfiguration={
"awsvpcConfiguration": {
"assignPublicIp": "ENABLED",
"subnets": subnets_list,
"securityGroups": [],
}
},
containerDefinitions=[{
"logConfiguration": {
"logDriver": "awslogs",
"options": {
"awslogs-region": aws_region,
"awslogs-group":
f"{environment}_dataflow_automation_agent",
"awslogs-stream-prefix": "workflow_start",
},
},
}],
labels=[f"{environment}_dataflow_automation"],
)
def test_deploy_flows_register_task_definition_all_args(
monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.run_task.return_value = {}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
task_role_arn="task_role_arn",
execution_role_arn="execution_role_arn",
region_name="region",
cluster="cluster",
subnets=["subnet"],
security_groups=["security_group"],
repository_credentials="repo",
assign_public_ip="DISABLED",
task_cpu="1",
task_memory="2",
)
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
})
])
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert (boto3_client.register_task_definition.call_args[1]["family"] ==
"prefect-task-id")
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"] == [{
"name":
"flow",
"image":
"test/name:tag",
"command": ["/bin/sh", "-c", "prefect execute cloud-flow"],
"environment": [
{
"name": "PREFECT__CLOUD__API",
"value": "https://api.prefect.io"
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
{
"name": "PREFECT__LOGGING__LEVEL",
"value": "DEBUG"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
],
"essential":
True,
"repositoryCredentials": {
"credentialsParameter": "repo"
},
}]
assert boto3_client.register_task_definition.call_args[1][
"requiresCompatibilities"] == ["FARGATE"]
assert boto3_client.register_task_definition.call_args[1][
"networkMode"] == "awsvpc"
assert boto3_client.register_task_definition.call_args[1]["cpu"] == "1"
assert boto3_client.register_task_definition.call_args[1]["memory"] == "2"
def test_deploy_flow_all_args(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
kwarg_dict = {
"taskRoleArn": "test",
"executionRoleArn": "test",
"volumes": "test",
"placementConstraints": "test",
"cpu": "test",
"memory": "test",
"tags": "test",
"pidMode": "test",
"ipcMode": "test",
"proxyConfiguration": "test",
"inferenceAccelerators": "test",
"cluster": "cluster",
"count": "test",
"startedBy": "test",
"group": "test",
"placementStrategy": "test",
"platformVersion": "test",
"networkConfiguration": {
"awsvpcConfiguration": {
"subnets": ["subnet"],
"assignPublicIp": "DISABLED",
"securityGroups": ["security_group"],
}
},
"enableECSManagedTags": "test",
"propagateTags": "test",
}
agent = FargateAgent(aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
region_name="region",
**kwarg_dict)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
}))
assert boto3_client.describe_task_definition.called
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["cluster"] == "cluster"
assert boto3_client.run_task.call_args[1][
"taskDefinition"] == "prefect-task-id"
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": ""
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
],
}]
}
assert boto3_client.run_task.call_args[1]["networkConfiguration"] == {
"awsvpcConfiguration": {
"subnets": ["subnet"],
"assignPublicIp": "DISABLED",
"securityGroups": ["security_group"],
}
}
def test_deploy_flows_all_args(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(
aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
task_role_arn="task_role_arn",
execution_role_arn="execution_role_arn",
region_name="region",
cluster="cluster",
subnets=["subnet"],
security_groups=["security_group"],
repository_credentials="repo",
assign_public_ip="DISABLED",
task_cpu="1",
task_memory="2",
)
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
})
])
assert boto3_client.describe_task_definition.called
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["cluster"] == "cluster"
assert boto3_client.run_task.call_args[1][
"taskDefinition"] == "prefect-task-id"
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": ""
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
],
}]
}
assert boto3_client.run_task.call_args[1]["networkConfiguration"] == {
"awsvpcConfiguration": {
"subnets": ["subnet"],
"assignPublicIp": "DISABLED",
"securityGroups": ["security_group"],
}
}
def test_deploy_flows_includes_agent_labels_in_environment(
monkeypatch, runner_token, flag):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
kwarg_dict = {
"taskRoleArn": "test",
"executionRoleArn": "test",
"volumes": "test",
"placementConstraints": "test",
"cpu": "1",
"memory": "2",
"tags": "test",
"pidMode": "test",
"ipcMode": "test",
"proxyConfiguration": "test",
"inferenceAccelerators": "test",
"cluster": "cluster",
"count": "test",
"startedBy": "test",
"group": "test",
"placementStrategy": "test",
"platformVersion": "test",
"networkConfiguration": {
"awsvpcConfiguration": {
"subnets": ["subnet"],
"assignPublicIp": "DISABLED",
"securityGroups": ["security_group"],
}
},
"enableECSManagedTags": "test",
"propagateTags": "test",
}
with set_temporary_config({"logging.log_to_cloud": flag}):
agent = FargateAgent(aws_access_key_id="id",
aws_secret_access_key="secret",
aws_session_token="token",
region_name="region",
labels=["aws", "staging"],
**kwarg_dict)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
}),
"id":
"id",
}))
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert (boto3_client.register_task_definition.call_args[1]["family"] ==
"prefect-task-id")
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"] == [{
"name":
"flow",
"image":
"test/name:tag",
"command": ["/bin/sh", "-c", "prefect execute cloud-flow"],
"environment": [
{
"name": "PREFECT__CLOUD__API",
"value": "https://api.prefect.io"
},
{
"name": "PREFECT__CLOUD__AGENT__LABELS",
"value": "['aws', 'staging']",
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": str(flag).lower()
},
{
"name": "PREFECT__LOGGING__LEVEL",
"value": "DEBUG"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
],
"essential":
True,
}]
assert boto3_client.register_task_definition.call_args[1][
"requiresCompatibilities"] == ["FARGATE"]
assert boto3_client.register_task_definition.call_args[1][
"networkMode"] == "awsvpc"
assert boto3_client.register_task_definition.call_args[1]["cpu"] == "1"
assert boto3_client.register_task_definition.call_args[1]["memory"] == "2"
def test_deploy_flows_enable_task_revisions_no_tags(monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {"tags": []}
boto3_client.run_task.return_value = {"tasks": [{"taskArn": "test"}]}
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
agent = FargateAgent(enable_task_revisions=True)
agent.deploy_flow(
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"id",
"version":
2,
"name":
"name",
}),
"id":
"id",
"name":
"name",
}))
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
boto3_client.register_task_definition.assert_called_with(
containerDefinitions=[{
"name":
"flow",
"image":
"test/name:tag",
"command": ["/bin/sh", "-c", "prefect execute cloud-flow"],
"environment": [
{
"name": "PREFECT__CLOUD__API",
"value": "https://api.prefect.io"
},
{
"name": "PREFECT__CLOUD__AGENT__LABELS",
"value": "[]"
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "false"
},
{
"name": "PREFECT__LOGGING__LEVEL",
"value": "DEBUG"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
],
"essential":
True,
}],
family="name",
networkMode="awsvpc",
requiresCompatibilities=["FARGATE"],
tags=[
{
"key": "PrefectFlowId",
"value": "id"
},
{
"key": "PrefectFlowVersion",
"value": "2"
},
],
)
assert boto3_client.run_task.called
assert agent.task_definition_name == "name"
