def make_pod(self,
namespace,
image,
pod_id,
cmds,
arguments,
labels,
kube_executor_config=None):
volumes, volume_mounts = 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
# )
return Pod(
namespace=namespace,
name=pod_id + "-" + str(uuid.uuid1())[:8],
image=image,
cmds=cmds,
args=arguments,
labels=labels,
envs=self.env_vars,
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=None)
def test_extract_service_account_name(self):
service_account_name = 'service_account_name'
pod = Pod('v3.14', {}, [], service_account_name=service_account_name)
self.expected['spec']['serviceAccountName'] = service_account_name
KubernetesRequestFactory.extract_service_account_name(
pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_volume_secrets(self):
# Test when secrets is not empty
secrets = [
Secret('volume', 'KEY1', 's1', 'key-1'),
Secret('env', 'KEY2', 's2'),
Secret('volume', 'KEY3', 's3', 'key-2')
]
pod = Pod('v3.14', {}, [], secrets=secrets)
self.expected['spec']['containers'][0]['volumeMounts'] = [{
'mountPath': 'KEY1',
'name': 'secretvol0',
'readOnly': True
}, {
'mountPath': 'KEY3',
'name': 'secretvol1',
'readOnly': True
}]
self.expected['spec']['volumes'] = [{
'name': 'secretvol0',
'secret': {
'secretName': 's1'
}
}, {
'name': 'secretvol1',
'secret': {
'secretName': 's3'
}
}]
KubernetesRequestFactory.extract_volume_secrets(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_affinity(self):
# Test when affinity is not empty
affinity = {'podAffinity': 'requiredDuringSchedulingIgnoredDuringExecution'}
pod = Pod('v3.14', {}, [], affinity=affinity)
self.expected['spec']['affinity'] = affinity
KubernetesRequestFactory.extract_affinity(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_node_selector(self):
# Test when affinity is not empty
node_selectors = {'disktype': 'ssd', 'accelerator': 'nvidia-tesla-p100'}
pod = Pod('v3.14', {}, [], node_selectors=node_selectors)
self.expected['spec']['nodeSelector'] = node_selectors
KubernetesRequestFactory.extract_node_selector(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_labels(self):
# Test when labels are not empty
labels = {'label_a': 'val_a', 'label_b': 'val_b'}
pod = Pod('v3.14', {}, [], labels=labels)
self.expected['metadata']['labels'] = labels
KubernetesRequestFactory.extract_labels(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_attach_volume_mounts(self):
# Test when volumes is not empty list
volume_mounts = ['vol_a', 'vol_b']
pod = Pod('v3.14', {}, [], volume_mounts=volume_mounts)
self.expected['spec']['containers'][0]['volumeMounts'] = volume_mounts
KubernetesRequestFactory.attach_volume_mounts(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_annotations(self):
# Test when annotations are not empty
annotations = {'annot_a': 'val_a', 'annot_b': 'val_b'}
pod = Pod('v3.14', {}, [], annotations=annotations)
self.expected['metadata']['annotations'] = annotations
KubernetesRequestFactory.extract_annotations(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_image_pull_policy(self):
# Test when pull policy is not none
pull_policy = 'IfNotPresent'
pod = Pod('v3.14', {}, [], image_pull_policy=pull_policy)
KubernetesRequestFactory.extract_image_pull_policy(pod, self.input_req)
self.expected['spec']['containers'][0]['imagePullPolicy'] = pull_policy
self.assertEqual(self.input_req, self.expected)
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),
'airflow_version': airflow_version.replace('+', '-'),
'kubernetes_executor': 'True',
}),
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())
def test_security_context(self):
security_context = {
'runAsUser': 1000,
'fsGroup': 2000
}
pod = Pod('v3.14', {}, [], security_context=security_context)
self.expected['spec']['securityContext'] = security_context
KubernetesRequestFactory.extract_security_context(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_image_pull_secrets(self):
image_pull_secrets = 'secret_a,secret_b,secret_c'
pod = Pod('v3.14', {}, [], image_pull_secrets=image_pull_secrets)
self.expected['spec']['imagePullSecrets'] = [
{'name': 'secret_a'},
{'name': 'secret_b'},
{'name': 'secret_c'},
]
KubernetesRequestFactory.extract_image_pull_secrets(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_tolerations(self):
tolerations = [{
'key': 'key',
'operator': 'Equal',
'value': 'value',
'effect': 'NoSchedule'
}]
pod = Pod('v3.14', {}, [], tolerations=tolerations)
self.expected['spec']['tolerations'] = tolerations
KubernetesRequestFactory.extract_tolerations(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_env_and_secrets(self):
# Test when secrets and envs are not empty
secrets = [
Secret('env', None, 's1'),
Secret('volume', 'KEY2', 's2', 'key-2'),
Secret('env', None, 's3')
]
envs = {
'ENV1': 'val1',
'ENV2': 'val2'
}
configmaps = ['configmap_a', 'configmap_b']
pod_runtime_envs = [PodRuntimeInfoEnv("ENV3", "status.podIP")]
pod = Pod(
image='v3.14',
envs=envs,
cmds=[],
secrets=secrets,
configmaps=configmaps,
pod_runtime_info_envs=pod_runtime_envs)
self.expected['spec']['containers'][0]['env'] = [
{'name': 'ENV1', 'value': 'val1'},
{'name': 'ENV2', 'value': 'val2'},
{
'name': 'ENV3',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
}
]
self.expected['spec']['containers'][0]['envFrom'] = [{
'secretRef': {
'name': 's1'
}
}, {
'secretRef': {
'name': 's3'
}
}, {
'configMapRef': {
'name': 'configmap_a'
}
}, {
'configMapRef': {
'name': 'configmap_b'
}
}]
KubernetesRequestFactory.extract_env_and_secrets(pod, self.input_req)
self.input_req['spec']['containers'][0]['env'].sort(key=lambda x: x['name'])
self.assertEqual(self.input_req, self.expected)
def test_extract_limits(self):
# Test when resources is not empty
resources = Resources(limit_memory='1Gi', limit_cpu=1)
pod = Pod('v3.14', {}, [], resources=resources)
self.expected['spec']['containers'][0]['resources'] = {
'limits': {
'memory': '1Gi',
'cpu': 1
}
}
KubernetesRequestFactory.extract_resources(pod, self.input_req)
self.assertEqual(self.expected, self.input_req)
def test_pod_mutation_hook(self):
"""
Tests that pods are mutated by the pod_mutation_hook
function in airflow_local_settings.
"""
with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"):
from airflow import settings
settings.import_local_settings() # pylint: ignore
pod = Pod(image="ubuntu", envs={}, cmds=['echo "1"'])
settings.pod_mutation_hook(pod)
assert pod.namespace == 'airflow-tests'
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes_dict, volume_mounts_dict = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts_dict))
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)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = dict(kube_executor_config.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={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
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)
def create_sync_pod(self):
return Pod(
namespace=configuration.conf.get("kubernetes", "namespace"),
name=AirflowKubernetesScheduler._create_pod_id("airflow-sync", "sync-worker"),
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy,
cmds=[
"sh",
"-c",
"tar -cz --exclude=__pycache__ -C $AIRFLOW_HOME/dags . | base64"
],
labels={
'cluster': "airflow",
},
service_account_name=self.kube_config.worker_service_account_name,
image_pull_secrets=self.kube_config.image_pull_secrets,
envs={}
)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id, execution_date,
airflow_command, kube_executor_config):
volumes, volume_mounts = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
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
)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
'iam.cloud.google.com/service-account': gcp_sa_key
} if gcp_sa_key else {}
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=['bash', '-cx', '--'],
args=[airflow_command],
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
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=kube_executor_config.affinity
)
def make_pod(self, namespace, image, pod_id, cmds, arguments, labels):
volumes, volume_mounts = self._get_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers()
return Pod(
namespace=namespace,
name=pod_id + "-" + str(uuid.uuid1())[:8],
image=image,
cmds=cmds,
args=arguments,
labels=labels,
envs={},
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=None)
def test_extract_env_and_secrets(self):
# Test when secrets and envs are not empty
secrets = [
Secret('env', None, 's1'),
Secret('volume', 'KEY2', 's2', 'key-2'),
Secret('env', None, 's3')
]
envs = {'ENV1': 'val1', 'ENV2': 'val2'}
configmaps = ['configmap_a', 'configmap_b']
pod = Pod('v3.14', envs, [], secrets=secrets, configmaps=configmaps)
self.expected['spec']['containers'][0]['env'] = [
{
'name': 'ENV1',
'value': 'val1'
},
{
'name': 'ENV2',
'value': 'val2'
},
]
self.expected['spec']['containers'][0]['envFrom'] = [{
'secretRef': {
'name': 's1'
}
}, {
'secretRef': {
'name': 's3'
}
}, {
'configMapRef': {
'name': 'configmap_a'
}
}, {
'configMapRef': {
'name': 'configmap_b'
}
}]
KubernetesRequestFactory.extract_env_and_secrets(pod, self.input_req)
self.input_req['spec']['containers'][0]['env'].sort(
key=lambda x: x['name'])
self.assertEqual(self.input_req, self.expected)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes, volume_mounts = self.init_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
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)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
'iam.cloud.google.com/service-account': gcp_sa_key
} if gcp_sa_key else {}
airflow_command = airflow_command.replace("-sd", "-i -sd")
airflow_path = airflow_command.split('-sd')[-1]
airflow_path = self.worker_airflow_home + airflow_path.split('/')[-1]
airflow_command = airflow_command.split(
'-sd')[0] + '-sd ' + airflow_path
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
cmds=['bash', '-cx', '--'],
args=[airflow_command],
labels={
'airflow-worker': worker_uuid,
'dag_id': dag_id,
'task_id': task_id,
'execution_date': execution_date
},
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)
def _convert_to_airflow_pod(pod):
"""
Converts a k8s V1Pod object into an `airflow.kubernetes.pod.Pod` object.
This function is purely for backwards compatibility
"""
base_container = pod.spec.containers[0] # type: k8s.V1Container
env_vars, secrets = _extract_env_vars_and_secrets(base_container.env)
volumes = _extract_volumes(pod.spec.volumes)
api_client = ApiClient()
init_containers = pod.spec.init_containers
image_pull_secrets = pod.spec.image_pull_secrets or []
if pod.spec.init_containers is not None:
init_containers = [
api_client.sanitize_for_serialization(i)
for i in pod.spec.init_containers
]
dummy_pod = Pod(
image=base_container.image,
envs=env_vars,
cmds=base_container.command,
args=base_container.args,
labels=pod.metadata.labels,
annotations=pod.metadata.annotations,
node_selectors=pod.spec.node_selector,
name=pod.metadata.name,
ports=_extract_ports(base_container.ports),
volumes=volumes,
volume_mounts=_extract_volume_mounts(base_container.volume_mounts),
namespace=pod.metadata.namespace,
image_pull_policy=base_container.image_pull_policy or 'IfNotPresent',
tolerations=pod.spec.tolerations,
init_containers=init_containers,
image_pull_secrets=",".join([i.name for i in image_pull_secrets]),
resources=base_container.resources,
service_account_name=pod.spec.service_account_name,
secrets=secrets,
affinity=api_client.sanitize_for_serialization(pod.spec.affinity),
hostnetwork=pod.spec.host_network,
security_context=_extract_security_context(pod.spec.security_context))
return dummy_pod
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 test_extract_env_and_secret_order(self):
envs = {
'ENV': 'val1',
}
pod_runtime_envs = [PodRuntimeInfoEnv('RUNTIME_ENV', 'status.podIP')]
pod = Pod(image='v3.14',
envs=envs,
cmds=[],
pod_runtime_info_envs=pod_runtime_envs)
self.expected['spec']['containers'][0]['env'] = [{
'name': 'RUNTIME_ENV',
'valueFrom': {
'fieldRef': {
'fieldPath': 'status.podIP'
}
}
}, {
'name': 'ENV',
'value': 'val1'
}]
KubernetesRequestFactory.extract_env_and_secrets(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id,
execution_date, airflow_command, kube_executor_config):
volumes, volume_mounts = self._get_volumes_and_mounts()
worker_init_container_spec = self._get_init_containers(
copy.deepcopy(volume_mounts))
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)
gcp_sa_key = kube_executor_config.gcp_service_account_key
annotations = {
"iam.cloud.google.com/service-account": gcp_sa_key
} if gcp_sa_key else {}
return Pod(
namespace=namespace,
name=pod_id,
image=kube_executor_config.image or self.kube_config.kube_image,
cmds=["bash", "-cx", "--"],
args=[airflow_command],
labels={
"airflow-worker": worker_uuid,
"dag_id": dag_id,
"task_id": task_id,
"execution_date": execution_date
},
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)
def test_extract_image(self):
image = 'v3.14'
pod = Pod(image, {}, [])
KubernetesRequestFactory.extract_image(pod, self.input_req)
self.expected['spec']['containers'][0]['image'] = image
self.assertEqual(self.input_req, self.expected)
def test_identity(self, name):
kube_request_factory_func = getattr(KubernetesRequestFactory, name)
pod = Pod('v3.14', {}, [])
kube_request_factory_func(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_hostnetwork(self):
hostnetwork = True
pod = Pod('v3.14', {}, [], hostnetwork=hostnetwork)
self.expected['spec']['hostNetwork'] = hostnetwork
KubernetesRequestFactory.extract_hostnetwork(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)
def test_extract_init_containers(self):
init_container = 'init_container'
pod = Pod('v3.14', {}, [], init_containers=init_container)
self.expected['spec']['initContainers'] = init_container
KubernetesRequestFactory.extract_init_containers(pod, self.input_req)
self.assertEqual(self.input_req, self.expected)