def test_environment_execute_calls_callbacks():
start_func = MagicMock()
exit_func = MagicMock()
with tempfile.TemporaryDirectory() as directory:
@prefect.task
def add_to_dict():
with open(path.join(directory, "output"), "w") as tmp:
tmp.write("success")
with open(path.join(directory, "flow_env.prefect"), "w+") as env:
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_path = path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
environment = RemoteEnvironment(on_start=start_func, on_exit=exit_func)
storage = Docker(registry_url="test")
environment.execute(storage, flow_path)
with open(path.join(directory, "output"), "r") as file:
assert file.read() == "success"
assert start_func.called
assert exit_func.called
def test_docker_agent_deploy_flow_no_registry_does_not_pull(
monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
}))
assert not api.pull.called
assert api.create_container.called
assert api.start.called
def test_k8s_agent_replace_yaml_responds_to_logging_config(
monkeypatch, runner_token, flag):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
"id":
"new_id",
}),
"id":
"id",
"name":
"name",
})
agent = KubernetesAgent(no_cloud_logs=flag)
job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[6]["value"] == str(not flag).lower()
def test_docker_agent_deploy_flow_no_pull_using_environment_metadata(
monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent(no_pull=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Local().serialize(),
"environment":
RemoteEnvironment(metadata={
"image": "name:tag"
}).serialize(),
}),
"id":
"id",
"name":
"name",
}))
assert not api.pull.called
assert api.create_container.called
assert api.start.called
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, runner_token):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "foo",
"environment": RemoteEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_k8s_agent_deploy_flow(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
batch_client = MagicMock()
batch_client.create_namespaced_job.return_value = {}
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(retrurn_value=batch_client))
agent = KubernetesAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
"id":
"id",
}),
"id":
"id",
}))
assert agent.batch_client.create_namespaced_job.called
assert (agent.batch_client.create_namespaced_job.call_args[1]["namespace"]
== "default")
assert (agent.batch_client.create_namespaced_job.call_args[1]["body"]
["apiVersion"] == "batch/v1")
def test_docker_agent_deploy_with_no_interface_check_linux(
monkeypatch, runner_token, linux_platform):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
get_ip = MagicMock()
monkeypatch.setattr("prefect.agent.docker.agent.get_docker_ip", get_ip)
agent = DockerAgent(docker_interface=False)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
}))
assert not get_ip.called
def test_docker_agent_network(monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
api.create_networking_config.return_value = {"test-network": "config"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent(network="test-network")
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
}))
assert agent.network == "test-network"
args, kwargs = api.create_container.call_args
assert kwargs["networking_config"] == {"test-network": "config"}
def test_k8s_agent_deploy_flow_uses_environment_metadata(
monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
batch_client = MagicMock()
batch_client.create_namespaced_job.return_value = {}
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(retrurn_value=batch_client))
agent = KubernetesAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage":
Local().serialize(),
"environment":
RemoteEnvironment(metadata={
"image": "repo/name:tag"
}).serialize(),
"id":
"id",
}),
"id":
"id",
}))
assert agent.batch_client.create_namespaced_job.called
assert (
agent.batch_client.create_namespaced_job.call_args[1]["body"]["spec"]
["template"]["spec"]["containers"][0]["image"] == "repo/name:tag")
def test_create_remote_environment():
environment = RemoteEnvironment()
assert environment
assert environment.executor == prefect.config.engine.executor.default_class
assert environment.executor_kwargs == {}
assert environment.labels == set()
assert environment.logger.name == "prefect.RemoteEnvironment"
def test_get_flow_image_raises_on_missing_info():
flow = Flow(
"test",
environment=RemoteEnvironment(),
storage=Local(),
)
with pytest.raises(ValueError):
image = get_flow_image(flow=flow)
def test_get_flow_image_env_metadata():
flow = Flow(
"test",
environment=RemoteEnvironment(metadata={"image": "repo/name:tag"}),
storage=Local(),
)
image = get_flow_image(flow=flow)
assert image == "repo/name:tag"
def test_create_remote_environment_populated():
environment = RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor",
executor_kwargs={"address": "test"},
)
assert environment
assert environment.executor == "prefect.engine.executors.DaskExecutor"
assert environment.executor_kwargs == {"address": "test"}
def test_k8s_agent_replace_yaml(monkeypatch, runner_token, 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":
RemoteEnvironment().serialize(),
"id":
"new_id",
}),
"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"
def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, runner_token,
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":
RemoteEnvironment().serialize(),
"id":
"new_id",
}),
"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
def test_get_flow_image_docker_storage():
flow = Flow(
"test",
environment=RemoteEnvironment(),
storage=Docker(registry_url="test", image_name="name",
image_tag="tag"),
)
image = get_flow_image(flow=flow)
assert image == "test/name:tag"
def test_create_remote_environment_populated():
environment = RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor",
executor_kwargs={"address": "test"},
labels=["foo", "bar", "good"],
)
assert environment
assert environment.executor == "prefect.engine.executors.DaskExecutor"
assert environment.executor_kwargs == {"address": "test"}
assert environment.labels == set(["foo", "bar", "good"])
def test_environment_execute():
with tempfile.TemporaryDirectory() as directory:
@prefect.task
def add_to_dict():
with open(path.join(directory, "output"), "w") as file:
file.write("success")
with open(path.join(directory, "flow_env.prefect"), "w+") as env:
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_path = path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
environment = RemoteEnvironment()
storage = Docker(registry_url="test")
environment.execute(storage, flow_path)
with open(path.join(directory, "output"), "r") as file:
assert file.read() == "success"
def test_get_flow_image_raises_on_missing_info():
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"environment": RemoteEnvironment().serialize(),
"id": "id",
}),
"id":
"id",
})
with pytest.raises(ValueError):
image = get_flow_image(flow_run=flow_run)
def test_get_flow_image_docker_storage():
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
"id":
"id",
}),
"id":
"id",
})
image = get_flow_image(flow_run=flow_run)
assert image == "test/name:tag"
def test_get_flow_image_env_metadata():
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Local().serialize(),
"environment":
RemoteEnvironment(metadata={
"image": "repo/name:tag"
}).serialize(),
"id":
"id",
}),
"id":
"id",
})
image = get_flow_image(flow_run=flow_run)
assert image == "repo/name:tag"
def test_create_remote_environment_populated():
def f():
pass
environment = RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor",
executor_kwargs={"address": "test"},
labels=["foo", "bar", "good"],
on_start=f,
on_exit=f,
metadata={"test": "here"},
)
assert environment
assert environment.executor == "prefect.engine.executors.DaskExecutor"
assert environment.executor_kwargs == {"address": "test"}
assert environment.labels == set(["foo", "bar", "good"])
assert environment.on_start is f
assert environment.on_exit is f
assert environment.metadata == {"test": "here"}
def test_docker_agent_deploy_flow_uses_environment_metadata(
monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
api.create_host_config.return_value = {"AutoRemove": True}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Local().serialize(),
"environment":
RemoteEnvironment(metadata={
"image": "repo/name:tag"
}).serialize(),
}),
"id":
"id",
"name":
"name",
}))
assert api.pull.called
assert api.create_container.called
assert api.start.called
assert api.create_host_config.call_args[1]["auto_remove"] is True
assert api.create_container.call_args[1][
"command"] == "prefect execute cloud-flow"
assert api.create_container.call_args[1]["host_config"][
"AutoRemove"] is True
assert api.start.call_args[1]["container"] == "container_id"
def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
"id":
"new_id",
}),
"id":
"id",
})
agent = KubernetesAgent(labels=["foo", "bar"])
job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[5]["value"] == "['foo', 'bar']"
def test_k8s_agent_replace_yaml_no_pull_secrets(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
"id":
"id",
}),
"id":
"id",
})
agent = KubernetesAgent()
job = agent.replace_job_spec_yaml(flow_run, image="test/name:tag")
assert not job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
def test_k8s_agent_deploy_flow_raises(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
batch_client = MagicMock()
batch_client.create_namespaced_job.return_value = {}
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(retrurn_value=batch_client))
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"environment": RemoteEnvironment().serialize(),
}),
"id":
"id",
}))
assert not agent.batch_client.create_namespaced_job.called
def test_docker_agent_deploy_flow_show_flow_logs(monkeypatch, runner_token):
process = MagicMock()
monkeypatch.setattr("multiprocessing.Process", process)
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent(show_flow_logs=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
RemoteEnvironment().serialize(),
}),
"id":
"id",
"name":
"name",
}))
process.assert_called_with(target=agent.stream_container_logs,
kwargs={"container_id": "container_id"})
assert len(agent.processes) == 1
assert api.create_container.called
assert api.start.called
def test_remote_environment_dependencies():
environment = RemoteEnvironment()
assert environment.dependencies == []
@task
def multiply(x):
if x == 9:
time.sleep(20)
time.sleep(1)
return x * 10
@task
def divide(x):
time.sleep(1)
return x / 10
@task
def reduce(vals):
print(sum(vals))
with Flow("mlm") as flow:
data = get_data()
m = multiply.map(data)
d = divide.map(m)
m2 = multiply.map(d)
d2 = divide.map(m2)
reduce(d2)
flow.environment = RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor")
flow.register(project_name="Demo")
def test_no_raise_on_remote_env(self):
flow = Flow("THIS IS A TEST", environment=RemoteEnvironment())
assert healthchecks.environment_dependency_check([flow]) is None