def test_make_pod_assert_labels(self):
# Tests the pod created has all the expected labels set
self.kube_config.dags_folder = 'dags'
worker_config = WorkerConfiguration(self.kube_config)
execution_date = parser.parse('2019-11-21 11:08:22.920875')
pod = PodGenerator.construct_pod(
"test_dag_id",
"test_task_id",
"test_pod_id",
1,
execution_date,
["bash -c 'ls /'"],
None,
worker_config.as_pod(),
"default",
"sample-uuid",
)
expected_labels = {
'airflow-worker': 'sample-uuid',
'airflow_version': airflow_version.replace('+', '-'),
'dag_id': 'test_dag_id',
'execution_date': datetime_to_label_safe_datestring(execution_date),
'kubernetes_executor': 'True',
'my_label': 'label_id',
'task_id': 'test_task_id',
'try_number': '1'
}
self.assertEqual(pod.metadata.labels, expected_labels)
def generate_pod_yaml(args):
"""Generates yaml files for each task in the DAG. Used for testing output of KubernetesExecutor"""
execution_date = args.execution_date
dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
yaml_output_path = args.output_path
kube_config = KubeConfig()
for task in dag.tasks:
ti = TaskInstance(task, execution_date)
pod = PodGenerator.construct_pod(
dag_id=args.dag_id,
task_id=ti.task_id,
pod_id=create_pod_id(args.dag_id, ti.task_id),
try_number=ti.try_number,
kube_image=kube_config.kube_image,
date=ti.execution_date,
args=ti.command_as_list(),
pod_override_object=PodGenerator.from_obj(ti.executor_config),
scheduler_job_id="worker-config",
namespace=kube_config.executor_namespace,
base_worker_pod=PodGenerator.deserialize_model_file(kube_config.pod_template_file),
)
pod_mutation_hook(pod)
api_client = ApiClient()
date_string = pod_generator.datetime_to_label_safe_datestring(execution_date)
yaml_file_name = f"{args.dag_id}_{ti.task_id}_{date_string}.yml"
os.makedirs(os.path.dirname(yaml_output_path + "/airflow_yaml_output/"), exist_ok=True)
with open(yaml_output_path + "/airflow_yaml_output/" + yaml_file_name, "w") as output:
sanitized_pod = api_client.sanitize_for_serialization(pod)
output.write(yaml.dump(sanitized_pod))
print(f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
def clear_not_launched_queued_tasks(self, session=None) -> None:
"""
Tasks can end up in a "Queued" state through either the executor being
abruptly shut down (leaving a non-empty task_queue on this executor)
or when a rescheduled/deferred operator comes back up for execution
(with the same try_number) before the pod of its previous incarnation
has been fully removed (we think).
This method checks each of those tasks to see if the corresponding pod
is around, and if not, and there's no matching entry in our own
task_queue, marks it for re-execution.
"""
self.log.debug("Clearing tasks that have not been launched")
if not self.kube_client:
raise AirflowException(NOT_STARTED_MESSAGE)
queued_tasks = session.query(TaskInstance).filter(TaskInstance.state == State.QUEUED).all()
self.log.info('Found %s queued task instances', len(queued_tasks))
# Go through the "last seen" dictionary and clean out old entries
allowed_age = self.kube_config.worker_pods_queued_check_interval * 3
for key, timestamp in list(self.last_handled.items()):
if time.time() - timestamp > allowed_age:
del self.last_handled[key]
for task in queued_tasks:
self.log.debug("Checking task %s", task)
# Check to see if we've handled it ourselves recently
if task.key in self.last_handled:
continue
# Build the pod selector
dict_string = "dag_id={},task_id={},airflow-worker={}".format(
pod_generator.make_safe_label_value(task.dag_id),
pod_generator.make_safe_label_value(task.task_id),
pod_generator.make_safe_label_value(str(self.scheduler_job_id)),
)
kwargs = dict(label_selector=dict_string)
if self.kube_config.kube_client_request_args:
kwargs.update(**self.kube_config.kube_client_request_args)
# Try run_id first
kwargs['label_selector'] += ',run_id=' + pod_generator.make_safe_label_value(task.run_id)
pod_list = self.kube_client.list_namespaced_pod(self.kube_config.kube_namespace, **kwargs)
if pod_list.items:
continue
# Fallback to old style of using execution_date
kwargs['label_selector'] = dict_string + ',exectuion_date={}'.format(
pod_generator.datetime_to_label_safe_datestring(task.execution_date)
)
pod_list = self.kube_client.list_namespaced_pod(self.kube_config.kube_namespace, **kwargs)
if pod_list.items:
continue
self.log.info('TaskInstance: %s found in queued state but was not launched, rescheduling', task)
session.query(TaskInstance).filter(
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.run_id == task.run_id,
).update({TaskInstance.state: State.SCHEDULED})
def test_execution_date_serialize_deserialize(self):
datetime_obj = datetime.now()
serialized_datetime = pod_generator.datetime_to_label_safe_datestring(
datetime_obj)
new_datetime_obj = pod_generator.label_safe_datestring_to_datetime(
serialized_datetime)
assert datetime_obj == new_datetime_obj
def clear_not_launched_queued_tasks(self, session=None) -> None:
"""
If the airflow scheduler restarts with pending "Queued" tasks, the tasks may or
may not
have been launched. Thus on starting up the scheduler let's check every
"Queued" task to
see if it has been launched (ie: if there is a corresponding pod on kubernetes)
If it has been launched then do nothing, otherwise reset the state to "None" so
the task
will be rescheduled
This will not be necessary in a future version of airflow in which there is
proper support
for State.LAUNCHED
"""
self.log.debug("Clearing tasks that have not been launched")
if not self.kube_client:
raise AirflowException(NOT_STARTED_MESSAGE)
queued_tasks = session \
.query(TaskInstance) \
.filter(TaskInstance.state == State.QUEUED).all()
self.log.info(
'When executor started up, found %s queued task instances',
len(queued_tasks)
)
for task in queued_tasks:
# pylint: disable=protected-access
self.log.debug("Checking task %s", task)
dict_string = (
"dag_id={},task_id={},execution_date={},airflow-worker={}".format(
pod_generator.make_safe_label_value(task.dag_id),
pod_generator.make_safe_label_value(task.task_id),
pod_generator.datetime_to_label_safe_datestring(
task.execution_date
),
self.scheduler_job_id
)
)
# pylint: enable=protected-access
kwargs = dict(label_selector=dict_string)
if self.kube_config.kube_client_request_args:
for key, value in self.kube_config.kube_client_request_args.items():
kwargs[key] = value
pod_list = self.kube_client.list_namespaced_pod(
self.kube_config.kube_namespace, **kwargs)
if not pod_list.items:
self.log.info(
'TaskInstance: %s found in queued state but was not launched, '
'rescheduling', task
)
session.query(TaskInstance).filter(
TaskInstance.dag_id == task.dag_id,
TaskInstance.task_id == task.task_id,
TaskInstance.execution_date == task.execution_date
).update({TaskInstance.state: State.NONE})
def generate_pod_yaml(args):
"""Generates yaml files for each task in the DAG. Used for testing output of KubernetesExecutor"""
from kubernetes.client.api_client import ApiClient
from airflow.executors.kubernetes_executor import AirflowKubernetesScheduler, KubeConfig
from airflow.kubernetes import pod_generator
from airflow.kubernetes.pod_generator import PodGenerator
from airflow.kubernetes.worker_configuration import WorkerConfiguration
from airflow.settings import pod_mutation_hook
execution_date = args.execution_date
dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
yaml_output_path = args.output_path
kube_config = KubeConfig()
for task in dag.tasks:
ti = TaskInstance(task, execution_date)
pod = PodGenerator.construct_pod(
dag_id=args.dag_id,
task_id=ti.task_id,
pod_id=AirflowKubernetesScheduler._create_pod_id( # pylint: disable=W0212
args.dag_id, ti.task_id),
try_number=ti.try_number,
kube_image=kube_config.kube_image,
date=ti.execution_date,
command=ti.command_as_list(),
pod_override_object=PodGenerator.from_obj(ti.executor_config),
worker_uuid="worker-config",
namespace=kube_config.executor_namespace,
base_worker_pod=WorkerConfiguration(
kube_config=kube_config).as_pod())
pod_mutation_hook(pod)
api_client = ApiClient()
date_string = pod_generator.datetime_to_label_safe_datestring(
execution_date)
yaml_file_name = f"{args.dag_id}_{ti.task_id}_{date_string}.yml"
os.makedirs(os.path.dirname(yaml_output_path +
"/airflow_yaml_output/"),
exist_ok=True)
with open(yaml_output_path + "/airflow_yaml_output/" + yaml_file_name,
"w") as output:
sanitized_pod = api_client.sanitize_for_serialization(pod)
output.write(yaml.dump(sanitized_pod))
print(
f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
def setUp(self):
self.static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48')
self.deserialize_result = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {'name': 'memory-demo', 'namespace': 'mem-example'},
'spec': {
'containers': [
{
'args': ['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'image': 'apache/airflow:stress-2020.07.10-1.0.4',
'name': 'memory-demo-ctr',
'resources': {'limits': {'memory': '200Mi'}, 'requests': {'memory': '100Mi'}},
}
]
},
}
self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'}
self.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'),
]
self.execution_date = parser.parse('2020-08-24 00:00:00.000000')
self.execution_date_label = datetime_to_label_safe_datestring(self.execution_date)
self.dag_id = 'dag_id'
self.task_id = 'task_id'
self.try_number = 3
self.labels = {
'airflow-worker': 'uuid',
'dag_id': self.dag_id,
'execution_date': self.execution_date_label,
'task_id': self.task_id,
'try_number': str(self.try_number),
'airflow_version': __version__.replace('+', '-'),
'kubernetes_executor': 'True',
}
self.annotations = {
'dag_id': self.dag_id,
'task_id': self.task_id,
'execution_date': self.execution_date.isoformat(),
'try_number': str(self.try_number),
}
self.metadata = {
'labels': self.labels,
'name': 'pod_id-' + self.static_uuid.hex,
'namespace': 'namespace',
'annotations': self.annotations,
}
self.resources = k8s.V1ResourceRequirements(
requests={
"cpu": 1,
"memory": "1Gi",
"ephemeral-storage": "2Gi",
},
limits={"cpu": 2, "memory": "2Gi", "ephemeral-storage": "4Gi", 'nvidia.com/gpu': 1},
)
self.k8s_client = ApiClient()
self.expected = k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
namespace="default",
name='myapp-pod-' + self.static_uuid.hex,
labels={'app': 'myapp'},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name='base',
image='busybox',
command=['sh', '-c', 'echo Hello Kubernetes!'],
env=[
k8s.V1EnvVar(name='ENVIRONMENT', value='prod'),
k8s.V1EnvVar(
name="LOG_LEVEL",
value='warning',
),
k8s.V1EnvVar(
name='TARGET',
value_from=k8s.V1EnvVarSource(
secret_key_ref=k8s.V1SecretKeySelector(name='secret_b', key='source_b')
),
),
],
env_from=[
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_a')),
k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name='configmap_b')),
k8s.V1EnvFromSource(secret_ref=k8s.V1SecretEnvSource(name='secret_a')),
],
ports=[k8s.V1ContainerPort(name="foo", container_port=1234)],
resources=k8s.V1ResourceRequirements(
requests={'memory': '100Mi'},
limits={
'memory': '200Mi',
},
),
)
],
security_context=k8s.V1PodSecurityContext(
fs_group=2000,
run_as_user=1000,
),
host_network=True,
image_pull_secrets=[
k8s.V1LocalObjectReference(name="pull_secret_a"),
k8s.V1LocalObjectReference(name="pull_secret_b"),
],
),
)
def test_pod_template_file_override_in_executor_config(
self, mock_get_kube_client, mock_run_pod_async):
current_folder = pathlib.Path(__file__).parent.absolute()
template_file = str(
(current_folder / "kubernetes_executor_template_files" /
"basic_template.yaml").absolute())
mock_kube_client = mock.patch('kubernetes.client.CoreV1Api',
autospec=True)
mock_get_kube_client.return_value = mock_kube_client
with conf_vars({('kubernetes', 'pod_template_file'): ''}):
executor = self.kubernetes_executor
executor.start()
assert executor.event_buffer == {}
assert executor.task_queue.empty()
execution_date = datetime.utcnow()
executor.execute_async(
key=('dag', 'task', execution_date, 1),
queue=None,
command=['airflow', 'tasks', 'run', 'true', 'some_parameter'],
executor_config={
"pod_template_file":
template_file,
"pod_override":
k8s.V1Pod(
metadata=k8s.V1ObjectMeta(
labels={"release": "stable"}),
spec=k8s.V1PodSpec(containers=[
k8s.V1Container(name="base", image="airflow:3.6")
], ),
),
},
)
assert not executor.task_queue.empty()
task = executor.task_queue.get_nowait()
_, _, expected_executor_config, expected_pod_template_file = task
# Test that the correct values have been put to queue
assert expected_executor_config.metadata.labels == {
'release': 'stable'
}
assert expected_pod_template_file == template_file
self.kubernetes_executor.kube_scheduler.run_next(task)
mock_run_pod_async.assert_called_once_with(
k8s.V1Pod(
api_version="v1",
kind="Pod",
metadata=k8s.V1ObjectMeta(
name=mock.ANY,
namespace="default",
annotations={
'dag_id': 'dag',
'execution_date': execution_date.isoformat(),
'task_id': 'task',
'try_number': '1',
},
labels={
'airflow-worker':
'5',
'airflow_version':
mock.ANY,
'dag_id':
'dag',
'execution_date':
datetime_to_label_safe_datestring(execution_date),
'kubernetes_executor':
'True',
'mylabel':
'foo',
'release':
'stable',
'task_id':
'task',
'try_number':
'1',
},
),
spec=k8s.V1PodSpec(
containers=[
k8s.V1Container(
name="base",
image="airflow:3.6",
args=[
'airflow', 'tasks', 'run', 'true',
'some_parameter'
],
env=[
k8s.V1EnvVar(
name='AIRFLOW_IS_K8S_EXECUTOR_POD',
value='True')
],
)
],
image_pull_secrets=[
k8s.V1LocalObjectReference(name='airflow-registry')
],
scheduler_name='default-scheduler',
security_context=k8s.V1PodSecurityContext(
fs_group=50000, run_as_user=50000),
),
))
def setUp(self):
self.static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48')
self.deserialize_result = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'memory-demo',
'namespace': 'mem-example'
},
'spec': {
'containers': [{
'args':
['--vm', '1', '--vm-bytes', '150M', '--vm-hang', '1'],
'command': ['stress'],
'image':
'polinux/stress',
'name':
'memory-demo-ctr',
'resources': {
'limits': {
'memory': '200Mi'
},
'requests': {
'memory': '100Mi'
}
}
}]
}
}
self.envs = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'}
self.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'),
]
self.execution_date = parser.parse('2020-08-24 00:00:00.000000')
self.execution_date_label = datetime_to_label_safe_datestring(
self.execution_date)
self.dag_id = 'dag_id'
self.task_id = 'task_id'
self.try_number = 3
self.labels = {
'airflow-worker': 'uuid',
'dag_id': self.dag_id,
'execution_date': self.execution_date_label,
'task_id': self.task_id,
'try_number': str(self.try_number),
'airflow_version': mock.ANY,
'kubernetes_executor': 'True'
}
self.annotations = {
'dag_id': self.dag_id,
'task_id': self.task_id,
'execution_date': self.execution_date.isoformat(),
'try_number': str(self.try_number),
}
self.metadata = {
'labels': self.labels,
'name': 'pod_id-' + self.static_uuid.hex,
'namespace': 'namespace',
'annotations': self.annotations,
}
self.resources = Resources('1Gi', 1, '2Gi', '2Gi', 2, 1, '4Gi')
self.k8s_client = ApiClient()
self.expected = {
'apiVersion': 'v1',
'kind': 'Pod',
'metadata': {
'name': 'myapp-pod-' + self.static_uuid.hex,
'labels': {
'app': 'myapp'
},
'namespace': 'default'
},
'spec': {
'containers': [{
'name':
'base',
'image':
'busybox',
'args': [],
'command': ['sh', '-c', 'echo Hello Kubernetes!'],
'env': [{
'name': 'ENVIRONMENT',
'value': 'prod'
}, {
'name': 'LOG_LEVEL',
'value': 'warning'
}, {
'name': 'TARGET',
'valueFrom': {
'secretKeyRef': {
'name': 'secret_b',
'key': 'source_b'
}
}
}],
'envFrom': [{
'configMapRef': {
'name': 'configmap_a'
}
}, {
'configMapRef': {
'name': 'configmap_b'
}
}, {
'secretRef': {
'name': 'secret_a'
}
}],
'resources': {
'requests': {
'memory': '1Gi',
'cpu': 1,
'ephemeral-storage': '2Gi'
},
'limits': {
'memory': '2Gi',
'cpu': 2,
'nvidia.com/gpu': 1,
'ephemeral-storage': '4Gi'
},
},
'ports': [{
'name': 'foo',
'containerPort': 1234
}],
'volumeMounts': [{
'mountPath':
'/etc/foo',
'name':
'secretvol' + str(self.static_uuid),
'readOnly':
True
}]
}],
'volumes': [{
'name': 'secretvol' + str(self.static_uuid),
'secret': {
'secretName': 'secret_b'
}
}],
'hostNetwork':
False,
'imagePullSecrets': [{
'name': 'pull_secret_a'
}, {
'name': 'pull_secret_b'
}],
'securityContext': {
'runAsUser': 1000,
'fsGroup': 2000,
},
}
}
