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,
},
}
}