def _deserialize(cls, encoded_var: Any) -> Any:
"""Helper function of depth first search for deserialization."""
# JSON primitives (except for dict) are not encoded.
if cls._is_primitive(encoded_var):
return encoded_var
elif isinstance(encoded_var, list):
return [cls._deserialize(v) for v in encoded_var]
if not isinstance(encoded_var, dict):
raise ValueError(
f"The encoded_var should be dict and is {type(encoded_var)}")
var = encoded_var[Encoding.VAR]
type_ = encoded_var[Encoding.TYPE]
if type_ == DAT.DICT:
return {k: cls._deserialize(v) for k, v in var.items()}
elif type_ == DAT.DAG:
return SerializedDAG.deserialize_dag(var)
elif type_ == DAT.OP:
return SerializedBaseOperator.deserialize_operator(var)
elif type_ == DAT.DATETIME:
return pendulum.from_timestamp(var)
elif type_ == DAT.POD:
if not HAS_KUBERNETES:
raise RuntimeError(
"Cannot deserialize POD objects without kubernetes libraries installed!"
)
pod = PodGenerator.deserialize_model_dict(var)
return pod
elif type_ == DAT.TIMEDELTA:
return datetime.timedelta(seconds=var)
elif type_ == DAT.TIMEZONE:
return decode_timezone(var)
elif type_ == DAT.RELATIVEDELTA:
return decode_relativedelta(var)
elif type_ == DAT.SET:
return {cls._deserialize(v) for v in var}
elif type_ == DAT.TUPLE:
return tuple(cls._deserialize(v) for v in var)
elif type_ == DAT.PARAM:
return cls._deserialize_param(var)
else:
raise TypeError(f'Invalid type {type_!s} in deserialization.')
def run_next(self, next_job: KubernetesJobType) -> None:
"""
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).replace("\n", " "))
key, command, kube_executor_config, pod_template_file = next_job
dag_id, task_id, run_id, try_number = key
if command[0:3] != ["airflow", "tasks", "run"]:
raise ValueError('The command must start with ["airflow", "tasks", "run"].')
base_worker_pod = get_base_pod_from_template(pod_template_file, self.kube_config)
if not base_worker_pod:
raise AirflowException(
f"could not find a valid worker template yaml at {self.kube_config.pod_template_file}"
)
pod = PodGenerator.construct_pod(
namespace=self.namespace,
scheduler_job_id=self.scheduler_job_id,
pod_id=create_pod_id(dag_id, task_id),
dag_id=dag_id,
task_id=task_id,
kube_image=self.kube_config.kube_image,
try_number=try_number,
date=None,
run_id=run_id,
args=command,
pod_override_object=kube_executor_config,
base_worker_pod=base_worker_pod,
)
# Reconcile the pod generated by the Operator and the Pod
# generated by the .cfg file
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.run_pod_async(pod, **self.kube_config.kube_client_request_args)
self.log.debug("Kubernetes Job created!")
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 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
if command[0:2] != ["airflow", "run"]:
raise ValueError('The command must start with ["airflow", "run"].')
pod = PodGenerator.construct_pod(
namespace=self.namespace,
worker_uuid=self.worker_uuid,
pod_id=self._create_pod_id(dag_id, task_id),
dag_id=pod_generator.make_safe_label_value(dag_id),
task_id=pod_generator.make_safe_label_value(task_id),
try_number=try_number,
kube_image=self.kube_config.kube_image,
date=execution_date,
command=command,
pod_override_object=kube_executor_config,
base_worker_pod=self.worker_configuration_pod)
sanitized_pod = self.launcher._client.api_client.sanitize_for_serialization(
pod)
json_pod = json.dumps(sanitized_pod, indent=2)
self.log.debug('Pod Creation Request before mutation: \n%s', json_pod)
# Reconcile the pod generated by the Operator and the Pod
# generated by the .cfg file
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 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 _default(obj):
"""Convert dates and numpy objects in a json serializable format."""
if isinstance(obj, datetime):
return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
elif isinstance(obj, date):
return obj.strftime('%Y-%m-%d')
elif isinstance(obj, Decimal):
_, _, exponent = obj.as_tuple()
if exponent >= 0: # No digits after the decimal point.
return int(obj)
# Technically lossy due to floating point errors, but the best we
# can do without implementing a custom encode function.
return float(obj)
elif isinstance(
obj,
(
np.int_,
np.intc,
np.intp,
np.int8,
np.int16,
np.int32,
np.int64,
np.uint8,
np.uint16,
np.uint32,
np.uint64,
),
):
return int(obj)
elif isinstance(obj, np.bool_):
return bool(obj)
elif isinstance(
obj, (np.float_, np.float16, np.float32, np.float64, np.complex_, np.complex64, np.complex128)
):
return float(obj)
elif k8s is not None and isinstance(obj, (k8s.V1Pod, k8s.V1ResourceRequirements)):
from airflow.kubernetes.pod_generator import PodGenerator
return PodGenerator.serialize_pod(obj)
raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable")
def _deserialize(cls, encoded_var: Any) -> Any: # pylint: disable=too-many-return-statements
"""Helper function of depth first search for deserialization."""
# JSON primitives (except for dict) are not encoded.
if cls._is_primitive(encoded_var):
return encoded_var
elif isinstance(encoded_var, list):
return [cls._deserialize(v) for v in encoded_var]
if not isinstance(encoded_var, dict):
raise ValueError(
f"The encoded_var should be dict and is {type(encoded_var)}")
var = encoded_var[Encoding.VAR]
type_ = encoded_var[Encoding.TYPE]
if type_ == DAT.DICT:
return {k: cls._deserialize(v) for k, v in var.items()}
elif type_ == DAT.DAG:
return SerializedDAG.deserialize_dag(var)
elif type_ == DAT.OP:
return SerializedBaseOperator.deserialize_operator(var)
elif type_ == DAT.DATETIME:
return pendulum.from_timestamp(var)
elif type_ == DAT.POD:
pod = PodGenerator.deserialize_model_dict(var)
return pod
elif type_ == DAT.TIMEDELTA:
return datetime.timedelta(seconds=var)
elif type_ == DAT.TIMEZONE:
return Timezone(var)
elif type_ == DAT.RELATIVEDELTA:
if 'weekday' in var:
var['weekday'] = relativedelta.weekday(
*var['weekday']) # type: ignore
return relativedelta.relativedelta(**var)
elif type_ == DAT.SET:
return {cls._deserialize(v) for v in var}
elif type_ == DAT.TUPLE:
return tuple([cls._deserialize(v) for v in var])
else:
raise TypeError(
'Invalid type {!s} in deserialization.'.format(type_))
def execute_async(
self,
key: TaskInstanceKey,
command: CommandType,
queue: Optional[str] = None,
executor_config: Optional[Any] = None,
) -> None:
"""Executes task asynchronously"""
self.log.info('Add task %s with command %s with executor_config %s',
key, command, executor_config)
kube_executor_config = PodGenerator.from_obj(executor_config)
if executor_config:
pod_template_file = executor_config.get("pod_template_override",
None)
else:
pod_template_file = None
if not self.task_queue:
raise AirflowException(NOT_STARTED_MESSAGE)
self.event_buffer[key] = (State.QUEUED, self.scheduler_job_id)
self.task_queue.put(
(key, command, kube_executor_config, pod_template_file))
def run_next(self, next_job: KubernetesJobType) -> None:
"""
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
if isinstance(command, str):
command = [command]
if command[0] != "airflow":
raise ValueError(
'The first element of command must be equal to "airflow".')
pod = PodGenerator.construct_pod(
namespace=self.namespace,
worker_uuid=self.worker_uuid,
pod_id=self._create_pod_id(dag_id, task_id),
dag_id=pod_generator.make_safe_label_value(dag_id),
task_id=pod_generator.make_safe_label_value(task_id),
try_number=try_number,
date=self._datetime_to_label_safe_datestring(execution_date),
command=command,
kube_executor_config=kube_executor_config,
worker_config=self.worker_configuration_pod)
# Reconcile the pod generated by the Operator and the Pod
# generated by the .cfg file
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 _default(obj):
"""Convert dates and numpy objects in a json serializable format."""
if isinstance(obj, datetime):
return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
elif isinstance(obj, date):
return obj.strftime('%Y-%m-%d')
elif isinstance(
obj, (np.int_, np.intc, np.intp, np.int8, np.int16, np.int32,
np.int64, np.uint8, np.uint16, np.uint32, np.uint64)):
return int(obj)
elif isinstance(obj, np.bool_):
return bool(obj)
elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64,
np.complex_, np.complex64, np.complex128)):
return float(obj)
elif k8s is not None and isinstance(obj, k8s.V1Pod):
from airflow.kubernetes.pod_generator import PodGenerator
return PodGenerator.serialize_pod(obj)
raise TypeError(
f"Object of type '{obj.__class__.__name__}' is not JSON serializable"
)
def test_deserialize_model_string(self):
fixture = """
apiVersion: v1
kind: Pod
metadata:
name: memory-demo
namespace: mem-example
spec:
containers:
- name: memory-demo-ctr
image: apache/airflow:stress-2020.07.10-1.0.4
resources:
limits:
memory: "200Mi"
requests:
memory: "100Mi"
command: ["stress"]
args: ["--vm", "1", "--vm-bytes", "150M", "--vm-hang", "1"]
"""
result = PodGenerator.deserialize_model_file(fixture)
sanitized_res = self.k8s_client.sanitize_for_serialization(result)
assert sanitized_res == self.deserialize_result
def _adopt_completed_pods(self, kube_client: kubernetes.client.CoreV1Api):
"""
Patch completed pod so that the KubernetesJobWatcher can delete it.
:param kube_client: kubernetes client for speaking to kube API
"""
kwargs = {
'field_selector': "status.phase=Succeeded",
'label_selector': 'kubernetes_executor=True',
}
pod_list = kube_client.list_namespaced_pod(namespace=self.kube_config.kube_namespace, **kwargs)
for pod in pod_list.items:
self.log.info("Attempting to adopt pod %s", pod.metadata.name)
pod.metadata.labels['airflow-worker'] = str(self.scheduler_job_id)
try:
kube_client.patch_namespaced_pod(
name=pod.metadata.name,
namespace=pod.metadata.namespace,
body=PodGenerator.serialize_pod(pod),
)
except ApiException as e:
self.log.info("Failed to adopt pod %s. Reason: %s", pod.metadata.name, e)
def execute_async(
self,
key: TaskInstanceKey,
command: CommandType,
queue: Optional[str] = None,
executor_config: Optional[Any] = None,
) -> None:
"""Executes task asynchronously"""
self.log.info('Add task %s with command %s with executor_config %s', key, command, executor_config)
try:
kube_executor_config = PodGenerator.from_obj(executor_config)
except Exception: # pylint: disable=broad-except
self.log.error("Invalid executor_config for %s", key)
self.fail(key=key, info="Invalid executor_config passed")
return
if executor_config:
pod_template_file = executor_config.get("pod_template_file", None)
else:
pod_template_file = None
if not self.task_queue:
raise AirflowException(NOT_STARTED_MESSAGE)
self.event_buffer[key] = (State.QUEUED, self.scheduler_job_id)
self.task_queue.put((key, command, kube_executor_config, pod_template_file))
def test_create_pod_id(self):
for dag_id, task_id in self._cases():
pod_name = PodGenerator.make_unique_pod_id(
create_pod_id(dag_id, task_id))
assert self._is_valid_pod_id(pod_name)
"_outlets": [],
"ui_color": "#f0ede4",
"ui_fgcolor": "#000",
"template_fields": ['bash_command', 'env'],
"template_fields_renderers": {'bash_command': 'bash', 'env': 'json'},
"bash_command": "echo {{ task.task_id }}",
'label': 'bash_task',
"_task_type": "BashOperator",
"_task_module": "airflow.operators.bash",
"pool": "default_pool",
"executor_config": {
'__type': 'dict',
'__var': {
"pod_override": {
'__type': 'k8s.V1Pod',
'__var': PodGenerator.serialize_pod(executor_config_pod),
}
},
},
},
{
"task_id": "custom_task",
"retries": 1,
"retry_delay": 300.0,
"sla": 100.0,
"_downstream_task_ids": [],
"_inlets": [],
"_is_dummy": False,
"_outlets": [],
"_operator_extra_links": [{"tests.test_utils.mock_operators.CustomOpLink": {}}],
"ui_color": "#fff",
def test_from_obj_with_resources_object(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
result = PodGenerator.from_obj({
"KubernetesExecutor": {
"annotations": {
"test": "annotation"
},
"volumes": [
{
"name": "example-kubernetes-test-volume",
"hostPath": {
"path": "/tmp/"
},
},
],
"volume_mounts": [
{
"mountPath": "/foo/",
"name": "example-kubernetes-test-volume",
},
],
"resources": {
"requests": {
"memory": "256Mi",
"cpu": "500m",
"ephemeral-storage": "2G",
"nvidia.com/gpu": "0"
},
"limits": {
"memory": "512Mi",
"cpu": "1000m",
"ephemeral-storage": "2G",
"nvidia.com/gpu": "0"
}
}
}
})
result = self.k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {
'test': 'annotation'
},
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'name':
'base',
'ports': [],
'volumeMounts':
[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume'
}],
'resources': {
'limits': {
'cpu': '1000m',
'ephemeral-storage': '2G',
'memory': '512Mi',
'nvidia.com/gpu': '0'
},
'requests': {
'cpu': '500m',
'ephemeral-storage': '2G',
'memory': '256Mi',
'nvidia.com/gpu': '0'
}
},
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume'
}],
}
}, result)
def test_from_obj_with_only_request_resources(self, mock_uuid):
self.maxDiff = None
mock_uuid.return_value = self.static_uuid
result = PodGenerator.from_obj({
"KubernetesExecutor": {
"annotations": {
"test": "annotation"
},
"volumes": [
{
"name": "example-kubernetes-test-volume",
"hostPath": {
"path": "/tmp/"
},
},
],
"volume_mounts": [
{
"mountPath": "/foo/",
"name": "example-kubernetes-test-volume",
},
],
'request_cpu':
"200m",
'request_memory':
"500Mi",
}
})
result = self.k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {
'test': 'annotation'
},
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'name':
'base',
'ports': [],
'resources': {
'requests': {
'cpu': '200m',
'memory': '500Mi',
},
},
'volumeMounts':
[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume'
}],
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume'
}],
}
}, result)
def patch_already_checked(self, pod: k8s.V1Pod):
"""Add an "already tried annotation to ensure we only retry once"""
pod.metadata.labels["already_checked"] = "True"
body = PodGenerator.serialize_pod(pod)
self.client.patch_namespaced_pod(pod.metadata.name,
pod.metadata.namespace, body)
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_deserialize_model_file(self):
path = sys.path[0] + '/tests/kubernetes/pod.yaml'
result = PodGenerator.deserialize_model_file(path)
sanitized_res = self.k8s_client.sanitize_for_serialization(result)
assert sanitized_res == self.deserialize_result
def patch_already_checked(self, pod: k8s.V1Pod):
"""Add an "already checked" annotation to ensure we don't reattach on retries"""
pod.metadata.labels[self.POD_CHECKED_KEY] = "True"
body = PodGenerator.serialize_pod(pod)
self.client.patch_namespaced_pod(pod.metadata.name,
pod.metadata.namespace, body)
def test_construct_pod_with_mutation(self, mock_uuid):
mock_uuid.return_value = self.static_uuid
worker_config = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name='gets-overridden-by-dynamic-args',
annotations={'should': 'stay'}),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name='doesnt-override',
resources=k8s.V1ResourceRequirements(
limits={
'cpu': '1m',
'memory': '1G'
}),
security_context=k8s.V1SecurityContext(
run_as_user=1))
]))
executor_config = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name='doesnt-override-either',
resources=k8s.V1ResourceRequirements(limits={
'cpu': '2m',
'memory': '2G'
}))
]))
result = PodGenerator.construct_pod(
'dag_id',
'task_id',
'pod_id',
3,
'date',
['command'],
executor_config,
worker_config,
'namespace',
'uuid',
)
sanitized_result = self.k8s_client.sanitize_for_serialization(result)
self.metadata.update({'annotations': {'should': 'stay'}})
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': self.metadata,
'spec': {
'containers': [{
'args': [],
'command': ['command'],
'env': [],
'envFrom': [],
'name': 'base',
'ports': [],
'resources': {
'limits': {
'cpu': '2m',
'memory': '2G'
}
},
'volumeMounts': [],
'securityContext': {
'runAsUser': 1
}
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': []
}
}, sanitized_result)
def test_attach_to_pod(self, mock_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()
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-' + static_uuid.hex
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [{
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'key': 'source_b',
'name': 'secret_b'
}
}
}],
'envFrom': [{
'secretRef': {
'name': 'secret_a'
}
}],
'image':
'airflow-worker:latest',
'name':
'base',
'ports': [],
'volumeMounts': [{
'mountPath': '/etc/foo',
'name': 'secretvol' + str(static_uuid),
'readOnly': True
}]
}],
'hostNetwork':
False,
'imagePullSecrets': [],
'volumes': [{
'name': 'secretvol' + str(static_uuid),
'secret': {
'secretName': 'secret_b'
}
}]
}
})
def test_create_pod_id(self):
for dag_id, task_id in self._cases():
pod_name = PodGenerator.make_unique_pod_id(
AirflowKubernetesScheduler._create_pod_id(dag_id, task_id))
self.assertTrue(self._is_valid_pod_id(pod_name))
def _serialize(
cls, var: Any
) -> Any: # Unfortunately there is no support for recursive types in mypy
"""Helper function of depth first search for serialization.
The serialization protocol is:
(1) keeping JSON supported types: primitives, dict, list;
(2) encoding other types as ``{TYPE: 'foo', VAR: 'bar'}``, the deserialization
step decode VAR according to TYPE;
(3) Operator has a special field CLASS to record the original class
name for displaying in UI.
"""
if cls._is_primitive(var):
# enum.IntEnum is an int instance, it causes json dumps error so we use its value.
if isinstance(var, enum.Enum):
return var.value
return var
elif isinstance(var, dict):
return cls._encode(
{str(k): cls._serialize(v)
for k, v in var.items()},
type_=DAT.DICT)
elif isinstance(var, list):
return [cls._serialize(v) for v in var]
elif HAS_KUBERNETES and isinstance(var, k8s.V1Pod):
json_pod = PodGenerator.serialize_pod(var)
return cls._encode(json_pod, type_=DAT.POD)
elif isinstance(var, DAG):
return SerializedDAG.serialize_dag(var)
elif isinstance(var, BaseOperator):
return SerializedBaseOperator.serialize_operator(var)
elif isinstance(var, cls._datetime_types):
return cls._encode(var.timestamp(), type_=DAT.DATETIME)
elif isinstance(var, datetime.timedelta):
return cls._encode(var.total_seconds(), type_=DAT.TIMEDELTA)
elif isinstance(var, Timezone):
return cls._encode(str(var.name), type_=DAT.TIMEZONE)
elif isinstance(var, relativedelta.relativedelta):
encoded = {
k: v
for k, v in var.__dict__.items() if not k.startswith("_") and v
}
if var.weekday and var.weekday.n:
# Every n'th Friday for example
encoded['weekday'] = [var.weekday.weekday, var.weekday.n]
elif var.weekday:
encoded['weekday'] = [var.weekday.weekday]
return cls._encode(encoded, type_=DAT.RELATIVEDELTA)
elif callable(var):
return str(get_python_source(var))
elif isinstance(var, set):
# FIXME: casts set to list in customized serialization in future.
try:
return cls._encode(sorted(cls._serialize(v) for v in var),
type_=DAT.SET)
except TypeError:
return cls._encode([cls._serialize(v) for v in var],
type_=DAT.SET)
elif isinstance(var, tuple):
# FIXME: casts tuple to list in customized serialization in future.
return cls._encode([cls._serialize(v) for v in var],
type_=DAT.TUPLE)
elif isinstance(var, TaskGroup):
return SerializedTaskGroup.serialize_task_group(var)
else:
log.debug('Cast type %s to str in serialization.', type(var))
return str(var)
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_from_obj(self):
result = PodGenerator.from_obj(
{
"pod_override": k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(name="foo", annotations={"test": "annotation"}),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
volume_mounts=[
k8s.V1VolumeMount(
mount_path="/foo/", name="example-kubernetes-test-volume"
)
],
)
],
volumes=[
k8s.V1Volume(
name="example-kubernetes-test-volume",
host_path=k8s.V1HostPathVolumeSource(path="/tmp/"),
)
],
),
)
}
)
result = self.k8s_client.sanitize_for_serialization(result)
assert {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'foo',
'annotations': {'test': 'annotation'},
},
'spec': {
'containers': [
{
'name': 'base',
'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}],
}
],
'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}],
},
} == result
result = PodGenerator.from_obj(
{
"KubernetesExecutor": {
"annotations": {"test": "annotation"},
"volumes": [
{
"name": "example-kubernetes-test-volume",
"hostPath": {"path": "/tmp/"},
},
],
"volume_mounts": [
{
"mountPath": "/foo/",
"name": "example-kubernetes-test-volume",
},
],
}
}
)
result_from_pod = PodGenerator.from_obj(
{
"pod_override": k8s.V1Pod(
metadata=k8s.V1ObjectMeta(annotations={"test": "annotation"}),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
volume_mounts=[
k8s.V1VolumeMount(
name="example-kubernetes-test-volume", mount_path="/foo/"
)
],
)
],
volumes=[k8s.V1Volume(name="example-kubernetes-test-volume", host_path="/tmp/")],
),
)
}
)
result = self.k8s_client.sanitize_for_serialization(result)
result_from_pod = self.k8s_client.sanitize_for_serialization(result_from_pod)
expected_from_pod = {
'metadata': {'annotations': {'test': 'annotation'}},
'spec': {
'containers': [
{
'name': 'base',
'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}],
}
],
'volumes': [{'hostPath': '/tmp/', 'name': 'example-kubernetes-test-volume'}],
},
}
assert (
result_from_pod == expected_from_pod
), "There was a discrepency between KubernetesExecutor and pod_override"
assert {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {'test': 'annotation'},
},
'spec': {
'containers': [
{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'name': 'base',
'ports': [],
'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}],
}
],
'hostNetwork': False,
'imagePullSecrets': [],
'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}],
},
} == 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.
"""
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="base",
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)
return pod
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_from_obj(self, mock_uuid):
mock_uuid.return_value = '0'
result = PodGenerator.from_obj({
"KubernetesExecutor": {
"annotations": {
"test": "annotation"
},
"volumes": [
{
"name": "example-kubernetes-test-volume",
"hostPath": {
"path": "/tmp/"
},
},
],
"volume_mounts": [
{
"mountPath": "/foo/",
"name": "example-kubernetes-test-volume",
},
],
"securityContext": {
"runAsUser": 1000
}
}
})
result = self.k8s_client.sanitize_for_serialization(result)
self.assertEqual(
{
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'annotations': {
'test': 'annotation'
},
},
'spec': {
'containers': [{
'args': [],
'command': [],
'env': [],
'envFrom': [],
'name':
'base',
'ports': [],
'volumeMounts':
[{
'mountPath': '/foo/',
'name': 'example-kubernetes-test-volume'
}],
}],
'imagePullSecrets': [],
'volumes': [{
'hostPath': {
'path': '/tmp/'
},
'name': 'example-kubernetes-test-volume'
}],
}
}, result)
