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 run_next(self, next_job):
"""
The run_next command will check the task_queue for any un-run jobs.
It will then create a unique job-id, launch that job in the cluster,
and store relevant info in the current_jobs map so we can track the job's
status
"""
self.log.info('Kubernetes job is %s', str(next_job))
key, command, kube_executor_config = next_job
dag_id, task_id, execution_date, try_number = key
config_pod = self.worker_configuration.make_pod(
namespace=self.namespace,
worker_uuid=self.worker_uuid,
pod_id=self._create_pod_id(dag_id, task_id),
dag_id=self._make_safe_label_value(dag_id),
task_id=self._make_safe_label_value(task_id),
try_number=try_number,
execution_date=self._datetime_to_label_safe_datestring(
execution_date),
airflow_command=command)
# Reconcile the pod generated by the Operator and the Pod
# generated by the .cfg file
pod = PodGenerator.reconcile_pods(config_pod, kube_executor_config)
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s",
pod.spec.containers[0].image)
# the watcher will monitor pods, so we do not block.
self.launcher.run_pod_async(
pod, **self.kube_config.kube_client_request_args)
self.log.debug("Kubernetes Job created!")
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.
"""
self.log.debug("Creating pod for K8sPodOperator task %s", self.task_id)
if self.pod_template_file:
self.log.debug("Pod template file found, will parse for base pod")
pod_template = pod_generator.PodGenerator.deserialize_model_file(
self.pod_template_file)
else:
pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name"))
pod = k8s.V1Pod(api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
namespace=self.namespace,
labels=self.labels,
name=self.name,
annotations=self.annotations,
),
spec=k8s.V1PodSpec(
node_selector=self.node_selectors,
affinity=self.affinity,
tolerations=self.tolerations,
init_containers=self.init_containers,
containers=[
k8s.V1Container(
image=self.image,
name="base",
command=self.cmds,
ports=self.ports,
resources=self.k8s_resources,
volume_mounts=self.volume_mounts,
args=self.arguments,
env=self.env_vars,
env_from=self.env_from,
)
],
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
host_network=self.hostnetwork,
security_context=self.security_context,
dns_policy=self.dnspolicy,
scheduler_name=self.schedulername,
restart_policy='Never',
priority_class_name=self.priority_class_name,
volumes=self.volumes,
))
pod = PodGenerator.reconcile_pods(pod_template, pod)
for secret in self.secrets:
self.log.debug("Adding secret to task %s", self.task_id)
pod = secret.attach_to_pod(pod)
if self.do_xcom_push:
self.log.debug("Adding xcom sidecar to task %s", self.task_id)
pod = PodGenerator.add_xcom_sidecar(pod)
return pod
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_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 _mutate_pod_backcompat(pod):
"""Backwards compatible Pod Mutation Hook"""
try:
dummy_pod = _convert_to_airflow_pod(pod)
settings.pod_mutation_hook(dummy_pod)
warnings.warn(
"Using `airflow.contrib.kubernetes.pod.Pod` is deprecated. "
"Please use `k8s.V1Pod` instead.", DeprecationWarning, stacklevel=2
)
dummy_pod = dummy_pod.to_v1_kubernetes_pod()
new_pod = PodGenerator.reconcile_pods(pod, dummy_pod)
except AttributeError as e:
try:
settings.pod_mutation_hook(pod)
return pod
except AttributeError as e2:
raise Exception([e, e2])
return new_pod
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 build_pod_request_obj(self, context=None):
"""
Returns V1Pod object based on pod template file, full pod spec, and other operator parameters.
The V1Pod attributes are derived (in order of precedence) from operator params, full pod spec, pod
template file.
"""
self.log.debug("Creating pod for KubernetesPodOperator task %s",
self.task_id)
if self.pod_template_file:
self.log.debug("Pod template file found, will parse for base pod")
pod_template = pod_generator.PodGenerator.deserialize_model_file(
self.pod_template_file)
if self.full_pod_spec:
pod_template = PodGenerator.reconcile_pods(
pod_template, self.full_pod_spec)
elif self.full_pod_spec:
pod_template = self.full_pod_spec
else:
pod_template = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="name"))
pod = k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
namespace=self.namespace,
labels=self.labels,
name=self.name,
annotations=self.annotations,
),
spec=k8s.V1PodSpec(
node_selector=self.node_selector,
affinity=self.affinity,
tolerations=self.tolerations,
init_containers=self.init_containers,
containers=[
k8s.V1Container(
image=self.image,
name=self.BASE_CONTAINER_NAME,
command=self.cmds,
ports=self.ports,
image_pull_policy=self.image_pull_policy,
resources=self.k8s_resources,
volume_mounts=self.volume_mounts,
args=self.arguments,
env=self.env_vars,
env_from=self.env_from,
)
],
image_pull_secrets=self.image_pull_secrets,
service_account_name=self.service_account_name,
host_network=self.hostnetwork,
security_context=self.security_context,
dns_policy=self.dnspolicy,
scheduler_name=self.schedulername,
restart_policy='Never',
priority_class_name=self.priority_class_name,
volumes=self.volumes,
),
)
pod = PodGenerator.reconcile_pods(pod_template, pod)
if self.random_name_suffix:
pod.metadata.name = PodGenerator.make_unique_pod_id(
pod.metadata.name)
for secret in self.secrets:
self.log.debug("Adding secret to task %s", self.task_id)
pod = secret.attach_to_pod(pod)
if self.do_xcom_push:
self.log.debug("Adding xcom sidecar to task %s", self.task_id)
pod = xcom_sidecar.add_xcom_sidecar(pod)
labels = self._get_ti_pod_labels(context)
self.log.info("Creating pod %s with labels: %s", pod.metadata.name,
labels)
# Merge Pod Identifying labels with labels passed to operator
pod.metadata.labels.update(labels)
# Add Airflow Version to the label
# And a label to identify that pod is launched by KubernetesPodOperator
pod.metadata.labels.update({
'airflow_version':
airflow_version.replace('+', '-'),
'kubernetes_pod_operator':
'True',
})
pod_mutation_hook(pod)
return pod
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)
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
