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()'"
]
