def test_k8s_agent_removes_yaml_no_volume(monkeypatch, cloud_api):
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent()
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
assert not job["spec"]["template"]["spec"].get("volumes", None)
assert not job["spec"]["template"]["spec"]["containers"][0].get(
"volumeMounts", None)
def test_k8s_agent_includes_agent_labels_in_job(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent(labels=["foo", "bar"])
job = agent.generate_job_spec_from_environment(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, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
agent = KubernetesAgent()
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
assert not job["spec"]["template"]["spec"].get("imagePullSecrets", None)
def test_k8s_agent_replace_yaml_responds_to_logging_config(
monkeypatch, cloud_api, flag):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
"name":
"name",
})
agent = KubernetesAgent(no_cloud_logs=flag)
job = agent.generate_job_spec_from_environment(flow_run,
image="test/name:tag")
env = job["spec"]["template"]["spec"]["containers"][0]["env"]
assert env[6]["value"] == str(not flag).lower()
def test_k8s_agent_replace_yaml_respects_multiple_image_secrets(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
monkeypatch.setenv("IMAGE_PULL_SECRETS", "some-secret,other-secret")
monkeypatch.setenv("IMAGE_PULL_POLICY", "custom_policy")
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Docker(
registry_url="test", image_name="name", image_tag="tag"
).serialize(),
"environment": LocalEnvironment().serialize(),
"id": "new_id",
"core_version": "0.13.0",
}
),
"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.generate_job_spec_from_environment(flow_run, image="test/name:tag")
expected_secrets = [{"name": "some-secret"}, {"name": "other-secret"}]
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == expected_secrets
def test_k8s_agent_replace_yaml(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
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":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
with set_temporary_config({
"cloud.agent.auth_token": "token",
"logging.log_to_cloud": True
}):
volume_mounts = [{"name": "my-vol", "mountPath": "/mnt/my-mount"}]
volumes = [{"name": "my-vol", "hostPath": "/host/folder"}]
agent = KubernetesAgent(volume_mounts=volume_mounts, volumes=volumes)
job = agent.generate_job_spec_from_environment(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"
volumeMounts = job["spec"]["template"]["spec"]["containers"][0][
"volumeMounts"]
assert volumeMounts[0]["name"] == "my-vol"
assert volumeMounts[0]["mountPath"] == "/mnt/my-mount"
assert (job["spec"]["template"]["spec"]["containers"][0]
["imagePullPolicy"] == "IfNotPresent")
volumes = job["spec"]["template"]["spec"]["volumes"]
assert volumes[0]["name"] == "my-vol"
assert volumes[0]["hostPath"] == "/host/folder"
assert job["spec"]["template"]["spec"].get("serviceAccountName",
None) is None
def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, cloud_api):
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
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")
monkeypatch.setenv("IMAGE_PULL_POLICY", "custom_policy")
monkeypatch.setenv("SERVICE_ACCOUNT_NAME", "svc_name")
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Docker(registry_url="test", image_name="name",
image_tag="tag").serialize(),
"environment":
LocalEnvironment().serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"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.generate_job_spec_from_environment(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
assert (job["spec"]["template"]["spec"]["containers"][0]
["imagePullPolicy"] == "custom_policy")
assert job["spec"]["template"]["spec"][
"serviceAccountName"] == "svc_name"
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [{
"name":
"my-secret"
}]
