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" }]