def test_create_namespaced_job_fails_outside_cluster(job_spec_file): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): environment.execute(Flow("test", storage=storage))
def test_populate_job_yaml_no_defaults(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) file_path = os.path.dirname( prefect.environments.execution.dask.k8s.__file__) with open(path.join(file_path, "job.yaml")) as job_file: job = yaml.safe_load(job_file) job["spec"]["template"]["spec"]["containers"][0] = {} del job["metadata"] del job["spec"]["template"]["metadata"] with set_temporary_config({ "cloud.graphql": "gql_test", "cloud.auth_token": "auth_test" }): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_job_spec_yaml( yaml_obj=job, docker_name="test1/test2:test3", flow_file_path="test4", ) assert (yaml_obj["metadata"]["labels"]["identifier"] == environment.identifier_label) assert yaml_obj["metadata"]["labels"]["flow_run_id"] == "id_test" assert (yaml_obj["spec"]["template"]["metadata"]["labels"] ["identifier"] == environment.identifier_label) env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "auth_test" assert env[2]["value"] == "id_test" assert env[3]["value"] == "namespace_test" assert env[4]["value"] == "test1/test2:test3" assert env[5]["value"] == "test4" assert env[10]["value"] == "[]" assert (yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] == "test1/test2:test3") assert yaml_obj["spec"]["template"]["spec"]["containers"][0][ "command"] == [ "/bin/sh", "-c", ] assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == [ "python -c 'import prefect; prefect.Flow.load(prefect.context.flow_file_path).environment.run_flow()'" ]
def test_create_k8s_job_environment_identifier_label(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) assert environment.identifier_label
def test_initialize_environment_with_spec_populates(monkeypatch, job_spec_file, initial_job_spec, default_command_args): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) assert environment._job_spec == initial_job_spec assert environment._job_spec["spec"]["template"]["spec"]["containers"][0][ "command"] == ["/bin/sh", "-c"] assert (environment._job_spec["spec"]["template"]["spec"]["containers"][0] ["args"] == default_command_args)
def test_create_k8s_job_environment(job_spec_file): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) assert environment.job_spec_file == job_spec_file assert environment.unique_job_name is False assert environment.executor is not None assert environment.labels == set() assert environment.on_start is None assert environment.on_exit is None assert environment.metadata == {} assert environment.logger.name == "prefect.KubernetesJobEnvironment"
def test_create_k8s_job_callbacks(job_spec_file): def f(): pass environment = KubernetesJobEnvironment( job_spec_file=job_spec_file, labels=["foo"], on_start=f, on_exit=f, ) assert environment.labels == set(["foo"]) assert environment.on_start is f assert environment.on_exit is f
def test_setup_k8s_job_environment_passes(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) environment.setup(storage=Docker()) assert environment
def test_execute_storage_missing_fields(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) with pytest.raises(ValueError): environment.execute(storage=Docker(), flow_location="")
def test_execute_improper_storage(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) with pytest.raises(TypeError): environment.execute(storage=Local(), flow_location="")
def test_populate_job_yaml(job_spec_file, job, default_command_args): environment = KubernetesJobEnvironment( job_spec_file=job_spec_file, unique_job_name=True ) job["spec"]["template"]["spec"]["containers"][0]["env"] = [] environment._job_spec = job with set_temporary_config( { "cloud.graphql": "gql_test", "cloud.auth_token": "auth_test", "logging.extra_loggers": "['test_logger']", } ): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_run_time_job_spec_details( docker_name="test1/test2:test3", ) assert "prefect-dask-job-" in yaml_obj["metadata"]["name"] assert len(yaml_obj["metadata"]["name"]) == 25 assert ( yaml_obj["metadata"]["labels"]["prefect.io/identifier"] == environment.identifier_label ) assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test" assert ( yaml_obj["spec"]["template"]["metadata"]["labels"]["prefect.io/identifier"] == environment.identifier_label ) env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "auth_test" assert env[2]["value"] == "id_test" assert env[3]["value"] == "namespace_test" assert env[4]["value"] == "test1/test2:test3" assert env[9]["value"] == "['test_logger']" assert ( yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] == "test1/test2:test3" ) assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["command"] == [ "/bin/sh", "-c", ] assert ( yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == default_command_args )
def test_initialize_environment_with_spec_populates(monkeypatch): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) assert environment._job_spec == "job"
def test_hash_flow(): flow = TEST_FLOW.copy() integration = prefect_saturn.PrefectCloudIntegration( prefect_cloud_project_name=TEST_PREFECT_PROJECT_NAME) with patch("prefect_saturn.core.Client", new=MockClient): flow_hash = integration._hash_flow(flow) assert isinstance(flow_hash, str) and len(flow_hash) > 0 # should be deterministic flow_hash_again = integration._hash_flow(flow) assert flow_hash == flow_hash_again # should not be impacted by storage flow.storage = Webhook( build_request_kwargs={}, build_request_http_method="POST", get_flow_request_kwargs={}, get_flow_request_http_method="GET", ) assert flow_hash == integration._hash_flow(flow) # should not be impacted by environment or run_config if RUN_CONFIG_AVAILABLE: flow.run_config = KubernetesRun() elif KUBE_JOB_ENV_AVAILABLE: flow.environment = KubernetesJobEnvironment() assert flow_hash == integration._hash_flow(flow) # should not change if you add a new task @task def goodbye_task(): logger = prefect.context.get("logger") logger.info("adios") flow.tasks = [hello_task, goodbye_task] new_flow_hash = integration._hash_flow(flow) assert isinstance(new_flow_hash, str) and len(new_flow_hash) > 0 assert new_flow_hash == flow_hash # should change if flow name changes flow.name = str(uuid.uuid4()) new_flow_hash = integration._hash_flow(flow) assert new_flow_hash != flow_hash # should change if project name changes previous_flow_hash = new_flow_hash integration.prefect_cloud_project_name = str(uuid.uuid4()) new_flow_hash = integration._hash_flow(flow) assert isinstance(new_flow_hash, str) and len(new_flow_hash) > 0 assert new_flow_hash != previous_flow_hash
def test_create_k8s_job_environment_with_deprecated_executor_kwargs(job_spec_file): with set_temporary_config( {"engine.executor.default_class": "prefect.engine.executors.LocalDaskExecutor"} ): with pytest.warns(UserWarning, match="executor_kwargs"): environment = KubernetesJobEnvironment( job_spec_file=job_spec_file, executor_kwargs={"scheduler": "synchronous"}, ) assert isinstance(environment.executor, LocalDaskExecutor) assert environment.executor.scheduler == "synchronous"
def test_create_k8s_job_environment(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) assert environment assert environment.job_spec_file == os.path.join(directory, "job.yaml") assert environment.labels == set() assert environment.logger.name == "prefect.KubernetesJobEnvironment"
def test_create_k8s_job_environment_with_executor_kwargs(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml"), executor_kwargs={"test": "here"}, ) assert environment assert environment.executor_kwargs == {"test": "here"}
def test_create_flow_run_job_fails_outside_cluster(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): environment.create_flow_run_job( docker_name="test1/test2:test3", flow_file_path="test4")
def test_roundtrip_cloudpickle(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) assert environment._job_spec == "job" new = cloudpickle.loads(cloudpickle.dumps(environment)) assert isinstance(new, KubernetesJobEnvironment) assert new._job_spec == "job"
def test_roundtrip_cloudpickle(job_spec_file, initial_job_spec): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) assert environment._job_spec == initial_job_spec new = cloudpickle.loads(cloudpickle.dumps(environment)) assert isinstance(new, KubernetesJobEnvironment) assert new._job_spec == initial_job_spec # Identifer labels do not persist assert environment.identifier_label assert new.identifier_label assert environment.identifier_label != new.identifier_label
def test_create_namespaced_job_fails_outside_cluster(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): environment.execute(Flow("test", storage=storage))
def test_create_k8s_job_environment(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml")) assert environment assert environment.job_spec_file == os.path.join(directory, "job.yaml") assert environment.unique_job_name == False assert environment.executor_kwargs == {} assert environment.labels == set() assert environment.on_start is None assert environment.on_exit is None assert environment.logger.name == "prefect.KubernetesJobEnvironment"
def test_populate_job_yaml_no_defaults(job_spec_file, job): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) job["spec"]["template"]["spec"]["containers"][0] = {} del job["metadata"] del job["spec"]["template"]["metadata"] with set_temporary_config({ "cloud.graphql": "gql_test", "cloud.auth_token": "auth_test" }): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_job_spec_yaml( yaml_obj=job, docker_name="test1/test2:test3", ) assert (yaml_obj["metadata"]["labels"]["prefect.io/identifier"] == environment.identifier_label) assert yaml_obj["metadata"]["labels"][ "prefect.io/flow_run_id"] == "id_test" assert (yaml_obj["spec"]["template"]["metadata"]["labels"] ["prefect.io/identifier"] == environment.identifier_label) env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "auth_test" assert env[2]["value"] == "id_test" assert env[3]["value"] == "namespace_test" assert env[4]["value"] == "test1/test2:test3" assert env[9]["value"] == "[]" assert (yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] == "test1/test2:test3") assert yaml_obj["spec"]["template"]["spec"]["containers"][0][ "command"] == [ "/bin/sh", "-c", ] assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == [ "python -c 'import prefect; prefect.environments.KubernetesJobEnvironment().run_flow()'" ]
def test_create_k8s_job_callbacks(): def f(): pass with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml"), labels=["foo"], on_start=f, on_exit=f, ) assert environment.labels == set(["foo"]) assert environment.on_start is f assert environment.on_exit is f
def test_create_flow_run_job(monkeypatch): file_path = os.path.dirname( prefect.environments.execution.dask.k8s.__file__) environment = KubernetesJobEnvironment(path.join(file_path, "job.yaml")) config = MagicMock() monkeypatch.setattr("kubernetes.config", config) batchv1 = MagicMock() monkeypatch.setattr("kubernetes.client", MagicMock(BatchV1Api=MagicMock(return_value=batchv1))) with set_temporary_config({"cloud.auth_token": "test"}): environment.create_flow_run_job(docker_name="test1/test2:test3", flow_file_path="test4") assert (batchv1.create_namespaced_job.call_args[1]["body"]["apiVersion"] == "batch/v1")
def test_populate_job_yaml_no_defaults(job_spec_file, job): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) # only command and args are set on the container when the instance # is initialized job["spec"]["template"]["spec"]["containers"][0] = { "command": ["/bin/sh", "-c"], "args": default_command_args, } del job["metadata"] del job["spec"]["template"]["metadata"] environment._job_spec = job with set_temporary_config( {"cloud.graphql": "gql_test", "cloud.auth_token": "auth_test"} ): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_run_time_job_spec_details( docker_name="test1/test2:test3", ) assert ( yaml_obj["metadata"]["labels"]["prefect.io/identifier"] == environment.identifier_label ) assert yaml_obj["metadata"]["labels"]["prefect.io/flow_run_id"] == "id_test" assert ( yaml_obj["spec"]["template"]["metadata"]["labels"]["prefect.io/identifier"] == environment.identifier_label ) env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"] assert env[0]["value"] == "gql_test" assert env[1]["value"] == "auth_test" assert env[2]["value"] == "id_test" assert env[3]["value"] == "namespace_test" assert env[4]["value"] == "test1/test2:test3" assert env[9]["value"] == "[]" assert ( yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] == "test1/test2:test3" )
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: 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" }), 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
def test_execute(monkeypatch): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = KubernetesJobEnvironment( job_spec_file=os.path.join(directory, "job.yaml") ) storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") create_flow_run = MagicMock() monkeypatch.setattr( "prefect.environments.KubernetesJobEnvironment.create_flow_run_job", create_flow_run, ) environment.execute(storage=storage, flow_location="") assert create_flow_run.call_args[1]["docker_name"] == "test1/test2:test3"
def test_execute(monkeypatch): file_path = os.path.dirname(prefect.environments.execution.dask.k8s.__file__) environment = KubernetesJobEnvironment(os.path.join(file_path, "job.yaml")) config = MagicMock() monkeypatch.setattr("kubernetes.config", config) batchv1 = MagicMock() monkeypatch.setattr( "kubernetes.client", MagicMock(BatchV1Api=MagicMock(return_value=batchv1)) ) storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") with set_temporary_config({"cloud.auth_token": "test"}): environment.execute(Flow("test", storage=storage)) assert ( batchv1.create_namespaced_job.call_args[1]["body"]["apiVersion"] == "batch/v1" )
def _get_environment( self, cluster_kwargs: Dict[str, Any], adapt_kwargs: Dict[str, Any], ): """ Get an environment that customizes the execution of a Prefect flow run. """ local_tmp_file = "/tmp/prefect-flow-run.yaml" with open(local_tmp_file, "w") as f: yaml.dump(self._flow_run_job_spec, stream=f, Dumper=yaml.RoundTripDumper) # saturn_flow_id is used by Saturn's custom Prefect agent k8s_environment = KubernetesJobEnvironment( metadata={"saturn_flow_id": self.flow_id, "image": self.image}, executor=DaskExecutor( cluster_class="dask_saturn.SaturnCluster", cluster_kwargs=cluster_kwargs, adapt_kwargs=adapt_kwargs, ), job_spec_file=local_tmp_file, labels=self._saturn_flow_labels, unique_job_name=True, ) # patch command and args to run the user's start script new_command = ["/bin/bash", "-ec"] k8s_environment._job_spec["spec"]["template"]["spec"]["containers"][0][ "command" ] = new_command args_from_prefect = k8s_environment._job_spec["spec"]["template"]["spec"]["containers"][ 0 ].get("args", []) args_from_prefect = " ".join(args_from_prefect) new_args = f"source /home/jovyan/.saturn/start_wrapper.sh; {args_from_prefect}" k8s_environment._job_spec["spec"]["template"]["spec"]["containers"][0]["args"] = [new_args] return k8s_environment
def test_environment_run(): class MyExecutor(LocalDaskExecutor): submit_called = False def submit(self, *args, **kwargs): self.submit_called = True return super().submit(*args, **kwargs) global_dict = {} @prefect.task def add_to_dict(): global_dict["run"] = True executor = MyExecutor() environment = KubernetesJobEnvironment(executor=executor) flow = prefect.Flow("test", tasks=[add_to_dict], environment=environment) environment.run(flow=flow) assert global_dict.get("run") is True assert executor.submit_called
def test_populate_job_yaml_command_and_args_not_overridden_at_run_time(job_spec_file): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) test_command = ["/bin/bash", "-acdefg"] test_args = "echo 'hello'; python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'" environment._job_spec["spec"]["template"]["spec"]["containers"][0][ "command" ] = test_command environment._job_spec["spec"]["template"]["spec"]["containers"][0][ "args" ] = test_args with set_temporary_config( {"cloud.graphql": "gql_test", "cloud.auth_token": "auth_test"} ): with prefect.context(flow_run_id="id_test", namespace="namespace_test"): yaml_obj = environment._populate_run_time_job_spec_details( docker_name="test1/test2:test3", ) assert ( yaml_obj["spec"]["template"]["spec"]["containers"][0]["command"] == test_command ) assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == test_args