def test_generate_job_spec_image_pull_secrets_empty_string_in_runconfig(
self, tmpdir):
"""Regression test for issue #5001."""
run_config = KubernetesRun(image_pull_secrets="")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert "imagePullSecrets" not in job["spec"]["template"]["spec"]
def test_generate_job_spec_image_pull_secrets_empty_string_in_env(
self, tmpdir, monkeypatch):
"""Regression test for issue #5001."""
run_config = KubernetesRun()
monkeypatch.setenv("IMAGE_PULL_SECRETS", "")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert "imagePullSecrets" not in job["spec"]["template"]["spec"]
def test_generate_job_spec_image_pull_secrets_from_env(
self, tmpdir, monkeypatch):
run_config = KubernetesRun()
monkeypatch.setenv("IMAGE_PULL_SECRETS", "in-env")
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [{
"name":
"in-env"
}]
def test_environment_has_api_key_from_config(self, config_with_api_key):
"""Check that the API key is passed through from the config via environ"""
flow_run = self.build_flow_run(KubernetesRun())
agent = KubernetesAgent(namespace="testing", )
job = agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env[
"PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def test_environment_has_tenant_id_from_server(self, config_with_api_key):
"""Check that the API key is passed through from the config via environ"""
flow_run = self.build_flow_run(KubernetesRun())
tenant_id = uuid.uuid4()
with set_temporary_config({"cloud.tenant_id": None}):
agent = KubernetesAgent(namespace="testing")
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
job = agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
def test_environment_has_api_key_from_disk(self, monkeypatch):
"""Check that the API key is passed through from the on disk cache"""
tenant_id = str(uuid.uuid4())
monkeypatch.setattr(
"prefect.Client.load_auth_from_disk",
MagicMock(return_value={
"api_key": "TEST_KEY",
"tenant_id": tenant_id
}),
)
flow_run = self.build_flow_run(KubernetesRun())
agent = KubernetesAgent(namespace="testing", )
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
job = agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env["PREFECT__CLOUD__API_KEY"] == "TEST_KEY"
assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_KEY"
assert env["PREFECT__CLOUD__TENANT_ID"] == tenant_id
class TestK8sAgentRunConfig:
def setup(self):
self.agent = KubernetesAgent(namespace="testing", )
def read_default_template(self):
from prefect.agent.kubernetes.agent import DEFAULT_JOB_TEMPLATE_PATH
with open(DEFAULT_JOB_TEMPLATE_PATH) as f:
return yaml.safe_load(f)
def build_flow_run(self, config, storage=None):
if storage is None:
storage = Local()
return GraphQLResult({
"flow":
GraphQLResult({
"storage": storage.serialize(),
"run_config": RunConfigSchema().dump(config),
"id": "new_id",
"core_version": "0.13.0",
}),
"id":
"id",
})
def test_generate_job_spec_uses_job_template_provided_in_run_config(self):
template = self.read_default_template()
labels = template.setdefault("metadata", {}).setdefault("labels", {})
labels["TEST"] = "VALUE"
flow_run = self.build_flow_run(KubernetesRun(job_template=template))
job = self.agent.generate_job_spec(flow_run)
assert job["metadata"]["labels"]["TEST"] == "VALUE"
def test_generate_job_spec_uses_job_template_path_provided_in_run_config(
self, tmpdir, monkeypatch):
path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
labels = template.setdefault("metadata", {}).setdefault("labels", {})
labels["TEST"] = "VALUE"
with open(path, "w") as f:
yaml.safe_dump(template, f)
template_path = f"agent://{path}"
flow_run = self.build_flow_run(
KubernetesRun(job_template_path=template_path))
mocked_read_bytes = MagicMock(wraps=read_bytes_from_path)
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.read_bytes_from_path",
mocked_read_bytes)
job = self.agent.generate_job_spec(flow_run)
assert job["metadata"]["labels"]["TEST"] == "VALUE"
assert mocked_read_bytes.call_args[0] == (template_path, )
def test_generate_job_spec_metadata(self, tmpdir):
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
job_labels = template.setdefault("metadata",
{}).setdefault("labels", {})
pod_labels = (template["spec"]["template"].setdefault(
"metadata", {}).setdefault("labels", {}))
job_labels.update({"JOB_LABEL": "VALUE1"})
pod_labels.update({"POD_LABEL": "VALUE2"})
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.job_template_path = template_path
flow_run = self.build_flow_run(KubernetesRun())
job = self.agent.generate_job_spec(flow_run)
identifier = job["metadata"]["labels"]["prefect.io/identifier"]
labels = {
"prefect.io/identifier": identifier,
"prefect.io/flow_run_id": flow_run.id,
"prefect.io/flow_id": flow_run.flow.id,
}
assert job["metadata"]["name"]
assert job["metadata"]["labels"] == dict(JOB_LABEL="VALUE1", **labels)
assert job["spec"]["template"]["metadata"]["labels"] == dict(
POD_LABEL="VALUE2", **labels)
@pytest.mark.parametrize(
"run_config, storage, expected",
[
(
KubernetesRun(),
Docker(registry_url="test", image_name="name",
image_tag="tag"),
"test/name:tag",
),
(KubernetesRun(image="myimage"), Local(), "myimage"),
(KubernetesRun(), Local(), "prefecthq/prefect:all_extras-0.13.0"),
],
ids=["on-storage", "on-run_config", "default"],
)
def test_generate_job_spec_image(self, run_config, storage, expected):
flow_run = self.build_flow_run(run_config, storage)
job = self.agent.generate_job_spec(flow_run)
image = job["spec"]["template"]["spec"]["containers"][0]["image"]
assert image == expected
def test_generate_job_spec_environment_variables(self, tmpdir):
"""Check that environment variables are set in precedence order
- CUSTOM1 & CUSTOM2 are set on the template
- CUSTOM2 & CUSTOM3 are set on the agent
- CUSTOM3 & CUSTOM4 are set on the RunConfig
"""
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
template_env = template["spec"]["template"]["spec"]["containers"][
0].setdefault("env", [])
template_env.extend([
{
"name": "CUSTOM1",
"value": "VALUE1"
},
{
"name": "CUSTOM2",
"value": "VALUE2"
},
])
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.job_template_path = template_path
self.agent.env_vars = {"CUSTOM2": "OVERRIDE2", "CUSTOM3": "VALUE3"}
run_config = KubernetesRun(image="test-image",
env={
"CUSTOM3": "OVERRIDE3",
"CUSTOM4": "VALUE4"
})
flow_run = self.build_flow_run(run_config)
job = self.agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env == {
"PREFECT__CLOUD__API": prefect.config.cloud.api,
"PREFECT__CLOUD__AUTH_TOKEN":
prefect.config.cloud.agent.auth_token,
"PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
"PREFECT__CONTEXT__FLOW_RUN_ID": flow_run.id,
"PREFECT__CONTEXT__FLOW_ID": flow_run.flow.id,
"PREFECT__CONTEXT__IMAGE": "test-image",
"PREFECT__LOGGING__LOG_TO_CLOUD":
str(self.agent.log_to_cloud).lower(),
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS":
"prefect.engine.cloud.CloudTaskRunner",
"CUSTOM1": "VALUE1",
"CUSTOM2":
"OVERRIDE2", # Agent env-vars override those in template
"CUSTOM3":
"OVERRIDE3", # RunConfig env-vars override those on agent and template
"CUSTOM4": "VALUE4",
}
def test_generate_job_spec_resources(self):
flow_run = self.build_flow_run(
KubernetesRun(cpu_request=1,
cpu_limit=2,
memory_request="4G",
memory_limit="8G"))
job = self.agent.generate_job_spec(flow_run)
resources = job["spec"]["template"]["spec"]["containers"][0][
"resources"]
assert resources == {
"limits": {
"cpu": "2",
"memory": "8G"
},
"requests": {
"cpu": "1",
"memory": "4G"
},
}
class TestK8sAgentRunConfig:
def setup(self):
self.agent = KubernetesAgent(
namespace="testing",
)
def read_default_template(self):
from prefect.agent.kubernetes.agent import DEFAULT_JOB_TEMPLATE_PATH
with open(DEFAULT_JOB_TEMPLATE_PATH) as f:
return yaml.safe_load(f)
def build_flow_run(self, config, storage=None, core_version="0.13.0"):
if storage is None:
storage = Local()
return GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": storage.serialize(),
"id": "new_id",
"core_version": core_version,
}
),
"run_config": None if config is None else config.serialize(),
"id": "id",
}
)
@pytest.mark.parametrize("run_config", [None, UniversalRun()])
def test_generate_job_spec_null_or_univeral_run_config(self, run_config):
self.agent.generate_job_spec_from_run_config = MagicMock(
wraps=self.agent.generate_job_spec_from_run_config
)
flow_run = self.build_flow_run(run_config)
self.agent.generate_job_spec(flow_run)
assert self.agent.generate_job_spec_from_run_config.called
def test_generate_job_spec_errors_if_non_kubernetesrun_run_config(self):
with pytest.raises(
TypeError,
match="`run_config` of type `LocalRun`, only `KubernetesRun` is supported",
):
self.agent.generate_job_spec(self.build_flow_run(LocalRun()))
def test_generate_job_spec_uses_job_template_provided_in_run_config(self):
template = self.read_default_template()
labels = template.setdefault("metadata", {}).setdefault("labels", {})
labels["TEST"] = "VALUE"
flow_run = self.build_flow_run(KubernetesRun(job_template=template))
job = self.agent.generate_job_spec(flow_run)
assert job["metadata"]["labels"]["TEST"] == "VALUE"
def test_generate_job_spec_uses_job_template_path_provided_in_run_config(
self, tmpdir, monkeypatch
):
path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
labels = template.setdefault("metadata", {}).setdefault("labels", {})
labels["TEST"] = "VALUE"
with open(path, "w") as f:
yaml.safe_dump(template, f)
template_path = f"agent://{path}"
flow_run = self.build_flow_run(KubernetesRun(job_template_path=template_path))
mocked_read_bytes = MagicMock(wraps=read_bytes_from_path)
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.read_bytes_from_path", mocked_read_bytes
)
job = self.agent.generate_job_spec(flow_run)
assert job["metadata"]["labels"]["TEST"] == "VALUE"
assert mocked_read_bytes.call_args[0] == (template_path,)
def test_generate_job_spec_metadata(self, tmpdir):
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
job_labels = template.setdefault("metadata", {}).setdefault("labels", {})
pod_labels = (
template["spec"]["template"]
.setdefault("metadata", {})
.setdefault("labels", {})
)
job_labels.update({"JOB_LABEL": "VALUE1"})
pod_labels.update({"POD_LABEL": "VALUE2"})
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.job_template_path = template_path
flow_run = self.build_flow_run(KubernetesRun())
job = self.agent.generate_job_spec(flow_run)
identifier = job["metadata"]["labels"]["prefect.io/identifier"]
labels = {
"prefect.io/identifier": identifier,
"prefect.io/flow_run_id": flow_run.id,
"prefect.io/flow_id": flow_run.flow.id,
}
assert job["metadata"]["name"]
assert job["metadata"]["labels"] == dict(JOB_LABEL="VALUE1", **labels)
assert job["spec"]["template"]["metadata"]["labels"] == dict(
POD_LABEL="VALUE2", **labels
)
assert job["spec"]["template"]["spec"]["restartPolicy"] == "Never"
@pytest.mark.parametrize(
"run_config, storage, expected",
[
(
KubernetesRun(),
Docker(registry_url="test", image_name="name", image_tag="tag"),
"test/name:tag",
),
(KubernetesRun(image="myimage"), Local(), "myimage"),
(KubernetesRun(), Local(), "prefecthq/prefect:0.13.0"),
],
ids=["on-storage", "on-run_config", "default"],
)
def test_generate_job_spec_image(self, run_config, storage, expected):
flow_run = self.build_flow_run(run_config, storage)
job = self.agent.generate_job_spec(flow_run)
image = job["spec"]["template"]["spec"]["containers"][0]["image"]
assert image == expected
@pytest.mark.parametrize(
"core_version, expected",
[
("0.12.0", "prefect execute cloud-flow"),
("0.14.0", "prefect execute flow-run"),
],
)
def test_generate_job_spec_container_args(self, core_version, expected):
flow_run = self.build_flow_run(KubernetesRun(), core_version=core_version)
job = self.agent.generate_job_spec(flow_run)
args = job["spec"]["template"]["spec"]["containers"][0]["args"]
assert args == expected.split()
def test_generate_job_spec_environment_variables(self, tmpdir):
"""Check that environment variables are set in precedence order
- CUSTOM1 & CUSTOM2 are set on the template
- CUSTOM2 & CUSTOM3 are set on the agent
- CUSTOM3 & CUSTOM4 are set on the RunConfig
"""
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
template_env = template["spec"]["template"]["spec"]["containers"][0].setdefault(
"env", []
)
template_env.extend(
[
{"name": "CUSTOM1", "value": "VALUE1"},
{"name": "CUSTOM2", "value": "VALUE2"},
]
)
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.job_template_path = template_path
self.agent.env_vars = {"CUSTOM2": "OVERRIDE2", "CUSTOM3": "VALUE3"}
run_config = KubernetesRun(
image="test-image", env={"CUSTOM3": "OVERRIDE3", "CUSTOM4": "VALUE4"}
)
flow_run = self.build_flow_run(run_config)
job = self.agent.generate_job_spec(flow_run)
env_list = job["spec"]["template"]["spec"]["containers"][0]["env"]
env = {item["name"]: item["value"] for item in env_list}
assert env == {
"PREFECT__CLOUD__API": prefect.config.cloud.api,
"PREFECT__CLOUD__AUTH_TOKEN": prefect.config.cloud.agent.auth_token,
"PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
"PREFECT__CONTEXT__FLOW_RUN_ID": flow_run.id,
"PREFECT__CONTEXT__FLOW_ID": flow_run.flow.id,
"PREFECT__CONTEXT__IMAGE": "test-image",
"PREFECT__LOGGING__LOG_TO_CLOUD": str(self.agent.log_to_cloud).lower(),
"PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
"PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
"CUSTOM1": "VALUE1",
"CUSTOM2": "OVERRIDE2", # Agent env-vars override those in template
"CUSTOM3": "OVERRIDE3", # RunConfig env-vars override those on agent and template
"CUSTOM4": "VALUE4",
}
def test_generate_job_spec_resources(self):
flow_run = self.build_flow_run(
KubernetesRun(
cpu_request=1, cpu_limit=2, memory_request="4G", memory_limit="8G"
)
)
job = self.agent.generate_job_spec(flow_run)
resources = job["spec"]["template"]["spec"]["containers"][0]["resources"]
assert resources == {
"limits": {"cpu": "2", "memory": "8G"},
"requests": {"cpu": "1", "memory": "4G"},
}
def test_generate_job_spec_service_account_name(self, tmpdir):
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
template["spec"]["template"]["spec"]["serviceAccountName"] = "on-agent-template"
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.service_account_name = "on-agent"
self.agent.job_template_path = template_path
template["spec"]["template"]["spec"][
"serviceAccountName"
] = "on-run-config-template"
run_config = KubernetesRun(
job_template=template, service_account_name="on-run-config"
)
# Check precedence order:
# 1. Explicit on run-config"
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["serviceAccountName"] == "on-run-config"
# 2. In job template on run-config
run_config.service_account_name = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert (
job["spec"]["template"]["spec"]["serviceAccountName"]
== "on-run-config-template"
)
# None in run-config job template is still used
run_config.job_template["spec"]["template"]["spec"]["serviceAccountName"] = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["serviceAccountName"] is None
# 3. Explicit on agent
# Not present in job template
run_config.job_template["spec"]["template"]["spec"].pop("serviceAccountName")
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["serviceAccountName"] == "on-agent"
# No job template present
run_config.job_template = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["serviceAccountName"] == "on-agent"
# 4. In job template on agent
self.agent.service_account_name = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert (
job["spec"]["template"]["spec"]["serviceAccountName"] == "on-agent-template"
)
def test_generate_job_spec_image_pull_secrets(self, tmpdir):
template_path = str(tmpdir.join("job.yaml"))
template = self.read_default_template()
template["spec"]["template"]["spec"]["imagePullSecrets"] = [
{"name": "on-agent-template"}
]
with open(template_path, "w") as f:
yaml.safe_dump(template, f)
self.agent.image_pull_secrets = ["on-agent"]
self.agent.job_template_path = template_path
template["spec"]["template"]["spec"]["imagePullSecrets"] = [
{"name": "on-run-config-template"}
]
run_config = KubernetesRun(
job_template=template, image_pull_secrets=["on-run-config"]
)
# Check precedence order:
# 1. Explicit on run-config"
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
{"name": "on-run-config"}
]
# 2. In job template on run-config
run_config.image_pull_secrets = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
{"name": "on-run-config-template"}
]
# None in run-config job template is still used
run_config.job_template["spec"]["template"]["spec"]["imagePullSecrets"] = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] is None
# 3. Explicit on agent
# Not present in job template
run_config.job_template["spec"]["template"]["spec"].pop("imagePullSecrets")
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
{"name": "on-agent"}
]
# No job template present
run_config.job_template = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
{"name": "on-agent"}
]
# 4. In job template on agent
self.agent.image_pull_secrets = None
job = self.agent.generate_job_spec(self.build_flow_run(run_config))
assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
{"name": "on-agent-template"}
]
def test_generate_job_spec_without_image_pull_secrets(self, tmpdir):
run_config = KubernetesRun()
agent = KubernetesAgent(namespace="testing")
job = agent.generate_job_spec(self.build_flow_run(run_config))
assert "imagePullSecrets" not in job["spec"]["template"]["spec"]
