def test_populate_job_yaml(): 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"), unique_job_name=True) 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]["env"] = [] 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_job_spec_yaml( yaml_obj=job, docker_name="test1/test2:test3", flow_file_path="test4", ) assert "prefect-dask-job-" in yaml_obj["metadata"]["name"] assert len(yaml_obj["metadata"]["name"]) == 25 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"] == "['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"] == [ "python -c 'import prefect; prefect.Flow.load(prefect.context.flow_file_path).environment.run_flow()'" ]
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", ) 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_populate_job_yaml(job_spec_file, job): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file, unique_job_name=True) job["spec"]["template"]["spec"]["containers"][0]["env"] = [] 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_job_spec_yaml( yaml_obj=job, 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"] == [ "python -c 'import prefect; prefect.environments.KubernetesJobEnvironment().run_flow()'" ]
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.execution.load_and_run_flow()'" ]
def test_populate_job_yaml_multiple_containers(job_spec_file, job): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) # Generate yaml object with multiple containers job["spec"]["template"]["spec"]["containers"][0]["env"] = [] job["spec"]["template"]["spec"]["containers"].append( copy.deepcopy(job["spec"]["template"]["spec"]["containers"][0])) job["spec"]["template"]["spec"]["containers"][1]["env"] = [] 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_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) # Assert First Container 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"] == [ "python -c 'import prefect; prefect.environments.KubernetesJobEnvironment().run_flow()'" ] # Assert Second Container env = yaml_obj["spec"]["template"]["spec"]["containers"][1]["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"][1]["image"] != "test1/test2:test3") assert yaml_obj["spec"]["template"]["spec"]["containers"][1]["args"] != [ "python -c 'import prefect; prefect.environments.KubernetesJobEnvironment().run_flow()'" ]
def test_populate_job_yaml_multiple_containers(): 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) # Generate yaml object with multiple containers job["spec"]["template"]["spec"]["containers"][0]["env"] = [] job["spec"]["template"]["spec"]["containers"].append( copy.deepcopy(job["spec"]["template"]["spec"]["containers"][0]) ) job["spec"]["template"]["spec"]["containers"][1]["env"] = [] 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 ) # Assert First Container 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 ( 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 'from prefect.environments import KubernetesJobEnvironment; KubernetesJobEnvironment().run_flow()'" ] # Assert Second Container env = yaml_obj["spec"]["template"]["spec"]["containers"][1]["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 ( yaml_obj["spec"]["template"]["spec"]["containers"][1]["image"] != "test1/test2:test3" ) assert yaml_obj["spec"]["template"]["spec"]["containers"][1]["args"] != [ "python -c 'from prefect.environments import KubernetesJobEnvironment; KubernetesJobEnvironment().run_flow()'" ]