def test_run_flow(monkeypatch, tmpdir, job_spec_file):
environment = KubernetesJobEnvironment(job_spec_file=job_spec_file)
flow_runner = MagicMock()
flow_runner_class = MagicMock(return_value=flow_runner)
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner_class),
)
d = Local(str(tmpdir))
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"}), prefect.context({"flow_run_id":
"id"}):
environment.run_flow()
assert flow_runner_class.call_args[1]["flow"].name == "name"
assert flow_runner.run.call_args[1]["executor"] is environment.executor
def test_get_flow_returns_flow():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
loc = storage.add_flow(f)
runner = storage.get_flow(loc)
assert runner == f
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 = Local(directory)
storage.add_flow(flow)
environment.execute(flow)
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_run_config(api, image_on_run_config):
if image_on_run_config:
storage = Local()
image = "on-run-config"
run = DockerRun(image=image, env={"TESTING": "VALUE"})
else:
storage = Docker(
registry_url="testing", image_name="on-storage", image_tag="tag"
)
image = "testing/on-storage:tag"
run = DockerRun(env={"TESTING": "VALUE"})
agent = DockerAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": storage.serialize(),
"run_config": run.serialize(),
"core_version": "0.13.11",
}
),
"id": "id",
"name": "name",
}
)
)
assert api.create_container.called
assert api.create_container.call_args[0][0] == image
assert api.create_container.call_args[1]["environment"]["TESTING"] == "VALUE"
def test_add_flow_raises_if_name_conflict():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
res = storage.add_flow(f)
g = Flow("test")
with pytest.raises(ValueError, match='name "test"'):
storage.add_flow(g)
def test_add_flow_with_weird_name_is_cleaned():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
flow = Flow("WELL what do you know?!~? looks like a test!!!!")
loc = s.add_flow(flow)
assert "?" not in loc
assert "!" not in loc
assert " " not in loc
assert "~" not in loc
def test_run_flow_no_flow_run_id_in_context(monkeypatch, tmpdir):
environment = Environment()
d = Local(str(tmpdir))
d.add_flow(Flow("name"))
with set_temporary_config({"cloud.auth_token": "test"}):
with pytest.raises(ValueError):
environment.run_flow()
def test_add_flow_raises_if_name_conflict():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
res = storage.add_flow(f)
g = Flow("test")
with pytest.raises(ValueError) as exc:
storage.add_flow(g)
assert 'name "test"' in str(exc.value)
def test_run_flow_no_flow_run_id_in_context(monkeypatch):
environment = Environment()
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(Flow("name"))
with set_temporary_config({"cloud.auth_token": "test"}):
with pytest.raises(ValueError):
environment.run_flow()
def test_containment():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
f = Flow("test")
s.add_flow(f)
assert True not in s
assert f not in s
assert "test" in s
assert Flow("other") not in s
assert "other" not in s
def build_flow_run(self, config, storage=None):
if storage is None:
storage = Local()
return GraphQLResult({
"flow":
GraphQLResult({
"storage": storage.serialize(),
"run_config": RunConfigSchema().dump(config),
"id": "new_id",
"core_version": "0.13.0",
}),
"id":
"id",
})
def test_get_flow_from_file_returns_flow(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)
f = Flow("test-flow")
storage = Local(stored_as_script=True, path=full_path)
storage.add_flow(f)
flow = storage.get_flow(full_path)
assert flow.run()
def test_add_flow_to_storage():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
assert f.name not in storage
res = storage.add_flow(f)
assert res.endswith("test.prefect")
assert f.name in storage
with open(os.path.join(tmpdir, "test.prefect"), "rb") as f:
wat = f.read()
assert isinstance(wat, bytes)
assert cloudpickle.loads(wat).name == "test"
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
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": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert not agent.batch_client.create_namespaced_job.called
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_docker_agent_deploy_flow_no_pull_using_environment_metadata(api):
agent = DockerAgent(no_pull=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"storage":
Local().serialize(),
"environment":
LocalEnvironment(metadata={
"image": "name:tag"
}).serialize(),
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
}))
assert not api.pull.called
assert api.create_container.called
assert api.start.called
def test_local_agent_deploy_no_existing_python_path(monkeypatch, cloud_api):
monkeypatch.delenv("PYTHONPATH", raising=False)
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"])
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local(directory="test").serialize(),
"id": "foo",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert popen.called
assert len(agent.processes) == 1
assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_run_flow(monkeypatch):
environment = DaskKubernetesEnvironment()
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
kube_cluster = MagicMock()
monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster)
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "flow_env.prefect"), "w+") as env:
storage = Local(directory)
flow = prefect.Flow("test", storage=storage)
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"
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x, y):
return x + y
with prefect.Flow(
"simple map",
environment=LocalEnvironment(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers1 = numbers()
flow.numbers2 = numbers()
flow.add = add.map(flow.numbers1, flow.numbers2)
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x):
if x == 2:
raise prefect.engine.signals.FAIL("Don't run for 2")
elif not isinstance(x, int):
return -99
return x + 10
with prefect.Flow(
"simple map",
environment=LocalEnvironment(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers = numbers()
flow.add1 = add.map(flow.numbers)
flow.add2 = add.map(flow.add1)
flow.add2.trigger = prefect.triggers.all_finished
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
batch_client = MagicMock()
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(return_value=batch_client))
core_client = MagicMock()
core_client.list_namespaced_pod.return_value = MagicMock(items=[])
monkeypatch.setattr("kubernetes.client.CoreV1Api",
MagicMock(return_value=core_client))
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
}))
assert not agent.batch_client.create_namespaced_job.called
def test_docker_agent_deploy_flow_unsupported_run_config(api):
agent = DockerAgent()
with pytest.raises(
TypeError,
match=
"`run_config` of type `LocalRun`, only `DockerRun` is supported",
):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"run_config": LocalRun().serialize(),
"id": "foo",
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_docker_agent_deploy_flow_no_pull_using_environment_metadata(
monkeypatch, cloud_api):
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":
LocalEnvironment(metadata={
"image": "name:tag"
}).serialize(),
"core_version":
"0.13.0",
}),
"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, cloud_api):
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": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
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_run_flow_calls_callbacks(monkeypatch):
start_func = MagicMock()
exit_func = MagicMock()
file_path = os.path.dirname(
prefect.environments.execution.dask.k8s.__file__)
environment = KubernetesJobEnvironment(path.join(file_path, "job.yaml"),
on_start=start_func,
on_exit=exit_func)
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "flow_env.prefect"), "w+") as env:
storage = Local(directory)
flow = prefect.Flow("test", storage=storage)
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"
assert start_func.called
assert exit_func.called
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, api):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
agent = DockerAgent()
with pytest.raises(ValueError):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"id": "foo",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
"name": "name",
"version": "version",
}
)
)
assert not api.pull.called
def test_create_local_storage_without_validation():
storage = Local(directory="C:\\Users\\chris\\.prefect\\flows", validate=False)
assert storage
assert storage.directory == "C:\\Users\\chris\\.prefect\\flows"
assert isinstance(storage.result, LocalResult)
assert storage.result.dir == storage.directory
def test_docker_agent_deploy_flow_uses_environment_metadata(api):
agent = DockerAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"id": "foo",
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "repo/name:tag"}
).serialize(),
"core_version": "0.13.0",
}
),
"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 flow-run"
assert api.create_container.call_args[1]["host_config"]["AutoRemove"] is True
assert api.start.call_args[1]["container"] == "container_id"
def test_run_flow_calls_callbacks(monkeypatch):
start_func = MagicMock()
exit_func = MagicMock()
environment = DaskKubernetesEnvironment(on_start=start_func, on_exit=exit_func)
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
kube_cluster = MagicMock()
monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster)
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.dask.k8s.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}), prefect.context(
{"flow_run_id": "id"}
):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
assert start_func.called
assert exit_func.called
