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)