def test_setup_definition_register(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.register_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
environment = FargateTaskEnvironment(
family="test",
containerDefinitions=[
{
"name": "flow-container",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}
],
)
environment.setup(
Flow(
"test",
storage=Docker(registry_url="test", image_name="image", image_tag="tag"),
)
)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1]["family"] == "test"
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"
] == [
{
"name": "flow-container",
"image": "test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.FargateTaskEnvironment().run_flow()'",
],
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": prefect.config.cloud.graphql,
},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{"name": "PREFECT__LOGGING__EXTRA_LOGGERS", "value": "[]",},
],
"essential": True,
}
]
def test_setup_definition_exists(monkeypatch):
existing_task_definition = {
"containerDefinitions": [{
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": config.cloud.graphql
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
],
"name":
"flow-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'",
],
}],
}
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {
"taskDefinition": existing_task_definition
}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
environment = FargateTaskEnvironment()
environment.setup(
Flow(
"test",
storage=Docker(registry_url="test",
image_name="image",
image_tag="tag"),
))
assert boto3_client.describe_task_definition.called
assert not boto3_client.register_task_definition.called
def test_setup_definition_exists(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
environment = FargateTaskEnvironment()
environment.setup(Docker(registry_url="test", image_name="image", image_tag="tag"))
assert boto3_client.describe_task_definition.called
def test_entire_environment_process_together(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.register_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "session")
monkeypatch.setenv("REGION_NAME", "region")
with prefect.context({"flow_run_id": "id"}), set_temporary_config(
{"cloud.auth_token": "test", "logging.extra_loggers": "['test_logger']",}
):
storage = Docker(registry_url="test", image_name="image", image_tag="tag")
flow = Flow("name", storage=storage)
environment = FargateTaskEnvironment(
containerDefinitions=[
{
"name": "flow-container",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}
],
cluster="test",
family="test",
taskDefinition="test",
)
assert environment
assert environment.aws_access_key_id == "id"
assert environment.aws_secret_access_key == "secret"
assert environment.aws_session_token == "session"
assert environment.region_name == "region"
environment.setup(flow=flow)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1]["family"] == "test"
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"
] == [
{
"name": "flow-container",
"image": "test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.FargateTaskEnvironment().run_flow()'",
],
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": prefect.config.cloud.graphql,
},
{"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": "['test_logger']",
},
],
"essential": True,
}
]
environment.execute(flow=flow)
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["taskDefinition"] == "test"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [
{
"name": "flow-container",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": prefect.config.cloud.get("auth_token"),
},
{"name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": "id"},
{"name": "PREFECT__CONTEXT__IMAGE", "value": "test/image:tag"},
],
}
]
}
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["cluster"] == "test"
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(prefect.Flow("name"))
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
flow_run=[
GraphQLResult(
{
"flow": GraphQLResult(
{"name": "name", "storage": d.serialize(),}
)
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
def test_validate_definition_not_changed_when_out_of_order_in_second_container(
monkeypatch, ):
existing_task_definition = {
"containerDefinitions": [
{
"environment": [
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
# This is added first in _render_task_definition_kwargs, so it's at the end now
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": config.cloud.graphql
},
],
"name":
"flow-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'",
],
},
{
"environment": [
{
"name": "foo",
"value": "bar",
},
{
"name": "foo2",
"value": "bar2",
},
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": config.cloud.graphql
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
],
"secrets": [
{
"name": "1",
"valueFrom": "1"
},
{
"name": "2",
"valueFrom": "2"
},
],
"mountPoints": [
{
"sourceVolume": "1",
"containerPath": "1",
"readOnly": False
},
{
"sourceVolume": "2",
"containerPath": "2",
"readOnly": False
},
],
"extraHosts": [
{
"hostname": "1",
"ipAddress": "1"
},
{
"hostname": "2",
"ipAddress": "2"
},
],
"volumesFrom": [
{
"sourceContainer": "1",
"readOnly": False
},
{
"sourceContainer": "2",
"readOnly": False
},
],
"ulimits": [
{
"name": "cpu",
"softLimit": 1,
"hardLimit": 1
},
{
"name": "memlock",
"softLimit": 2,
"hardLimit": 2
},
],
"portMappings": [
{
"containerPort": 80,
"hostPort": 80,
"protocol": "tcp"
},
{
"containerPort": 81,
"hostPort": 81,
"protocol": "tcp"
},
],
"logConfiguration": {
"logDriver":
"awslogs",
"options": {},
"secretOptions": [
{
"name": "1",
"valueFrom": "1"
},
{
"name": "2",
"valueFrom": "2"
},
],
},
"name":
"some-other-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
],
},
],
"memory":
256,
"cpu":
512,
}
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {
"taskDefinition": existing_task_definition
}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
environment = FargateTaskEnvironment(
memory=256,
cpu=512,
containerDefinitions=[
{},
{
"environment": [
{
"name": "foo2",
"value": "bar2",
},
{
"name": "foo",
"value": "bar",
},
],
"secrets": [
{
"name": "2",
"valueFrom": "2"
},
{
"name": "1",
"valueFrom": "1"
},
],
"mountPoints": [
{
"sourceVolume": "2",
"containerPath": "2",
"readOnly": False
},
{
"sourceVolume": "1",
"containerPath": "1",
"readOnly": False
},
],
"extraHosts": [
{
"hostname": "2",
"ipAddress": "2"
},
{
"hostname": "1",
"ipAddress": "1"
},
],
"volumesFrom": [
{
"sourceContainer": "2",
"readOnly": False
},
{
"sourceContainer": "1",
"readOnly": False
},
],
"ulimits": [
{
"name": "memlock",
"softLimit": 2,
"hardLimit": 2
},
{
"name": "cpu",
"softLimit": 1,
"hardLimit": 1
},
],
"portMappings": [
{
"containerPort": 81,
"hostPort": 81,
"protocol": "tcp"
},
{
"containerPort": 80,
"hostPort": 80,
"protocol": "tcp"
},
],
"logConfiguration": {
"logDriver":
"awslogs",
"options": {},
"secretOptions": [
{
"name": "2",
"valueFrom": "2"
},
{
"name": "1",
"valueFrom": "1"
},
],
},
"name":
"some-other-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
],
},
],
)
environment.setup(
Flow(
"test",
storage=Docker(registry_url="test",
image_name="image",
image_tag="tag"),
))
assert boto3_client.describe_task_definition.called
assert not boto3_client.register_task_definition.called
def test_entire_environment_process_together(monkeypatch):
boto3_client = MagicMock()
boto3_client.describe_task_definition.side_effect = ClientError({}, None)
boto3_client.register_task_definition.return_value = {}
boto3_client.run_task.return_value = {}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
monkeypatch.setenv("AWS_SESSION_TOKEN", "session")
monkeypatch.setenv("REGION_NAME", "region")
with prefect.context({"flow_run_id": "id"}):
storage = Docker(registry_url="test",
image_name="image",
image_tag="tag")
environment = FargateTaskEnvironment(
containerDefinitions=[{
"name": "flow-container",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}],
cluster="test",
family="test",
taskDefinition="test",
)
assert environment
assert environment.aws_access_key_id == "id"
assert environment.aws_secret_access_key == "secret"
assert environment.aws_session_token == "session"
assert environment.region_name == "region"
environment.setup(storage=storage)
assert boto3_client.describe_task_definition.called
assert boto3_client.register_task_definition.called
assert boto3_client.register_task_definition.call_args[1][
"family"] == "test"
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"] == [{
"name":
"flow-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.Flow.load(prefect.context.flow_file_path).environment.run_flow()'",
],
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": prefect.config.cloud.graphql,
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__LOGGING__LOG_TO_CLOUD",
"value": "true"
},
],
"essential":
True,
}]
environment.execute(storage=storage, flow_location=".prefect/flows")
assert boto3_client.run_task.called
assert boto3_client.run_task.call_args[1]["taskDefinition"] == "test"
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow-container",
"environment": [
{
"name": "PREFECT__CLOUD__AUTH_TOKEN",
"value": prefect.config.cloud.agent.auth_token,
},
{
"name": "PREFECT__CONTEXT__FLOW_RUN_ID",
"value": "id"
},
{
"name": "PREFECT__CONTEXT__IMAGE",
"value": "test/image:tag"
},
{
"name": "PREFECT__CONTEXT__FLOW_FILE_PATH",
"value": ".prefect/flows",
},
],
}]
}
assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
assert boto3_client.run_task.call_args[1]["cluster"] == "test"
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "flow_env.prefect"), "w+"):
flow = prefect.Flow("test")
flow_path = os.path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
with set_temporary_config({"cloud.auth_token": "test"}):
with prefect.context(flow_file_path=os.path.join(
directory, "flow_env.prefect")):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "test"
def test_validate_definition_not_changed_when_names_are_in_arn(monkeypatch):
existing_task_definition = {
"containerDefinitions": [{
"environment": [
{
"name": "PREFECT__CLOUD__GRAPHQL",
"value": config.cloud.graphql
},
{
"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
"value": "false"
},
{
"name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudFlowRunner",
},
{
"name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
"value": "prefect.engine.cloud.CloudTaskRunner",
},
{
"name": "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS",
"value": "true"
},
{
"name": "PREFECT__LOGGING__EXTRA_LOGGERS",
"value": str(config.logging.extra_loggers),
},
],
"name":
"flow-container",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'",
],
}],
"taskRoleArn":
"arn:aws:iam::000000000000:role/my-role-name",
"memory":
256,
"cpu":
512,
}
boto3_client = MagicMock()
boto3_client.describe_task_definition.return_value = {
"taskDefinition": existing_task_definition
}
monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))
environment = FargateTaskEnvironment(memory=256,
cpu=512,
taskRoleArn="my-role-name")
environment.setup(
Flow(
"test",
storage=Docker(registry_url="test",
image_name="image",
image_tag="tag"),
))
assert boto3_client.describe_task_definition.called
assert not boto3_client.register_task_definition.called
