def create_security_context(volume_disabled: bool,
volume_owner: int) -> k8s.V1PodSecurityContext:
"""
Creates security context based on volume information
:param volume_disabled:
:param volume_owner:
:return:
"""
return (k8s.V1PodSecurityContext(fs_group=volume_owner)
if not volume_disabled else k8s.V1PodSecurityContext())
def test_gen_pod(self, mock_uuid):
mock_uuid.return_value = '0'
pod_generator = PodGenerator(
labels={'app': 'myapp'},
name='myapp-pod',
image_pull_secrets='pull_secret_a,pull_secret_b',
image='busybox',
envs=self.envs,
cmds=['sh', '-c', 'echo Hello Kubernetes!'],
security_context=k8s.V1PodSecurityContext(
run_as_user=1000,
fs_group=2000,
),
namespace='default',
ports=[k8s.V1ContainerPort(name='foo', container_port=1234)],
configmaps=['configmap_a', 'configmap_b'])
result = pod_generator.gen_pod()
result = append_to_pod(result, self.secrets)
result = self.resources.attach_to_pod(result)
result_dict = self.k8s_client.sanitize_for_serialization(result)
# sort
result_dict['spec']['containers'][0]['env'].sort(
key=lambda x: x['name'])
result_dict['spec']['containers'][0]['envFrom'].sort(
key=lambda x: list(x.values())[0]['name'])
self.assertDictEqual(result_dict, self.expected)
def test_gen_pod_extract_xcom(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
pod_generator = PodGenerator(
labels={'app': 'myapp'},
name='myapp-pod',
image_pull_secrets='pull_secret_a,pull_secret_b',
image='busybox',
envs=self.envs,
cmds=['sh', '-c', 'echo Hello Kubernetes!'],
namespace='default',
security_context=k8s.V1PodSecurityContext(
run_as_user=1000,
fs_group=2000,
),
ports=[k8s.V1ContainerPort(name='foo', container_port=1234)],
configmaps=['configmap_a', 'configmap_b'],
extract_xcom=True)
result = pod_generator.gen_pod()
result = append_to_pod(result, self.secrets)
result = self.resources.attach_to_pod(result)
result_dict = self.k8s_client.sanitize_for_serialization(result)
container_two = {
'name': 'airflow-xcom-sidecar',
'image': "alpine",
'command': ['sh', '-c', PodDefaults.XCOM_CMD],
'volumeMounts': [{
'name': 'xcom',
'mountPath': '/airflow/xcom'
}],
'resources': {
'requests': {
'cpu': '1m'
}
},
}
self.expected['spec']['containers'].append(container_two)
self.expected['spec']['containers'][0]['volumeMounts'].insert(
0, {
'name': 'xcom',
'mountPath': '/airflow/xcom'
})
self.expected['spec']['volumes'].insert(0, {
'name': 'xcom',
'emptyDir': {}
})
result_dict['spec']['containers'][0]['env'].sort(
key=lambda x: x['name'])
self.assertEqual(result_dict, self.expected)
def _get_security_context(self) -> k8s.V1PodSecurityContext:
"""Defines the security context"""
security_context = k8s.V1PodSecurityContext()
if self.kube_config.worker_run_as_user != "":
security_context.run_as_user = self.kube_config.worker_run_as_user
if self.kube_config.worker_fs_group != "":
security_context.fs_group = self.kube_config.worker_fs_group
# set fs_group to 65533 if not explicitly specified and using git ssh keypair auth
if self.kube_config.git_ssh_key_secret_name and security_context.fs_group is None:
security_context.fs_group = 65533
return security_context
def setUp(self):
self.static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48')
self.deserialize_result = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {'name': 'memory-demo', 'namespace': 'mem-example'},
'spec': {
'containers': [
{
'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'image': 'apache/airflow:stress-2020.07.10-1.0.4',
'name': 'memory-demo-ctr',
'resources': {'limits': {'memory': '200Mi'}, 'requests': {'memory': '100Mi'}},
}
]
},
}
self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'}
self.secrets = [
# This should be a secretRef
Secret('env', None, 'secret_a'),
# This should be a single secret mounted in volumeMounts
Secret('volume', '/etc/foo', 'secret_b'),
# This should produce a single secret mounted in env
Secret('env', 'TARGET', 'secret_b', 'source_b'),
]
self.execution_date = parser.parse('2020-08-24 00:00:00.000000')
self.execution_date_label = datetime_to_label_safe_datestring(self.execution_date)
self.dag_id = 'dag_id'
self.task_id = 'task_id'
self.try_number = 3
self.labels = {
'airflow-worker': 'uuid',
'dag_id': self.dag_id,
'execution_date': self.execution_date_label,
'task_id': self.task_id,
'try_number': str(self.try_number),
'airflow_version': __version__.replace('+', '-'),
'kubernetes_executor': 'True',
}
self.annotations = {
'dag_id': self.dag_id,
'task_id': self.task_id,
'execution_date': self.execution_date.isoformat(),
'try_number': str(self.try_number),
}
self.metadata = {
'labels': self.labels,
'name': 'pod_id-' + self.static_uuid.hex,
'namespace': 'namespace',
'annotations': self.annotations,
}
self.resources = k8s.V1ResourceRequirements(
requests={
"cpu": 1,
"memory": "1Gi",
"ephemeral-storage": "2Gi",
},
limits={"cpu": 2, "memory": "2Gi", "ephemeral-storage": "4Gi", 'nvidia.com/gpu': 1},
)
self.k8s_client = ApiClient()
self.expected = k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
namespace="default",
name='myapp-pod-' + self.static_uuid.hex,
labels={'app': 'myapp'},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name='base',
image='busybox',
command=['sh', '-c', 'echo Hello Kubernetes!'],
env=[
k8s.V1EnvVar(name='ENVIRONMENT', value='prod'),
k8s.V1EnvVar(
name="LOG_LEVEL",
value='warning',
),
k8s.V1EnvVar(
name='TARGET',
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(name='secret_b', key='source_b')
),
),
],
env_from=[
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_a')),
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_b')),
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name='secret_a')),
],
ports=[k8s.V1ContainerPort(name="foo", container_port=1234)],
resources=k8s.V1ResourceRequirements(
requests={'memory': '100Mi'},
limits={
'memory': '200Mi',
},
),
)
],
security_context=k8s.V1PodSecurityContext(
fs_group=2000,
run_as_user=1000,
),
host_network=True,
image_pull_secrets=[
k8s.V1LocalObjectReference(name="pull_secret_a"),
k8s.V1LocalObjectReference(name="pull_secret_b"),
],
),
)
def test_pod_template_file_override_in_executor_config(
self, mock_get_kube_client, mock_run_pod_async):
current_folder = pathlib.Path(__file__).parent.absolute()
template_file = str(
(current_folder / "kubernetes_executor_template_files" /
"basic_template.yaml").absolute())
mock_kube_client = mock.patch('kubernetes.client.CoreV1Api',
autospec=True)
mock_get_kube_client.return_value = mock_kube_client
with conf_vars({('kubernetes', 'pod_template_file'): ''}):
executor = self.kubernetes_executor
executor.start()
assert executor.event_buffer == {}
assert executor.task_queue.empty()
execution_date = datetime.utcnow()
executor.execute_async(
key=('dag', 'task', execution_date, 1),
queue=None,
command=['airflow', 'tasks', 'run', 'true', 'some_parameter'],
executor_config={
"pod_template_file":
template_file,
"pod_override":
k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
labels={"release": "stable"}),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name="base", image="airflow:3.6")
], ),
),
},
)
assert not executor.task_queue.empty()
task = executor.task_queue.get_nowait()
_, _, expected_executor_config, expected_pod_template_file = task
# Test that the correct values have been put to queue
assert expected_executor_config.metadata.labels == {
'release': 'stable'
}
assert expected_pod_template_file == template_file
self.kubernetes_executor.kube_scheduler.run_next(task)
mock_run_pod_async.assert_called_once_with(
k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
name=mock.ANY,
namespace="default",
annotations={
'dag_id': 'dag',
'execution_date': execution_date.isoformat(),
'task_id': 'task',
'try_number': '1',
},
labels={
'airflow-worker':
'5',
'airflow_version':
mock.ANY,
'dag_id':
'dag',
'execution_date':
datetime_to_label_safe_datestring(execution_date),
'kubernetes_executor':
'True',
'mylabel':
'foo',
'release':
'stable',
'task_id':
'task',
'try_number':
'1',
},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
image="airflow:3.6",
args=[
'airflow', 'tasks', 'run', 'true',
'some_parameter'
],
env=[
k8s.V1EnvVar(
name='AIRFLOW_IS_K8S_EXECUTOR_POD',
value='True')
],
)
],
image_pull_secrets=[
k8s.V1LocalObjectReference(name='airflow-registry')
],
scheduler_name='default-scheduler',
security_context=k8s.V1PodSecurityContext(
fs_group=50000, run_as_user=50000),
),
))
def handlePrivileged(self, key: str, value: bool, full_path: str) -> None:
self.container.securityContext = K.V1PodSecurityContext(
dict(privileged=value))
def test_pod_mutation_to_k8s_pod(self):
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK,
"airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
from airflow.kubernetes.pod_launcher import PodLauncher
self.mock_kube_client = Mock()
self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client)
init_container = k8s.V1Container(name="init-container",
volume_mounts=[
k8s.V1VolumeMount(
mount_path="/tmp",
name="init-secret")
])
pod = pod_generator.PodGenerator(
image="foo",
name="bar",
namespace="baz",
image_pull_policy="Never",
init_containers=[init_container],
cmds=["foo"],
args=["/bin/sh", "-c", "touch /tmp/healthy"],
tolerations=[{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}],
volume_mounts=[{
"name": "foo",
"mountPath": "/mnt",
"subPath": "/",
"readOnly": True
}],
security_context=k8s.V1PodSecurityContext(fs_group=0,
run_as_user=1),
volumes=[k8s.V1Volume(name="foo")]).gen_pod()
sanitized_pod_pre_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_pre_mutation,
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': mock.ANY,
'namespace': 'baz'
},
'spec': {
'containers': [{
'args': ['/bin/sh', '-c', 'touch /tmp/healthy'],
'command': ['foo'],
'env': [],
'envFrom': [],
'image':
'foo',
'imagePullPolicy':
'Never',
'name':
'base',
'ports': [],
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}]
}],
'initContainers': [{
'name':
'init-container',
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'tolerations': [{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}],
'volumes': [{
'name': 'foo'
}],
'securityContext': {
'fsGroup': 0,
'runAsUser': 1
}
}
},
)
# Apply Pod Mutation Hook
pod = self.pod_launcher._mutate_pod_backcompat(pod)
sanitized_pod_post_mutation = api_client.sanitize_for_serialization(
pod)
self.assertEqual(
sanitized_pod_post_mutation, {
"apiVersion": "v1",
"kind": "Pod",
'metadata': {
'labels': {
'test_label': 'test_value'
},
'name': mock.ANY,
'namespace': 'airflow-tests'
},
'spec': {
'affinity': {
'nodeAffinity': {
'requiredDuringSchedulingIgnoredDuringExecution':
{
'nodeSelectorTerms': [{
'matchExpressions': [{
'key': 'test/dynamic-pods',
'operator': 'In',
'values': ['true']
}]
}]
}
}
},
'containers': [{
'args': ['/bin/sh', '-c', 'touch /tmp/healthy2'],
'command': ['foo'],
'env': [{
'name': 'TEST_USER',
'value': 'ADMIN'
}],
'image':
'my_image',
'imagePullPolicy':
'Never',
'name':
'base',
'ports': [{
'containerPort': 8080
}, {
'containerPort': 8081
}],
'resources': {
'limits': {
'nvidia.com/gpu': '200G'
},
'requests': {
'cpu': '200Mi',
'memory': '2G'
}
},
'volumeMounts': [{
'mountPath': '/mnt',
'name': 'foo',
'readOnly': True,
'subPath': '/'
}, {
'mountPath': '/opt/airflow/secrets/',
'name': 'airflow-secrets-mount',
'readOnly': True
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'initContainers': [{
'name':
'init-container',
'securityContext': {
'runAsGroup': 50000,
'runAsUser': 50000
},
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
'tolerations': [{
'effect': 'NoSchedule',
'key': 'static-pods',
'operator': 'Equal',
'value': 'true'
}, {
'effect': 'NoSchedule',
'key': 'dynamic-pods',
'operator': 'Equal',
'value': 'true'
}],
'volumes': [
{
'name': 'airflow-secrets-mount',
'secret': {
'secretName': 'airflow-test-secrets'
}
},
{
'name': 'bar'
},
{
'name': 'foo'
},
],
'securityContext': {
'runAsUser': 1
}
}
})
def test_convert_to_airflow_pod(self):
input_pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name="foo", namespace="bar"),
spec=k8s.V1PodSpec(
init_containers=[
k8s.V1Container(name="init-container",
volume_mounts=[
k8s.V1VolumeMount(mount_path="/tmp",
name="init-secret")
])
],
containers=[
k8s.V1Container(
name="base",
command=["foo"],
image="myimage",
env=[
k8s.V1EnvVar(
name="AIRFLOW_SECRET",
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(
name="ai", key="secret_key")))
],
ports=[
k8s.V1ContainerPort(
name="myport",
container_port=8080,
)
],
volume_mounts=[
k8s.V1VolumeMount(name="myvolume",
mount_path="/tmp/mount",
read_only="True"),
k8s.V1VolumeMount(name='airflow-config',
mount_path='/config',
sub_path='airflow.cfg',
read_only=True),
k8s.V1VolumeMount(name="airflow-secret",
mount_path="/opt/mount",
read_only=True)
])
],
security_context=k8s.V1PodSecurityContext(
run_as_user=0,
fs_group=0,
),
volumes=[
k8s.V1Volume(name="myvolume"),
k8s.V1Volume(
name="airflow-config",
config_map=k8s.V1ConfigMap(data="airflow-data")),
k8s.V1Volume(name="airflow-secret",
secret=k8s.V1SecretVolumeSource(
secret_name="secret-name", )),
k8s.V1Volume(name="init-secret",
secret=k8s.V1SecretVolumeSource(
secret_name="init-secret", ))
]))
result_pod = _convert_to_airflow_pod(input_pod)
expected = Pod(
name="foo",
namespace="bar",
envs={},
init_containers=[{
'name':
'init-container',
'volumeMounts': [{
'mountPath': '/tmp',
'name': 'init-secret'
}]
}],
cmds=["foo"],
image="myimage",
ports=[Port(name="myport", container_port=8080)],
volume_mounts=[
VolumeMount(name="myvolume",
mount_path="/tmp/mount",
sub_path=None,
read_only="True"),
VolumeMount(name="airflow-config",
read_only=True,
mount_path="/config",
sub_path="airflow.cfg"),
VolumeMount(name="airflow-secret",
mount_path="/opt/mount",
sub_path=None,
read_only=True)
],
secrets=[Secret("env", "AIRFLOW_SECRET", "ai", "secret_key")],
security_context={
'fsGroup': 0,
'runAsUser': 0
},
volumes=[
Volume(name="myvolume", configs={'name': 'myvolume'}),
Volume(name="airflow-config",
configs={
'configMap': {
'data': 'airflow-data'
},
'name': 'airflow-config'
}),
Volume(name='airflow-secret',
configs={
'name': 'airflow-secret',
'secret': {
'secretName': 'secret-name'
}
}),
Volume(name='init-secret',
configs={
'name': 'init-secret',
'secret': {
'secretName': 'init-secret'
}
})
],
)
expected_dict = expected.as_dict()
result_dict = result_pod.as_dict()
print(result_pod.volume_mounts)
parsed_configs = self.pull_out_volumes(result_dict)
result_dict['volumes'] = parsed_configs
self.assertEqual(result_dict['secrets'], expected_dict['secrets'])
self.assertDictEqual(expected_dict, result_dict)
