def __init__(self, kube_config, task_queue, result_queue, kube_client, worker_uuid): self.log.debug("Creating Kubernetes executor") self.kube_config = kube_config self.task_queue = task_queue self.result_queue = result_queue self.namespace = self.kube_config.kube_namespace self.log.debug("Kubernetes using namespace %s", self.namespace) self.kube_client = kube_client self.launcher = PodLauncher(kube_client=self.kube_client) self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config) self._manager = multiprocessing.Manager() self.watcher_queue = self._manager.Queue() self.worker_uuid = worker_uuid self.kube_watcher = self._make_kube_watcher()
def __init__(self, kube_config: Any, task_queue: 'Queue[KubernetesJobType]', result_queue: 'Queue[KubernetesResultsType]', kube_client: client.CoreV1Api, scheduler_job_id: str): super().__init__() self.log.debug("Creating Kubernetes executor") self.kube_config = kube_config self.task_queue = task_queue self.result_queue = result_queue self.namespace = self.kube_config.kube_namespace self.log.debug("Kubernetes using namespace %s", self.namespace) self.kube_client = kube_client self.launcher = PodLauncher(kube_client=self.kube_client) self._manager = multiprocessing.Manager() self.watcher_queue = self._manager.Queue() self.scheduler_job_id = scheduler_job_id self.kube_watcher = self._make_kube_watcher()
def setUp(self): self.mock_kube_client = mock.Mock() self.pod_launcher = PodLauncher(kube_client=self.mock_kube_client)
def test_pod_mutation_v1_pod(self): with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK_V1_POD, "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) pod = pod_generator.PodGenerator(image="myimage", cmds=["foo"], namespace="baz", volume_mounts=[{ "name": "foo", "mountPath": "/mnt", "subPath": "/", "readOnly": True }], volumes=[{ "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': { 'namespace': 'baz' }, 'spec': { 'containers': [{ 'args': [], 'command': ['foo'], 'env': [], 'envFrom': [], 'image': 'myimage', 'name': 'base', 'ports': [], 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'name': 'foo' }] } }) # 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': { 'namespace': 'airflow-tests' }, 'spec': { 'containers': [{ 'args': [], 'command': ['foo'], 'env': [{ 'name': 'TEST_USER', 'value': 'ADMIN' }], 'envFrom': [], 'image': 'test-image', 'name': 'base', 'ports': [{ 'containerPort': 8080 }, { 'containerPort': 8081 }], 'volumeMounts': [{ 'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/' }, { 'mountPath': '/opt/airflow/secrets/', 'name': 'airflow-secrets-mount', 'readOnly': True }] }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'name': 'foo' }, { 'name': 'airflow-secrets-mount', 'secret': { 'secretName': 'airflow-test-secrets' } }] } })
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 } } })