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