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_register_task_definition(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()
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")
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_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_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_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_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_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"
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_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_disable_task_revisions_with_external_kwargs(
monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_resource = MagicMock()
streaming_body = MagicMock()
streaming_body.read.return_value.decode.return_value = '{"cpu": "256", "networkConfiguration": "test", "tags": [{"key": "test", "value": "test"}]}'
boto3_resource.return_value.Object.return_value.get.return_value = {
"Body": streaming_body
}
boto3_client.describe_task_definition.return_value = {}
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))
monkeypatch.setattr("boto3.resource", boto3_resource)
agent = FargateAgent(
enable_task_revisions=False,
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",
cluster="test",
labels=["aws", "staging"],
)
agent.deploy_flow(
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"new_id",
"version":
6,
"name":
"name",
}),
"id":
"id",
"name":
"name",
}))
assert agent.task_definition_kwargs == {}
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.not_called
boto3_client.run_task.assert_called_with(
launchType="FARGATE",
networkConfiguration="test",
cluster="test",
overrides={
"containerOverrides": [{
"name":
"flow",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": ""
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
],
}]
},
taskDefinition="prefect-task-new_id",
tags=[{
"key": "test",
"value": "test"
}],
)
assert agent.task_definition_name == "prefect-task-new_id"
def test_deploy_flows_enable_task_revisions_with_external_kwargs(
monkeypatch, runner_token):
boto3_client = MagicMock()
boto3_resource = MagicMock()
streaming_body = MagicMock()
streaming_body.read.return_value.decode.return_value = '{"cpu": "256", "networkConfiguration": "test", "tags": [{"key": "test", "value": "test"}]}'
boto3_resource.return_value.Object.return_value.get.return_value = {
"Body": streaming_body
}
boto3_client.describe_task_definition.return_value = {
"tags": [
{
"key": "PrefectFlowId",
"value": "id"
},
{
"key": "PrefectFlowVersion",
"value": "5"
},
]
}
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))
monkeypatch.setattr("boto3.resource", boto3_resource)
agent = FargateAgent(
enable_task_revisions=True,
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",
cluster="test",
tags=[{
"key": "team",
"value": "data"
}],
labels=["aws", "staging"],
)
agent.deploy_flow(
GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"id":
"new_id",
"version":
6,
"name":
"name",
}),
"id":
"id",
"name":
"name",
}))
assert boto3_client.describe_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": "['aws', 'staging']",
},
{
"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,
}],
cpu="256",
family="name",
networkMode="awsvpc",
requiresCompatibilities=["FARGATE"],
tags=[
{
"key": "test",
"value": "test"
},
{
"key": "PrefectFlowId",
"value": "new_id"
},
{
"key": "PrefectFlowVersion",
"value": "6"
},
],
)
boto3_client.run_task.assert_called_with(
launchType="FARGATE",
networkConfiguration="test",
cluster="test",
overrides={
"containerOverrides": [{
"name":
"flow",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": ""
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
],
}]
},
taskDefinition="name",
tags=[{
"key": "test",
"value": "test"
}],
)
assert agent.task_definition_name == "name"
