def execute(self, context):
try:
client = kube_client.get_kube_client(
in_cluster=self.in_cluster,
cluster_context=self.cluster_context,
config_file=self.config_file)
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
resources=self.resources,
pod=self.full_pod_spec,
).gen_pod()
pod = append_to_pod(pod, self.ports)
pod = append_to_pod(pod, self.pod_runtime_info_envs)
pod = append_to_pod(pod, self.volumes)
pod = append_to_pod(pod, self.volume_mounts)
pod = append_to_pod(pod, self.secrets)
self.pod = pod
launcher = pod_launcher.PodLauncher(kube_client=client,
extract_xcom=self.do_xcom_push)
try:
(final_state, result) = launcher.run_pod(
pod,
startup_timeout=self.startup_timeout_seconds,
get_logs=self.get_logs)
finally:
if self.is_delete_operator_pod:
launcher.delete_pod(pod)
if final_state != State.SUCCESS:
raise AirflowException(
'Pod returned a failure: {state}'.format(
state=final_state))
return result
except AirflowException as ex:
raise AirflowException(
'Pod Launching failed: {error}'.format(error=ex))
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),
'airflow_version': airflow_version.replace('+', '-'),
'kubernetes_executor': 'True',
},
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 as_pod(self) -> k8s.V1Pod:
"""Creates POD."""
if self.kube_config.pod_template_file:
return PodGenerator(pod_template_file=self.kube_config.
pod_template_file).gen_pod()
pod = PodGenerator(
image=self.kube_config.kube_image,
image_pull_policy=self.kube_config.kube_image_pull_policy
or 'IfNotPresent',
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(),
labels=self.kube_config.kube_labels,
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
or 'default',
restart_policy='Never',
resources=self._get_resources(),
).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):
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_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 create_pod_request_obj(self) -> k8s.V1Pod:
"""
Creates a V1Pod based on user parameters. Note that a `pod` or `pod_template_file`
will supersede all other values.
"""
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
schedulername=self.schedulername,
init_containers=self.init_containers,
restart_policy='Never',
priority_class_name=self.priority_class_name,
pod_template_file=self.pod_template_file,
pod=self.full_pod_spec,
).gen_pod()
pod = append_to_pod(
pod,
self.pod_runtime_info_envs + self.ports + # type: ignore
self.resources + self.secrets + # type: ignore
self.volumes + # type: ignore
self.volume_mounts # type: ignore
)
return pod
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 create_new_pod_for_operator(
self, labels, launcher) -> Tuple[State, k8s.V1Pod, Optional[str]]:
"""
Creates a new pod and monitors for duration of task
:param labels: labels used to track pod
:param launcher: pod launcher that will manage launching and monitoring pods
:return:
"""
if not (self.full_pod_spec or self.pod_template_file):
# Add Airflow Version to the label
# And a label to identify that pod is launched by KubernetesPodOperator
self.labels.update({
'airflow_version':
airflow_version.replace('+', '-'),
'kubernetes_pod_operator':
'True',
})
self.labels.update(labels)
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
schedulername=self.schedulername,
init_containers=self.init_containers,
restart_policy='Never',
priority_class_name=self.priority_class_name,
pod_template_file=self.pod_template_file,
pod=self.full_pod_spec,
).gen_pod()
# noinspection PyTypeChecker
pod = append_to_pod(
pod,
self.pod_runtime_info_envs + self.ports + # type: ignore
self.resources + self.secrets + # type: ignore
self.volumes + # type: ignore
self.volume_mounts # type: ignore
)
self.pod = pod
self.log.debug("Starting pod:\n%s", yaml.safe_dump(pod.to_dict()))
try:
launcher.start_pod(pod,
startup_timeout=self.startup_timeout_seconds)
final_state, result = launcher.monitor_pod(pod=pod,
get_logs=self.get_logs)
except AirflowException:
if self.log_events_on_failure:
for event in launcher.read_pod_events(pod).items:
self.log.error("Pod Event: %s - %s", event.reason,
event.message)
raise
finally:
if self.is_delete_operator_pod:
launcher.delete_pod(pod)
return final_state, pod, result
def execute(self, context):
try:
if self.in_cluster is not None:
client = kube_client.get_kube_client(
in_cluster=self.in_cluster,
cluster_context=self.cluster_context,
config_file=self.config_file)
else:
client = kube_client.get_kube_client(
cluster_context=self.cluster_context,
config_file=self.config_file)
if not (self.full_pod_spec or self.pod_template_file):
# Add Airflow Version to the label
# And a label to identify that pod is launched by KubernetesPodOperator
self.labels.update({
'airflow_version':
airflow_version.replace('+', '-'),
'kubernetes_pod_operator':
'True',
})
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
schedulername=self.schedulername,
init_containers=self.init_containers,
restart_policy='Never',
priority_class_name=self.priority_class_name,
pod_template_file=self.pod_template_file,
pod=self.full_pod_spec,
).gen_pod()
pod = append_to_pod(
pod, self.pod_runtime_info_envs + self.ports + self.resources +
self.secrets + self.volumes + self.volume_mounts)
self.pod = pod
launcher = pod_launcher.PodLauncher(kube_client=client,
extract_xcom=self.do_xcom_push)
final_state, result = None, None
try:
(final_state, result) = launcher.run_pod(
pod,
startup_timeout=self.startup_timeout_seconds,
get_logs=self.get_logs)
finally:
if final_state != State.SUCCESS:
# Before deleting the pod we get events and status of the pod (events can be fetched
# after pod deletion but not statuses. For consistency both are fetched before deletion.
self._log_pod_events_on_failure(pod, launcher)
self._log_pod_status_on_failure(pod, launcher)
if self.is_delete_operator_pod:
launcher.delete_pod(pod)
if final_state != State.SUCCESS:
if self.retry_only_on_pod_launching_failure:
self.log.info(
'Task failure due to task image, avoiding retries and failing the task.'
)
ti = context.get('task_instance')
ti.error()
raise AirflowException(
'Pod returned a failure: {state}'.format(
state=final_state))
return result
except AirflowException as ex:
raise AirflowException(
'Pod Launching failed: {error}'.format(error=ex))
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 execute(self, context):
try:
if self.in_cluster is not None:
client = kube_client.get_kube_client(
in_cluster=self.in_cluster,
cluster_context=self.cluster_context,
config_file=self.config_file)
else:
client = kube_client.get_kube_client(
cluster_context=self.cluster_context,
config_file=self.config_file)
if not (self.full_pod_spec or self.pod_template_file):
# Add Airflow Version to the label
# And a label to identify that pod is launched by KubernetesPodOperator
self.labels.update({
'airflow_version':
airflow_version.replace('+', '-'),
'kubernetes_pod_operator':
'True',
})
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
schedulername=self.schedulername,
init_containers=self.init_containers,
restart_policy='Never',
priority_class_name=self.priority_class_name,
pod_template_file=self.pod_template_file,
pod=self.full_pod_spec,
).gen_pod()
pod = append_to_pod(
pod, self.pod_runtime_info_envs + self.ports + self.resources +
self.secrets + self.volumes + self.volume_mounts)
self.pod = pod
launcher = pod_launcher.PodLauncher(kube_client=client,
extract_xcom=self.do_xcom_push)
try:
(final_state, result) = launcher.run_pod(
pod,
startup_timeout=self.startup_timeout_seconds,
get_logs=self.get_logs)
except AirflowException:
if self.log_events_on_failure:
for event in launcher.read_pod_events(pod).items:
self.log.error("Pod Event: %s - %s", event.reason,
event.message)
raise
finally:
if self.is_delete_operator_pod:
launcher.delete_pod(pod)
if final_state != State.SUCCESS:
if self.log_events_on_failure:
for event in launcher.read_pod_events(pod).items:
self.log.error("Pod Event: %s - %s", event.reason,
event.message)
raise AirflowException(
'Pod returned a failure: {state}'.format(
state=final_state))
return result
except AirflowException as ex:
raise AirflowException(
'Pod Launching failed: {error}'.format(error=ex))
def create_pod_request_obj(self):
"""
Creates a V1Pod based on user parameters. Note that a `pod` or `pod_template_file`
will supersede all other values.
"""
if self.pod_template_file:
pod_template = pod_generator.PodGenerator.deserialize_model_file(
self.pod_template_file)
else:
pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name"))
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
cmds=self.cmds,
args=self.arguments,
labels=self.labels,
name=self.name,
envs=self.env_vars,
extract_xcom=self.do_xcom_push,
image_pull_policy=self.image_pull_policy,
node_selectors=self.node_selectors,
annotations=self.annotations,
affinity=self.affinity,
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
hostnetwork=self.hostnetwork,
tolerations=self.tolerations,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
init_containers=self.init_containers,
restart_policy='Never',
schedulername=self.schedulername,
priority_class_name=self.priority_class_name,
).gen_pod()
# noinspection PyTypeChecker
pod = append_to_pod(
pod,
self.pod_runtime_info_envs + # type: ignore
self.ports + # type: ignore
self.resources + # type: ignore
self.secrets + # type: ignore
self.volumes + # type: ignore
self.volume_mounts # type: ignore
)
env_from = pod.spec.containers[0].env_from or []
for configmap in self.configmaps:
env_from.append(
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(
name=configmap)))
pod.spec.containers[0].env_from = env_from
if self.full_pod_spec:
pod_template = PodGenerator.reconcile_pods(pod_template,
self.full_pod_spec)
pod = PodGenerator.reconcile_pods(pod_template, pod)
# if self.do_xcom_push:
# pod = PodGenerator.add_sidecar(pod)
return pod