def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, try_number, airflow_command) -> k8s.V1Pod:
"""Creates POD."""
pod_generator = PodGenerator(
namespace=namespace,
name=pod_id,
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy,
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date,
'try_number': str(try_number),
},
cmds=airflow_command,
volumes=self._get_volumes(),
volume_mounts=self._get_volume_mounts(),
init_containers=self._get_init_containers(),
annotations=self.kube_config.kube_annotations,
affinity=self.kube_config.kube_affinity,
tolerations=self.kube_config.kube_tolerations,
envs=self._get_environment(),
node_selectors=self.kube_config.kube_node_selectors,
service_account_name=self.kube_config.worker_service_account_name,
)
pod = pod_generator.gen_pod()
pod.spec.containers[0].env_from = pod.spec.containers[0].env_from or []
pod.spec.containers[0].env_from.extend(self._get_env_from())
pod.spec.security_context = self._get_security_context()
return append_to_pod(pod, self._get_secrets())
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_reconcile_pods_empty_mutator_pod(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
base_pod = PodGenerator(
image='image1',
name='name1',
envs={
'key1': 'val1'
},
cmds=['/bin/command1.sh', 'arg1'],
ports=[k8s.V1ContainerPort(name='port', container_port=2118)],
volumes=[{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume1'
}],
volume_mounts=[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}],
).gen_pod()
mutator_pod = None
name = 'name1-' + self.static_uuid.hex
base_pod.metadata.name = name
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
self.assertEqual(base_pod, result)
mutator_pod = k8s.V1Pod()
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
self.assertEqual(base_pod, result)
def test_reconcile_pods(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
path = sys.path[
0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml'
base_pod = PodGenerator(pod_template_file=path,
extract_xcom=False).gen_pod()
mutator_pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
name="name2",
labels={"bar": "baz"},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
image='',
name='name',
command=['/bin/command2.sh', 'arg2'],
volume_mounts=[
k8s.V1VolumeMount(
mount_path="/foo/",
name="example-kubernetes-test-volume2")
],
)
],
volumes=[
k8s.V1Volume(
host_path=k8s.V1HostPathVolumeSource(path="/tmp/"),
name="example-kubernetes-test-volume2",
)
],
),
)
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
expected: k8s.V1Pod = self.expected
expected.metadata.name = "name2"
expected.metadata.labels['bar'] = 'baz'
expected.spec.volumes = expected.spec.volumes or []
expected.spec.volumes.append(
k8s.V1Volume(host_path=k8s.V1HostPathVolumeSource(path="/tmp/"),
name="example-kubernetes-test-volume2"))
base_container: k8s.V1Container = expected.spec.containers[0]
base_container.command = ['/bin/command2.sh', 'arg2']
base_container.volume_mounts = [
k8s.V1VolumeMount(mount_path="/foo/",
name="example-kubernetes-test-volume2")
]
base_container.name = "name"
expected.spec.containers[0] = base_container
result_dict = self.k8s_client.sanitize_for_serialization(result)
expected_dict = self.k8s_client.sanitize_for_serialization(expected)
assert result_dict == expected_dict
def test_attach_to_pod(self, mock_uuid):
mock_uuid.return_value = '0'
pod = PodGenerator(image='airflow-worker:latest',
name='base').gen_pod()
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'),
]
k8s_client = ApiClient()
result = append_to_pod(pod, secrets)
result = k8s_client.sanitize_for_serialization(result)
self.assertEqual(result, {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {'name': 'base-0'},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [{
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'key': 'source_b',
'name': 'secret_b'
}
}
}],
'envFrom': [{'secretRef': {'name': 'secret_a'}}],
'image': 'airflow-worker:latest',
'imagePullPolicy': 'IfNotPresent',
'name': 'base',
'ports': [],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'secretvol0',
'readOnly': True}]
}],
'hostNetwork': False,
'imagePullSecrets': [],
'restartPolicy': 'Never',
'volumes': [{
'name': 'secretvol0',
'secret': {'secretName': 'secret_b'}
}]
}
})
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 test_validate_pod_generator(self):
with self.assertRaises(AirflowConfigException):
PodGenerator(image='k', pod=k8s.V1Pod())
with self.assertRaises(AirflowConfigException):
PodGenerator(pod=k8s.V1Pod(), pod_template_file='k')
with self.assertRaises(AirflowConfigException):
PodGenerator(image='k', pod_template_file='k')
PodGenerator(image='k')
PodGenerator(pod_template_file='tests/kubernetes/pod.yaml')
PodGenerator(pod=k8s.V1Pod())
def test_reconcile_pods_empty_mutator_pod(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
path = sys.path[0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml'
pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True)
base_pod = pod_generator.gen_pod()
mutator_pod = None
name = 'name1-' + self.static_uuid.hex
base_pod.metadata.name = name
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
assert base_pod == result
mutator_pod = k8s.V1Pod()
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
assert base_pod == result
def test_port_attach_to_pod(self, mock_uuid):
mock_uuid.return_value = '0'
pod = PodGenerator(image='airflow-worker:latest',
name='base').gen_pod()
ports = [Port('https', 443), Port('http', 80)]
k8s_client = ApiClient()
result = append_to_pod(pod, ports)
result = k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'base-0'
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'image':
'airflow-worker:latest',
'imagePullPolicy':
'IfNotPresent',
'name':
'base',
'ports': [{
'name': 'https',
'containerPort': 443
}, {
'name': 'http',
'containerPort': 80
}],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'restartPolicy':
'Never',
'volumes': []
}
}, result)
def test_port_attach_to_pod(self, mock_uuid):
import uuid
static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48')
mock_uuid.return_value = static_uuid
pod = PodGenerator(image='airflow-worker:latest',
name='base').gen_pod()
ports = [Port('https', 443), Port('http', 80)]
k8s_client = ApiClient()
result = append_to_pod(pod, ports)
result = k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'base-' + static_uuid.hex
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'image':
'airflow-worker:latest',
'name':
'base',
'ports': [{
'name': 'https',
'containerPort': 443
}, {
'name': 'http',
'containerPort': 80
}],
'volumeMounts': [],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': []
}
}, result)
def test_make_pod_with_executor_config(self):
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
config_pod = PodGenerator(
image='',
affinity=self.affinity_config,
tolerations=self.tolerations_config,
).gen_pod()
pod = worker_config.as_pod()
result = PodGenerator.reconcile_pods(pod, config_pod)
self.assertTrue(result.spec.affinity['podAntiAffinity'] is not None)
self.assertEqual(
'app', result.spec.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']['matchExpressions'][0]['key'])
self.assertEqual(2, len(result.spec.tolerations))
self.assertEqual('prod', result.spec.tolerations[1]['key'])
def test_gen_pod_extract_xcom(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
path = sys.path[
0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml'
pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True)
result = pod_generator.gen_pod()
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)
base_container: k8s.V1Container = self.expected.spec.containers[0]
base_container.volume_mounts = base_container.volume_mounts or []
base_container.volume_mounts.append(
k8s.V1VolumeMount(name="xcom", mount_path="/airflow/xcom"))
self.expected.spec.containers[0] = base_container
self.expected.spec.volumes = self.expected.spec.volumes or []
self.expected.spec.volumes.append(
k8s.V1Volume(
name='xcom',
empty_dir={},
))
result_dict = self.k8s_client.sanitize_for_serialization(result)
expected_dict = self.k8s_client.sanitize_for_serialization(
self.expected)
assert result_dict == expected_dict
def as_pod(self) -> k8s.V1Pod:
"""Creates POD."""
pod_generator = PodGenerator(
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy,
image_pull_secrets=self.kube_config.image_pull_secrets,
volumes=self._get_volumes(),
volume_mounts=self._get_volume_mounts(),
init_containers=self._get_init_containers(),
annotations=self.kube_config.kube_annotations,
affinity=self.kube_config.kube_affinity,
tolerations=self.kube_config.kube_tolerations,
envs=self._get_environment(),
node_selectors=self.kube_config.kube_node_selectors,
service_account_name=self.kube_config.worker_service_account_name,
)
pod = pod_generator.gen_pod()
pod.spec.containers[0].env_from = pod.spec.containers[0].env_from or []
pod.spec.containers[0].env_from.extend(self._get_env_from())
pod.spec.security_context = self._get_security_context()
return append_to_pod(pod, self._get_secrets())
def test_attach_to_pod(self, mock_uuid):
static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48')
mock_uuid.return_value = static_uuid
path = sys.path[0] + '/tests/kubernetes/pod_generator_base.yaml'
pod = PodGenerator(pod_template_file=path).gen_pod()
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'),
]
k8s_client = ApiClient()
pod = append_to_pod(pod, secrets)
result = k8s_client.sanitize_for_serialization(pod)
self.assertEqual(
result, {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'labels': {
'app': 'myapp'
},
'name': 'myapp-pod-cf4a56d281014217b0272af6216feb48',
'namespace': 'default'
},
'spec': {
'containers': [{
'command': ['sh', '-c', 'echo Hello Kubernetes!'],
'env': [{
'name': 'ENVIRONMENT',
'value': 'prod'
}, {
'name': 'LOG_LEVEL',
'value': 'warning'
}, {
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'key': 'source_b',
'name': 'secret_b'
}
}
}],
'envFrom': [{
'configMapRef': {
'name': 'configmap_a'
}
}, {
'secretRef': {
'name': 'secret_a'
}
}],
'image':
'busybox',
'name':
'base',
'ports': [{
'containerPort': 1234,
'name': 'foo'
}],
'resources': {
'limits': {
'memory': '200Mi'
},
'requests': {
'memory': '100Mi'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}, {
'mountPath': '/etc/foo',
'name': 'secretvol' + str(static_uuid),
'readOnly': True
}]
}, {
'command': [
'sh', '-c',
'trap "exit 0" INT; while true; do sleep '
'30; done;'
],
'image':
'alpine',
'name':
'airflow-xcom-sidecar',
'resources': {
'requests': {
'cpu': '1m'
}
},
'volumeMounts': [{
'mountPath': '/airflow/xcom',
'name': 'xcom'
}]
}],
'hostNetwork':
True,
'imagePullSecrets': [{
'name': 'pull_secret_a'
}, {
'name': 'pull_secret_b'
}],
'securityContext': {
'fsGroup': 2000,
'runAsUser': 1000
},
'volumes': [{
'emptyDir': {},
'name': 'xcom'
}, {
'name': 'secretvol' + str(static_uuid),
'secret': {
'secretName': 'secret_b'
}
}]
}
})
def test_reconcile_pods(self):
with mock.patch('uuid.uuid4') as mock_uuid:
mock_uuid.return_value = '0'
base_pod = PodGenerator(
image='image1',
name='name1',
envs={
'key1': 'val1'
},
cmds=['/bin/command1.sh', 'arg1'],
ports=k8s.V1ContainerPort(name='port', container_port=2118),
volumes=[{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume1'
}],
volume_mounts=[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}],
).gen_pod()
mutator_pod = PodGenerator(envs={
'key2': 'val2'
},
image='',
name='name2',
cmds=['/bin/command2.sh', 'arg2'],
volumes=[{
'hostPath': {
'path': '/tmp/'
},
'name':
'example-kubernetes-test-volume2'
}],
volume_mounts=[{
'mountPath':
'/foo/',
'name':
'example-kubernetes-test-volume2'
}]).gen_pod()
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
result = self.k8s_client.sanitize_for_serialization(result)
self.assertEqual(
result, {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'name2-0'
},
'spec': {
'containers': [{
'args': [],
'command': ['/bin/command1.sh', 'arg1'],
'env': [{
'name': 'key1',
'value': 'val1'
}, {
'name': 'key2',
'value': 'val2'
}],
'envFrom': [],
'image':
'image1',
'imagePullPolicy':
'IfNotPresent',
'name':
'base',
'ports': {
'containerPort': 2118,
'name': 'port',
},
'volumeMounts': [
{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}, {
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume2'
}
]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'restartPolicy':
'Never',
'volumes': [{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume1'
}, {
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume2'
}]
}
})
def test_reconcile_pods(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
base_pod = PodGenerator(
image='image1',
name='name1',
envs={
'key1': 'val1'
},
cmds=['/bin/command1.sh', 'arg1'],
ports=[k8s.V1ContainerPort(name='port', container_port=2118)],
volumes=[{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume1'
}],
labels={
"foo": "bar"
},
volume_mounts=[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}],
).gen_pod()
mutator_pod = PodGenerator(envs={
'key2': 'val2'
},
image='',
name='name2',
labels={
"bar": "baz"
},
cmds=['/bin/command2.sh', 'arg2'],
volumes=[{
'hostPath': {
'path': '/tmp/'
},
'name':
'example-kubernetes-test-volume2'
}],
volume_mounts=[{
'mountPath':
'/foo/',
'name':
'example-kubernetes-test-volume2'
}]).gen_pod()
result = PodGenerator.reconcile_pods(base_pod, mutator_pod)
result = self.k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'name2-' + self.static_uuid.hex,
'labels': {
'foo': 'bar',
"bar": "baz"
}
},
'spec': {
'containers': [{
'args': [],
'command': ['/bin/command2.sh', 'arg2'],
'env': [{
'name': 'key1',
'value': 'val1'
}, {
'name': 'key2',
'value': 'val2'
}],
'envFrom': [],
'image':
'image1',
'name':
'base',
'ports': [{
'containerPort': 2118,
'name': 'port',
}],
'volumeMounts': [
{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume1'
}, {
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume2'
}
]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume1'
}, {
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume2'
}]
}
}, result)