def __init__( self, image, envs, cmds, args=None, secrets=None, labels=None, node_selectors=None, name=None, ports=None, volumes=None, volume_mounts=None, namespace='default', result=None, image_pull_policy='IfNotPresent', image_pull_secrets=None, init_containers=None, service_account_name=None, resources=None, annotations=None, affinity=None, hostnetwork=False, tolerations=None, security_context=None, configmaps=None, pod_runtime_info_envs=None, dnspolicy=None ): warnings.warn( "Using `airflow.contrib.kubernetes.pod.Pod` is deprecated. Please use `k8s.V1Pod`.", DeprecationWarning, stacklevel=2 ) self.image = image self.envs = envs or {} self.cmds = cmds self.args = args or [] self.secrets = secrets or [] self.result = result self.labels = labels or {} self.name = name self.ports = ports or [] self.volumes = volumes or [] self.volume_mounts = volume_mounts or [] self.node_selectors = node_selectors or {} self.namespace = namespace self.image_pull_policy = image_pull_policy self.image_pull_secrets = image_pull_secrets self.init_containers = init_containers self.service_account_name = service_account_name self.resources = resources or Resources() self.annotations = annotations or {} self.affinity = affinity or {} self.hostnetwork = hostnetwork or False self.tolerations = tolerations or [] self.security_context = security_context or {} self.configmaps = configmaps or [] self.pod_runtime_info_envs = pod_runtime_info_envs or [] self.dnspolicy = dnspolicy
def _set_resources(self, resources): if isinstance(resources, Resources): return resources else: inputResource = Resources() if resources: for item in resources.keys(): setattr(inputResource, item, resources[item]) return inputResource
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id, execution_date, try_number, airflow_command, kube_executor_config): volumes_dict, volume_mounts_dict = self._get_volumes_and_mounts() worker_init_container_spec = self._get_init_containers() resources = Resources( request_memory=kube_executor_config.request_memory, request_cpu=kube_executor_config.request_cpu, limit_memory=kube_executor_config.limit_memory, limit_cpu=kube_executor_config.limit_cpu, limit_gpu=kube_executor_config.limit_gpu) gcp_sa_key = kube_executor_config.gcp_service_account_key annotations = dict(kube_executor_config.annotations ) or self.kube_config.kube_annotations if gcp_sa_key: annotations['iam.cloud.google.com/service-account'] = gcp_sa_key volumes = [value for value in volumes_dict.values() ] + kube_executor_config.volumes volume_mounts = [value for value in volume_mounts_dict.values() ] + kube_executor_config.volume_mounts affinity = kube_executor_config.affinity or self.kube_config.kube_affinity tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations return Pod( namespace=namespace, name=pod_id, image=kube_executor_config.image or self.kube_config.kube_image, image_pull_policy=(kube_executor_config.image_pull_policy or self.kube_config.kube_image_pull_policy), cmds=airflow_command, labels=self._get_labels( kube_executor_config.labels, { 'airflow-worker': worker_uuid, 'dag_id': dag_id, 'task_id': task_id, 'execution_date': execution_date, 'try_number': str(try_number), }), envs=self._get_environment(), secrets=self._get_secrets(), service_account_name=self.kube_config.worker_service_account_name, image_pull_secrets=self.kube_config.image_pull_secrets, init_containers=worker_init_container_spec, volumes=volumes, volume_mounts=volume_mounts, resources=resources, annotations=annotations, node_selectors=(kube_executor_config.node_selectors or self.kube_config.kube_node_selectors), affinity=affinity, tolerations=tolerations, security_context=self._get_security_context(), configmaps=self._get_configmaps(), host_aliases=self._get_host_aliases())
def test_extract_limits_resources(self): # Test when resources is not empty resources = Resources(limit_memory='2Gi', limit_cpu=2) pod = Pod('v3.14', {}, [], resources=resources) self.expected['spec']['containers'][0]['resources'] = { 'limits': { 'memory': '2Gi', 'cpu': 2 } } KubernetesRequestFactory.extract_resources(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def _extract_resources(resources): if isinstance(resources, k8s.V1ResourceRequirements): requests = resources.requests or {} limits = resources.limits or {} return Resources( request_memory=requests.get('memory', None), request_cpu=requests.get('cpu', None), request_ephemeral_storage=requests.get('ephemeral-storage', None), limit_memory=limits.get('memory', None), limit_cpu=limits.get('cpu', None), limit_ephemeral_storage=limits.get('ephemeral-storage', None), limit_gpu=limits.get('nvidia.com/gpu')) elif isinstance(resources, Resources): return resources
def test_extract_all_resources(self): # Test when resources is not empty resources = Resources(request_memory='1Gi', request_cpu=1, limit_memory='2Gi', limit_cpu=2, limit_gpu=3) pod = Pod('v3.14', {}, [], resources=resources) self.expected['spec']['containers'][0]['resources'] = { 'requests': { 'memory': '1Gi', 'cpu': 1 }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 3 } } KubernetesRequestFactory.extract_resources(pod, self.input_req) self.assertEqual(self.input_req, self.expected)
def _set_resources(resources): if not resources: return [] return [Resources(**resources)]
def _set_resources(self, resources): inputResource = Resources() if resources: for item in resources.keys(): setattr(inputResource, item, resources[item]) return inputResource
def test_to_v1_pod(self, mock_uuid): from airflow.contrib.kubernetes.pod import Pod as DeprecatedPod from airflow.kubernetes.volume import Volume from airflow.kubernetes.volume_mount import VolumeMount from airflow.kubernetes.secret import Secret from airflow.kubernetes.pod import Resources import uuid static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48') mock_uuid.return_value = static_uuid pod = DeprecatedPod( image="foo", name="bar", namespace="baz", image_pull_policy="Never", envs={"test_key": "test_value"}, cmds=["airflow"], resources=Resources(request_memory="1G", request_cpu="100Mi", limit_gpu="100G"), init_containers=k8s.V1Container(name="test-container", volume_mounts=k8s.V1VolumeMount( mount_path="/foo/bar", name="init-volume-secret")), volumes=[ Volume(name="foo", configs={}), { "name": "bar", 'secret': { 'secretName': 'volume-secret' } } ], secrets=[ Secret("volume", None, "init-volume-secret"), Secret('env', "AIRFLOW_SECRET", 'secret_name', "airflow_config"), Secret("volume", "/opt/airflow", "volume-secret", "secret-key") ], volume_mounts=[ VolumeMount(name="foo", mount_path="/mnt", sub_path="/", read_only=True) ]) k8s_client = ApiClient() result = pod.to_v1_kubernetes_pod() result = k8s_client.sanitize_for_serialization(result) expected = \ {'metadata': {'labels': {}, 'name': 'bar', 'namespace': 'baz'}, 'spec': {'affinity': {}, 'containers': [{'args': [], 'command': ['airflow'], 'env': [{'name': 'test_key', 'value': 'test_value'}, {'name': 'AIRFLOW_SECRET', 'valueFrom': {'secretKeyRef': {'key': 'airflow_config', 'name': 'secret_name'}}}], 'image': 'foo', 'imagePullPolicy': 'Never', 'name': 'base', 'resources': {'limits': {'nvidia.com/gpu': '100G'}, 'requests': {'cpu': '100Mi', 'memory': '1G'}}, 'volumeMounts': [{'mountPath': '/mnt', 'name': 'foo', 'readOnly': True, 'subPath': '/'}, {'mountPath': '/opt/airflow', 'name': 'secretvol' + str(static_uuid), 'readOnly': True}]}], 'hostNetwork': False, 'initContainers': {'name': 'test-container', 'volumeMounts': {'mountPath': '/foo/bar', 'name': 'init-volume-secret'}}, 'securityContext': {}, 'tolerations': [], 'volumes': [{'name': 'foo'}, {'name': 'bar', 'secret': {'secretName': 'volume-secret'}}, {'name': 'secretvolcf4a56d2-8101-4217-b027-2af6216feb48', 'secret': {'secretName': 'init-volume-secret'}}, {'name': 'secretvol' + str(static_uuid), 'secret': {'secretName': 'volume-secret'}} ]}} self.maxDiff = None self.assertEqual(expected, result)
def setUp(self): 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.resources = Resources('1Gi', 1, '2Gi', 2, 1) self.k8s_client = ApiClient() self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod-0', 'labels': { 'app': 'myapp' }, 'namespace': 'default' }, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'args': [], 'command': ['sh', '-c', 'echo Hello Kubernetes!'], 'imagePullPolicy': 'IfNotPresent', 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }, { 'secretRef': { 'name': 'secret_a' } }], 'resources': { 'requests': { 'memory': '1Gi', 'cpu': 1 }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 1 }, }, 'ports': [{ 'name': 'foo', 'containerPort': 1234 }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol0', 'readOnly': True }] }], 'restartPolicy': 'Never', 'volumes': [{ 'name': 'secretvol0', 'secret': { 'secretName': 'secret_b' } }], 'hostNetwork': False, 'imagePullSecrets': [{ 'name': 'pull_secret_a' }, { 'name': 'pull_secret_b' }], 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } }
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': 'polinux/stress', '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.labels = { 'airflow-worker': 'uuid', 'dag_id': 'dag_id', 'execution_date': 'date', 'task_id': 'task_id', 'try_number': '3', 'airflow_version': mock.ANY, 'kubernetes_executor': 'True' } self.metadata = { 'labels': self.labels, 'name': 'pod_id-' + self.static_uuid.hex, 'namespace': 'namespace' } self.resources = Resources('1Gi', 1, '2Gi', '2Gi', 2, 1, '4Gi') self.k8s_client = ApiClient() self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod-' + self.static_uuid.hex, 'labels': { 'app': 'myapp' }, 'namespace': 'default' }, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'args': [], 'command': ['sh', '-c', 'echo Hello Kubernetes!'], 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }, { 'secretRef': { 'name': 'secret_a' } }], 'resources': { 'requests': { 'memory': '1Gi', 'cpu': 1, 'ephemeral-storage': '2Gi' }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 1, 'ephemeral-storage': '4Gi' }, }, 'ports': [{ 'name': 'foo', 'containerPort': 1234 }], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol' + str(self.static_uuid), 'readOnly': True }] }], 'volumes': [{ 'name': 'secretvol' + str(self.static_uuid), 'secret': { 'secretName': 'secret_b' } }], 'hostNetwork': False, 'imagePullSecrets': [{ 'name': 'pull_secret_a' }, { 'name': 'pull_secret_b' }], 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } }
def _set_resources(resources): return [Resources(**resources) if resources else Resources()]
def test_display_resources(self): resources_string = str(Resources('1Gi', 1)) self.assertEqual( resources_string, "Request: [cpu: 1, memory: 1Gi], Limit: [cpu: None, memory: None, gpu: None]" )
def setUp(self): self.simple_pod_request_factory = SimplePodRequestFactory() self.xcom_pod_request_factory = ExtractXcomPodRequestFactory() self.pod = Pod( image='busybox', envs={ 'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning' }, name='myapp-pod', cmds=['sh', '-c', 'echo Hello Kubernetes!'], labels={'app': 'myapp'}, image_pull_secrets='pull_secret_a,pull_secret_b', configmaps=['configmap_a', 'configmap_b'], ports=[{'name': 'foo', 'containerPort': 1234}], resources=Resources('1Gi', 1, '2Gi', 2, 1), 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'), ], security_context={ 'runAsUser': 1000, 'fsGroup': 2000, } ) self.maxDiff = None self.expected = { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'myapp-pod', 'labels': {'app': 'myapp'}, 'annotations': {}}, 'spec': { 'containers': [{ 'name': 'base', 'image': 'busybox', 'command': [ 'sh', '-c', 'echo Hello Kubernetes!' ], 'imagePullPolicy': 'IfNotPresent', 'args': [], 'env': [{ 'name': 'ENVIRONMENT', 'value': 'prod' }, { 'name': 'LOG_LEVEL', 'value': 'warning' }, { 'name': 'TARGET', 'valueFrom': { 'secretKeyRef': { 'name': 'secret_b', 'key': 'source_b' } } }], 'envFrom': [{ 'secretRef': { 'name': 'secret_a' } }, { 'configMapRef': { 'name': 'configmap_a' } }, { 'configMapRef': { 'name': 'configmap_b' } }], 'resources': { 'requests': { 'memory': '1Gi', 'cpu': 1 }, 'limits': { 'memory': '2Gi', 'cpu': 2, 'nvidia.com/gpu': 1 }, }, 'ports': [{'name': 'foo', 'containerPort': 1234}], 'volumeMounts': [{ 'mountPath': '/etc/foo', 'name': 'secretvol0', 'readOnly': True }] }], 'restartPolicy': 'Never', 'nodeSelector': {}, 'volumes': [{ 'name': 'secretvol0', 'secret': { 'secretName': 'secret_b' } }], 'imagePullSecrets': [ {'name': 'pull_secret_a'}, {'name': 'pull_secret_b'} ], 'affinity': {}, 'securityContext': { 'runAsUser': 1000, 'fsGroup': 2000, }, } }