def test_docker_init_responds_to_python_version(monkeypatch, version_info):
version_mock = MagicMock(major=version_info[0], minor=version_info[1])
monkeypatch.setattr(sys, "version_info", version_mock)
storage = Docker()
assert storage.base_image == "python:{}.{}".format(*version_info)
def test_k8s_agent_replace_yaml(monkeypatch, cloud_api):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
monkeypatch.setenv("JOB_CPU_LIMIT", "cl")
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
with set_temporary_config({
"cloud.agent.auth_token": "token",
"logging.log_to_cloud": True
}):
agent = KubernetesAgent()
job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
assert (job["spec"]["template"]["metadata"]["labels"]
["prefect.io/flow_run_id"] == "id")
assert (job["spec"]["template"]["spec"]["containers"][0]["image"] ==
"test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[0]["value"] == "https://api.prefect.io"
assert env[1]["value"] == "token"
assert env[2]["value"] == "id"
assert env[3]["value"] == "new_id"
assert env[4]["value"] == "default"
assert env[5]["value"] == "[]"
assert env[6]["value"] == "true"
assert (job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
== "my-secret")
resources = job["spec"]["template"]["spec"]["containers"][0][
"resources"]
assert resources["requests"]["memory"] == "mr"
assert resources["limits"]["memory"] == "ml"
assert resources["requests"]["cpu"] == "cr"
assert resources["limits"]["cpu"] == "cl"
assert (job["spec"]["template"]["spec"]["containers"][0]
["imagePullPolicy"] == "IfNotPresent")
assert job["spec"]["template"]["spec"].get("serviceAccountName",
None) is None
def test_setup_dask_environment_passes():
environment = DaskKubernetesEnvironment()
environment.setup(storage=Docker())
assert environment
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 boto3_client.run_task.called_with(taskDefinition="name")
name="example-selenium",
schedule=Schedule(clocks=[
# TODO: specify the schedule you want this to run, and with what parameters
# https://docs.prefect.io/core/concepts/schedules.html
CronClock(cron='0 0 * * *',
parameter_defaults=dict(
home_page='https://www.metacritic.com/',
gaming_platform='Switch')),
]),
# TODO: specify the environment you want to execute the Flow in (from Prefect Cloud)
environment=KubernetesJobEnvironment(job_spec_file='job_spec.yaml', ),
storage=Docker(
# TODO: change to your docker registry:
# https://docs.prefect.io/cloud/recipes/configuring_storage.html
registry_url='szelenka',
# TODO: need to specify a base Docker image which has the chromedriver dependencies already installed
base_image='szelenka/python-selenium-chromium:3.7.4',
# TODO: 'pin' the exact versions you used on your development machine
python_dependencies=['selenium==3.141.0', 'sqlalchemy==1.3.15'],
),
# TODO: specify how you want to handle results
# https://docs.prefect.io/core/concepts/results.html#results-and-result-handlers
result_handler=LocalResultHandler()) as flow:
# specify the DAG input parameters
_path_to_chromedriver = Parameter('path_to_chromedriver',
default='/usr/bin/chromedriver')
_home_page_url = Parameter('home_page',
default='https://www.metacritic.com/')
_gaming_platform = Parameter('gaming_platform', default='Switch')
_db_file = Parameter("db_file",
default='game_reviews.sqlite',
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 boto3_client.run_task.called_with(
taskDefinition="prefect-task-new_id")
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_docker_storage_allows_for_user_provided_config_locations():
storage = Docker(env_vars={"PREFECT__USER_CONFIG_PATH": "1"})
assert storage.env_vars == {"PREFECT__USER_CONFIG_PATH": "1"}
def test_create_docker_storage():
storage = Docker(secrets=["cloud_creds"])
assert storage
assert storage.logger
assert len(storage.secrets) == 1
assert storage.secrets == ["cloud_creds"]
def test_docker_init_responds_to_python_version(monkeypatch, version_info):
version_mock = MagicMock(major=version_info[0], minor=version_info[1])
monkeypatch.setattr(sys, "version_info", version_mock)
monkeypatch.setattr(prefect, "__version__", "0.9.2+c2394823")
storage = Docker()
assert storage.base_image == "python:{}.{}-slim".format(*version_info)
def test_docker_init_responds_to_prefect_version(monkeypatch, version):
monkeypatch.setattr(prefect, "__version__", version[0])
storage = Docker()
assert storage.prefect_version == version[1]
class TestGenerateTaskDefinition:
def generate_task_definition(self, run_config, storage=None, **kwargs):
if storage is None:
storage = Local()
agent = ECSAgent(**kwargs)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage": storage.serialize(),
"run_config": run_config.serialize(),
"id": "flow-id",
"version": 1,
"name": "Test Flow",
"core_version": "0.13.0",
}),
"id":
"flow-run-id",
})
return agent.generate_task_definition(flow_run, run_config)
@pytest.mark.parametrize("use_path", [False, True])
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_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"
},
]
@pytest.mark.parametrize(
"run_config, storage, expected",
[
(
ECSRun(),
Docker(registry_url="test", image_name="name",
image_tag="tag"),
"test/name:tag",
),
(ECSRun(image="myimage"), Local(), "myimage"),
(ECSRun(), Local(), "prefecthq/prefect:all_extras-0.13.0"),
],
ids=["on-storage", "on-run_config", "default"],
)
def test_generate_task_definition_image(self, run_config, storage,
expected):
taskdef = self.generate_task_definition(run_config, storage)
assert taskdef["containerDefinitions"][0]["image"] == expected
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_generate_task_definition_resources(self):
taskdef = self.generate_task_definition(
ECSRun(cpu="2048", memory="4096"))
assert taskdef["cpu"] == "2048"
assert taskdef["memory"] == "4096"
@pytest.mark.parametrize(
"on_run_config, on_agent, expected",
[
(None, None, None),
("task-role-1", None, "task-role-1"),
(None, "task-role-2", "task-role-2"),
("task-role-1", "task-role-2", "task-role-1"),
],
)
def test_generate_task_definition_task_role_arn(self, on_run_config,
on_agent, expected):
taskdef = self.generate_task_definition(
ECSRun(task_role_arn=on_run_config), task_role_arn=on_agent)
assert taskdef.get("taskRoleArn") == 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_generate_task_definition_multiple_containers(self):
"""A container with the name "flow" is used for prefect stuff"""
run_config = ECSRun(
task_definition={
"containerDefinitions": [
{
"name": "other",
"image": "other-image"
},
{
"name": "flow",
"cpu": 1234
},
]
},
image="flow-image",
)
taskdef = self.generate_task_definition(run_config)
assert taskdef["containerDefinitions"][0]["name"] == "other"
assert taskdef["containerDefinitions"][0]["image"] == "other-image"
assert taskdef["containerDefinitions"][1]["name"] == "flow"
assert taskdef["containerDefinitions"][1]["cpu"] == 1234
assert taskdef["containerDefinitions"][1]["image"] == "flow-image"
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.execution.load_and_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 register_workflow(prefect_register_token_secret_name: str):
"""
Registers the workflow to Prefect Cloud
Parameters:
prefect_register_token_secret_name [str]
-- name of aws secrets manager secret where prefect register token is stored
"""
flow_module = __import__("flow")
flow_name = f"{env}_{flow_module.flow.name}"
flow_module.flow.name = flow_name
flow_module.flow.environment = FargateTaskEnvironment(
requiresCompatibilities=["FARGATE"],
region=aws_region,
labels=[f"{env}_dataflow_automation"],
taskDefinition=flow_name,
family=flow_name,
cpu="512",
memory="3072",
networkMode="awsvpc",
networkConfiguration={
"awsvpcConfiguration": {
"assignPublicIp": "ENABLED",
"subnets": subnets,
"securityGroups": [],
}
},
containerDefinitions=[{
"logConfiguration": {
"logDriver": "awslogs",
"options": {
"awslogs-region": aws_region,
"awslogs-group": f"{env}_dataflow_automation_workflows",
"awslogs-stream-prefix": flow_name,
},
}
}],
executionRoleArn=execution_role_arn,
taskRoleArn=task_role_arn,
cluster=f"{env}_dataflow_automation_workflows",
)
# Set the flow storage. Where to get the code from
flow_module.flow.storage = Docker(
registry_url=f"{account_id}.dkr.ecr.{aws_region}.amazonaws.com",
image_name=flow_name,
image_tag="latest",
python_dependencies=["boto3"],
env_vars={"PYTHONPATH": "/opt/prefect/flows"},
)
# Authenticate to ECR as the registration process pushes the image to AWS
ecr_authenticate()
# Instantiate the prefect client
prefect_client = Client(api_token=get_prefect_token(
secret_name=prefect_register_token_secret_name))
# Create ECR repository
create_ecr_repository(flow_name=flow_name)
# Register the Workflow
prefect_client.register(flow=flow_module.flow,
project_name=f"{env}_dataflow_automation")
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_cant_create_docker_with_both_base_image_and_dockerfile():
with pytest.raises(ValueError):
Docker(dockerfile="path/to/file", base_image="python:3.6")
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 boto3_client.run_task.called_with(taskDefinition="name")
def test_create_dockerfile_with_flow_file(no_docker_host_var, tmpdir):
contents = """from prefect import Flow\nf=Flow('test-flow')"""
full_path = os.path.join(tmpdir, "flow.py")
with open(full_path, "w") as f:
f.write(contents)
with open(os.path.join(tmpdir, "test"), "w+") as t:
t.write("asdf")
with tempfile.TemporaryDirectory() as tempdir_inside:
storage = Docker(files={full_path: "flow.py"},
stored_as_script=True,
path="flow.py")
f = Flow("test-flow")
storage.add_flow(f)
dpath = storage.create_dockerfile_object(directory=tempdir_inside)
with open(dpath, "r") as dockerfile:
output = dockerfile.read()
assert "COPY flow.py flow.py" in output
storage = Docker(files={full_path: "flow.py"}, stored_as_script=True)
f = Flow("test-flow")
storage.add_flow(f)
with pytest.raises(ValueError):
storage.create_dockerfile_object(directory=tempdir_inside)
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_create_docker_storage():
storage = Docker()
assert storage
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_serialize_docker_storage():
storage = Docker()
serialized_storage = storage.serialize()
assert serialized_storage["type"] == "Docker"
def test_setup_cloud_environment_passes():
environment = CloudEnvironment()
environment.setup(storage=Docker())
assert environment
def test_docker_storage_doesnt_have_get_flow_method():
storage = Docker(base_image="python:3.6")
with pytest.raises(NotImplementedError):
storage.get_flow("")
def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, cloud_api):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
monkeypatch.setenv("JOB_CPU_LIMIT", "cl")
monkeypatch.setenv("IMAGE_PULL_POLICY", "custom_policy")
monkeypatch.setenv("SERVICE_ACCOUNT_NAME", "svc_name")
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
with set_temporary_config({
"cloud.agent.auth_token": "token",
"logging.log_to_cloud": True
}):
agent = KubernetesAgent(
env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar"))
job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
assert (job["spec"]["template"]["metadata"]["labels"]
["prefect.io/flow_run_id"] == "id")
assert (job["spec"]["template"]["spec"]["containers"][0]["image"] ==
"test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[0]["value"] == "https://api.prefect.io"
assert env[1]["value"] == "token"
assert env[2]["value"] == "id"
assert env[3]["value"] == "new_id"
assert env[4]["value"] == "default"
assert env[5]["value"] == "[]"
assert env[6]["value"] == "true"
user_vars = [
dict(name="AUTH_THING", value="foo"),
dict(name="PKG_SETTING", value="bar"),
]
assert env[-1] in user_vars
assert env[-2] in user_vars
assert (job["spec"]["template"]["spec"]["containers"][0]
["imagePullPolicy"] == "custom_policy")
assert job["spec"]["template"]["spec"][
"serviceAccountName"] == "svc_name"
def test_add_similar_flows_fails():
storage = Docker()
flow = prefect.Flow("test")
storage.add_flow(flow)
with pytest.raises(ValueError):
storage.add_flow(flow)
def test_execute_storage_missing_fields():
environment = DaskKubernetesEnvironment()
with pytest.raises(ValueError):
environment.execute(storage=Docker(), flow_location="")
def test_files_not_absolute_path():
with pytest.raises(ValueError):
storage = Docker(files={"test": "test"})
return 1
@task
def transform(x):
return x + 1
@task(tags=['testtagherelol'])
def pdata(orig, trans):
print(orig)
print(trans)
with Flow(
"testflow-fargate",
environment=RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor",
executor_kwargs={'n_workers': 16},
),
storage=Docker(registry_url="joshmeek18", image_name="flows"),
) as flow:
x = data()
t = transform(x)
p1 = pdata(x, t)
p2 = pdata(x, t)
# flow.run()
flow.register(project_name="Demo")
# flow.visualize()
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("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",
"image": "image",
"command": [],
"environment": [],
"essential": True,
}],
cluster="test",
)
assert environment
assert environment.aws_access_key_id == "id"
assert environment.aws_secret_access_key == "secret"
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"] ==
"prefect-task-id-custom")
assert boto3_client.register_task_definition.call_args[1][
"containerDefinitions"] == [{
"name":
"flow",
"image":
"test/image:tag",
"command": [
"/bin/sh",
"-c",
"python -c 'from prefect.environments import FargateTaskEnvironment; 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"
},
],
"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"] ==
"prefect-task-id-custom")
assert boto3_client.run_task.call_args[1]["overrides"] == {
"containerOverrides": [{
"name":
"flow",
"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"
