def test_worker_git_dags(self):
# Tests persistence volume config created when `git_repo` is set
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_folder = '/usr/local/airflow/dags'
self.kube_config.worker_dags_folder = '/usr/local/airflow/dags'
self.kube_config.git_sync_container_repository = 'gcr.io/google-containers/git-sync-amd64'
self.kube_config.git_sync_container_tag = 'v2.0.5'
self.kube_config.git_sync_container = 'gcr.io/google-containers/git-sync-amd64:v2.0.5'
self.kube_config.git_sync_init_container_name = 'git-sync-clone'
self.kube_config.git_subpath = 'dags_folder'
self.kube_config.git_sync_root = '/git'
self.kube_config.git_dags_folder_mount_point = '/usr/local/airflow/dags/repo/dags_folder'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']
self.assertTrue('emptyDir' in dag_volume[0])
self.assertEqual(self.kube_config.git_dags_folder_mount_point, dag_volume_mount[0]['mountPath'])
self.assertTrue(dag_volume_mount[0]['readOnly'])
init_container = worker_config._get_init_containers(volume_mounts)[0]
init_container_volume_mount = [mount for mount in init_container['volumeMounts']
if mount['name'] == 'airflow-dags']
self.assertEqual('git-sync-clone', init_container['name'])
self.assertEqual('gcr.io/google-containers/git-sync-amd64:v2.0.5', init_container['image'])
self.assertEqual(1, len(init_container_volume_mount))
self.assertFalse(init_container_volume_mount[0]['readOnly'])
def test_make_pod_git_sync_rev(self):
# Tests the pod created with git_sync_credentials_secret will get into the init container
self.kube_config.git_sync_rev = 'sampletag'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
self.kube_config.git_dags_folder_mount_point = 'dags'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'path'
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1, "bash -c 'ls /'",
kube_executor_config)
rev_env = {
'name': 'GIT_SYNC_REV',
'value': self.kube_config.git_sync_rev
}
self.assertIn(rev_env, pod.init_containers[0]["env"],
'The git_sync_rev env did not get into the init container')
def test_get_secrets(self):
# Test when secretRef is None and kube_secrets is not empty
self.kube_config.kube_secrets = {
'AWS_SECRET_KEY': 'airflow-secret=aws_secret_key',
'POSTGRES_PASSWORD': '******'
}
self.kube_config.env_from_secret_ref = None
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
secrets.sort(key=lambda secret: secret.deploy_target)
expected = [
Secret('env', 'AWS_SECRET_KEY', 'airflow-secret', 'aws_secret_key'),
Secret('env', 'POSTGRES_PASSWORD', 'airflow-secret', 'postgres_credentials')
]
self.assertListEqual(expected, secrets)
# Test when secret is not empty and kube_secrets is empty dict
self.kube_config.kube_secrets = {}
self.kube_config.env_from_secret_ref = 'secret_a,secret_b'
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
expected = [
Secret('env', None, 'secret_a'),
Secret('env', None, 'secret_b')
]
self.assertListEqual(expected, secrets)
def test_make_pod_git_sync_ssh_without_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct without known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'", kube_executor_config)
init_containers = worker_config._get_init_containers()
git_ssh_key_file = next((x['value'] for x in init_containers[0]['env']
if x['name'] == 'GIT_SSH_KEY_FILE'), None)
volume_mount_ssh_key = next(
(x['mountPath'] for x in init_containers[0]['volumeMounts']
if x['name'] == worker_config.git_sync_ssh_secret_volume_name),
None)
self.assertTrue(git_ssh_key_file)
self.assertTrue(volume_mount_ssh_key)
self.assertEqual(65533, pod.security_context['fsGroup'])
self.assertEqual(
git_ssh_key_file, volume_mount_ssh_key,
'The location where the git ssh secret is mounted'
' needs to be the same as the GIT_SSH_KEY_FILE path')
def test_worker_git_dags(self):
# Tests persistence volume config created when `git_repo` is set
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_folder = '/usr/local/airflow/dags'
self.kube_config.worker_dags_folder = '/usr/local/airflow/dags'
self.kube_config.git_sync_container_repository = 'gcr.io/google-containers/git-sync-amd64'
self.kube_config.git_sync_container_tag = 'v2.0.5'
self.kube_config.git_sync_container = 'gcr.io/google-containers/git-sync-amd64:v2.0.5'
self.kube_config.git_sync_init_container_name = 'git-sync-clone'
self.kube_config.git_subpath = 'dags_folder'
self.kube_config.git_sync_root = '/git'
self.kube_config.git_sync_run_as_user = 65533
self.kube_config.git_dags_folder_mount_point = '/usr/local/airflow/dags/repo/dags_folder'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']
self.assertTrue('emptyDir' in dag_volume[0])
self.assertEqual(self.kube_config.git_dags_folder_mount_point, dag_volume_mount[0]['mountPath'])
self.assertTrue(dag_volume_mount[0]['readOnly'])
init_container = worker_config._get_init_containers()[0]
init_container_volume_mount = [mount for mount in init_container['volumeMounts']
if mount['name'] == 'airflow-dags']
self.assertEqual('git-sync-clone', init_container['name'])
self.assertEqual('gcr.io/google-containers/git-sync-amd64:v2.0.5', init_container['image'])
self.assertEqual(1, len(init_container_volume_mount))
self.assertFalse(init_container_volume_mount[0]['readOnly'])
self.assertEqual(65533, init_container['securityContext']['runAsUser'])
def test_init_environment_using_git_sync_ssh_with_known_hosts(self):
# Tests the init environment created with git-sync SSH authentication option is correct
# with known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
self.kube_config.git_ssh_known_hosts_configmap_name = 'airflow-configmap'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0]['env']
self.assertTrue({
'name': 'GIT_SSH_KEY_FILE',
'value': '/etc/git-secret/ssh'
} in env)
self.assertTrue({'name': 'GIT_KNOWN_HOSTS', 'value': 'true'} in env)
self.assertTrue({
'name': 'GIT_SSH_KNOWN_HOSTS_FILE',
'value': '/etc/git-secret/known_hosts'
} in env)
self.assertTrue({'name': 'GIT_SYNC_SSH', 'value': 'true'} in env)
def test_make_pod_git_sync_ssh_without_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct without known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1, "bash -c 'ls /'",
kube_executor_config)
init_containers = worker_config._get_init_containers()
git_ssh_key_file = next((x['value'] for x in init_containers[0]['env']
if x['name'] == 'GIT_SSH_KEY_FILE'), None)
volume_mount_ssh_key = next((x['mountPath'] for x in init_containers[0]['volumeMounts']
if x['name'] == worker_config.git_sync_ssh_secret_volume_name),
None)
self.assertTrue(git_ssh_key_file)
self.assertTrue(volume_mount_ssh_key)
self.assertEqual(65533, pod.security_context['fsGroup'])
self.assertEqual(git_ssh_key_file,
volume_mount_ssh_key,
'The location where the git ssh secret is mounted'
' needs to be the same as the GIT_SSH_KEY_FILE path')
def test_init_environment_using_git_sync_user_with_known_hosts(self):
# Tests the init environment created with git-sync User authentication option is correct
# with known hosts file
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_user = '******'
self.kube_config.git_password = '******'
self.kube_config.git_ssh_known_hosts_configmap_name = 'airflow-configmap'
self.kube_config.git_ssh_key_secret_name = None
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0]['env']
self.assertFalse({'name': 'GIT_SSH_KEY_FILE', 'value': '/etc/git-secret/ssh'} in env)
self.assertTrue({'name': 'GIT_SYNC_USERNAME', 'value': 'git_user'} in env)
self.assertTrue({'name': 'GIT_SYNC_PASSWORD', 'value': 'git_password'} in env)
self.assertTrue({'name': 'GIT_KNOWN_HOSTS', 'value': 'true'} in env)
self.assertTrue({'name': 'GIT_SSH_KNOWN_HOSTS_FILE',
'value': '/etc/git-secret/known_hosts'} in env)
self.assertFalse({'name': 'GIT_SYNC_SSH', 'value': 'true'} in env)
def test_make_pod_git_sync_ssh_with_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct with known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
git_ssh_known_hosts_file = next(
(x['value'] for x in init_containers[0]['env']
if x['name'] == 'GIT_SSH_KNOWN_HOSTS_FILE'), None)
volume_mount_ssh_known_hosts_file = next(
(x['mountPath'] for x in init_containers[0]['volumeMounts']
if x['name'] == worker_config.git_sync_ssh_known_hosts_volume_name
), None)
self.assertTrue(git_ssh_known_hosts_file)
self.assertTrue(volume_mount_ssh_known_hosts_file)
self.assertEqual(
git_ssh_known_hosts_file, volume_mount_ssh_known_hosts_file,
'The location where the git known hosts file is mounted'
' needs to be the same as the GIT_SSH_KNOWN_HOSTS_FILE path')
def test_worker_configuration_no_subpaths(self):
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
for volume_or_mount in volumes + volume_mounts:
if volume_or_mount['name'] != 'airflow-config':
self.assertNotIn('subPath', volume_or_mount,
"subPath shouldn't be defined")
def test_worker_environment_no_dags_folder(self):
self.kube_config.airflow_configmap = ''
self.kube_config.git_dags_folder_mount_point = ''
self.kube_config.dags_folder = ''
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertNotIn('AIRFLOW__CORE__DAGS_FOLDER', env)
def test_worker_environment_no_dags_folder(self):
self.kube_config.airflow_configmap = ''
self.kube_config.git_dags_folder_mount_point = ''
self.kube_config.dags_folder = ''
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertNotIn('AIRFLOW__CORE__DAGS_FOLDER', env)
def test_worker_environment_when_dags_folder_specified(self):
dags_folder = '/workers/path/to/dags'
self.kube_config.worker_dags_folder = dags_folder
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_set_airflow_local_settings_configmap(self):
"""
Test that airflow_local_settings.py can be set via configmap by
checking volume & volume-mounts are set correctly.
"""
self.kube_config.airflow_home = '/usr/local/airflow'
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.airflow_local_settings_configmap = 'airflow-configmap'
self.kube_config.dags_folder = '/workers/path/to/dags'
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1, "bash -c 'ls /'",
kube_executor_config)
airflow_config_volume = [
volume for volume in pod.volumes if volume["name"] == 'airflow-config'
]
# Test that volume_name is found
self.assertEqual(1, len(airflow_config_volume))
# Test that config map exists
self.assertEqual(
{'configMap': {'name': 'airflow-configmap'}, 'name': 'airflow-config'},
airflow_config_volume[0]
)
# Test that 2 Volume Mounts exists and has 2 different mount-paths
# One for airflow.cfg
# Second for airflow_local_settings.py
volume_mounts = [
volume_mount for volume_mount in pod.volume_mounts
if volume_mount['name'] == 'airflow-config'
]
self.assertEqual(2, len(volume_mounts))
six.assertCountEqual(
self,
[
{
'mountPath': '/usr/local/airflow/airflow.cfg',
'name': 'airflow-config',
'readOnly': True,
'subPath': 'airflow.cfg',
},
{
'mountPath': '/usr/local/airflow/config/airflow_local_settings.py',
'name': 'airflow-config',
'readOnly': True,
'subPath': 'airflow_local_settings.py',
}
],
volume_mounts
)
def test_get_labels(self):
worker_config = WorkerConfiguration(self.kube_config)
labels = worker_config._get_labels({
'dag_id': 'override_dag_id',
})
self.assertEqual({
'my_label': 'label_id',
'dag_id': 'override_dag_id'
}, labels)
def test_worker_configuration_no_subpaths(self):
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
volumes_list = [value for value in volumes.values()]
volume_mounts_list = [value for value in volume_mounts.values()]
for volume_or_mount in volumes_list + volume_mounts_list:
if volume_or_mount['name'] != 'airflow-config':
self.assertNotIn('subPath', volume_or_mount,
"subPath shouldn't be defined")
def test_worker_configuration_no_subpaths(self):
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
for volume_or_mount in volumes + volume_mounts:
if volume_or_mount['name'] != 'airflow-config':
self.assertNotIn(
'subPath', volume_or_mount,
"subPath shouldn't be defined"
)
def test_worker_environment_when_dags_folder_specified(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_dags_folder_mount_point = ''
dags_folder = '/workers/path/to/dags'
self.kube_config.dags_folder = dags_folder
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_worker_environment_when_dags_folder_specified(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_dags_folder_mount_point = ''
dags_folder = '/workers/path/to/dags'
self.kube_config.dags_folder = dags_folder
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_worker_configuration_no_subpaths(self):
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
volumes_list = [value for value in volumes.values()]
volume_mounts_list = [value for value in volume_mounts.values()]
for volume_or_mount in volumes_list + volume_mounts_list:
if volume_or_mount['name'] != 'airflow-config':
self.assertNotIn(
'subPath', volume_or_mount,
"subPath shouldn't be defined"
)
def test_get_configmaps(self):
# Test when configmap is empty
self.kube_config.env_from_configmap_ref = ''
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_configmaps()
self.assertListEqual([], configmaps)
# test when configmap is not empty
self.kube_config.env_from_configmap_ref = 'configmap_a,configmap_b'
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_configmaps()
self.assertListEqual(['configmap_a', 'configmap_b'], configmaps)
def test_worker_pvc_dags(self):
# Tests persistence volume config created when `dags_volume_claim` is set
self.kube_config.dags_volume_claim = 'airflow-dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
dag_volume = [volume for volume in volumes if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts if mount['name'] == 'airflow-dags']
self.assertEqual('airflow-dags', dag_volume[0]['persistentVolumeClaim']['claimName'])
self.assertEqual(1, len(dag_volume_mount))
def __init__(self, kube_config, task_queue, result_queue, kube_client, worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config)
self.watcher_queue = multiprocessing.Queue()
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def test_worker_environment_dags_folder_using_git_sync(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'dags'
self.kube_config.git_dags_folder_mount_point = '/workers/path/to/dags'
dags_folder = '{}/{}/{}'.format(
self.kube_config.git_dags_folder_mount_point,
self.kube_config.git_sync_dest, self.kube_config.git_subpath)
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_worker_environment_dags_folder_using_git_sync(self):
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_sync_dest = 'repo'
self.kube_config.git_subpath = 'dags'
self.kube_config.git_dags_folder_mount_point = '/workers/path/to/dags'
dags_folder = '{}/{}/{}'.format(self.kube_config.git_dags_folder_mount_point,
self.kube_config.git_sync_dest,
self.kube_config.git_subpath)
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(dags_folder, env['AIRFLOW__CORE__DAGS_FOLDER'])
def test_init_environment_using_git_sync_run_as_user_empty(self):
# Tests if git_syn_run_as_user is none, then no securityContext created in init container
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.git_sync_run_as_user = ''
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
self.assertNotIn('securityContext', init_containers[0],
"securityContext shouldn't be defined")
def test_worker_container_dags(self):
# Tests that the 'airflow-dags' persistence volume is NOT created when `dags_in_image` is set
self.kube_config.dags_in_image = True
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']
init_containers = worker_config._get_init_containers(volume_mounts)
self.assertEqual(0, len(dag_volume))
self.assertEqual(0, len(dag_volume_mount))
self.assertEqual(0, len(init_containers))
def test_worker_container_dags(self):
# Tests that the 'airflow-dags' persistence volume is NOT created when `dags_in_image` is set
self.kube_config.dags_in_image = True
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']
init_containers = worker_config._get_init_containers()
self.assertEqual(0, len(dag_volume))
self.assertEqual(0, len(dag_volume_mount))
self.assertEqual(0, len(init_containers))
def test_worker_pvc_dags(self):
# Tests persistence volume config created when `dags_volume_claim` is set
self.kube_config.dags_volume_claim = 'airflow-dags'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
init_containers = worker_config._get_init_containers()
dag_volume = [volume for volume in volumes.values() if volume['name'] == 'airflow-dags']
dag_volume_mount = [mount for mount in volume_mounts.values() if mount['name'] == 'airflow-dags']
self.assertEqual('airflow-dags', dag_volume[0]['persistentVolumeClaim']['claimName'])
self.assertEqual(1, len(dag_volume_mount))
self.assertTrue(dag_volume_mount[0]['readOnly'])
self.assertEqual(0, len(init_containers))
def test_worker_with_subpaths(self):
self.kube_config.dags_volume_subpath = 'dags'
self.kube_config.logs_volume_subpath = 'logs'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config._get_volumes_and_mounts()
for volume in [value for value in volumes.values()]:
self.assertNotIn('subPath', volume,
"subPath isn't valid configuration for a volume")
for volume_mount in [value for value in volume_mounts.values()]:
if volume_mount['name'] != 'airflow-config':
self.assertIn(
'subPath', volume_mount,
"subPath should've been passed to volumeMount configuration"
)
def test_make_pod_git_sync_credentials_secret(self):
# Tests the pod created with git_sync_credentials_secret will get into the init container
self.kube_config.git_sync_credentials_secret = 'airflow-git-creds-secret'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'", kube_executor_config)
username_env = {
'name': 'GIT_SYNC_USERNAME',
'valueFrom': {
'secretKeyRef': {
'name': self.kube_config.git_sync_credentials_secret,
'key': 'GIT_SYNC_USERNAME'
}
}
}
password_env = {
'name': 'GIT_SYNC_PASSWORD',
'valueFrom': {
'secretKeyRef': {
'name': self.kube_config.git_sync_credentials_secret,
'key': 'GIT_SYNC_PASSWORD'
}
}
}
self.assertIn(
username_env, pod.init_containers[0]["env"],
'The username env for git credentials did not get into the init container'
)
self.assertIn(
password_env, pod.init_containers[0]["env"],
'The password env for git credentials did not get into the init container'
)
def test_make_pod_run_as_user_0(self):
# Tests the pod created with run-as-user 0 actually gets that in it's config
self.kube_config.worker_run_as_user = 0
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
self.kube_config.worker_fs_group = None
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1, "bash -c 'ls /'",
kube_executor_config)
self.assertEqual(0, pod.security_context['runAsUser'])
def test_worker_with_subpaths(self):
self.kube_config.dags_volume_subpath = 'dags'
self.kube_config.logs_volume_subpath = 'logs'
worker_config = WorkerConfiguration(self.kube_config)
volumes, volume_mounts = worker_config.init_volumes_and_mounts()
for volume in volumes:
self.assertNotIn(
'subPath', volume,
"subPath isn't valid configuration for a volume"
)
for volume_mount in volume_mounts:
if volume_mount['name'] != 'airflow-config':
self.assertIn(
'subPath', volume_mount,
"subPath should've been passed to volumeMount configuration"
)
def test_worker_generate_dag_volume_mount_path(self):
self.kube_config.git_dags_folder_mount_point = '/root/airflow/git/dags'
self.kube_config.dags_folder = '/root/airflow/dags'
worker_config = WorkerConfiguration(self.kube_config)
self.kube_config.dags_volume_claim = 'airflow-dags'
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = '/host/airflow/dags'
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path,
self.kube_config.git_dags_folder_mount_point)
def test_worker_generate_dag_volume_mount_path(self):
self.kube_config.git_dags_folder_mount_point = '/root/airflow/git/dags'
self.kube_config.dags_folder = '/root/airflow/dags'
worker_config = WorkerConfiguration(self.kube_config)
self.kube_config.dags_volume_claim = 'airflow-dags'
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = '/host/airflow/dags'
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path, self.kube_config.dags_folder)
self.kube_config.dags_volume_claim = ''
self.kube_config.dags_volume_host = ''
dag_volume_mount_path = worker_config.generate_dag_volume_mount_path()
self.assertEqual(dag_volume_mount_path,
self.kube_config.git_dags_folder_mount_point)
def test_kubernetes_environment_variables(self):
# Tests the kubernetes environment variables get copied into the worker pods
input_environment = {
'ENVIRONMENT': 'prod',
'LOG_LEVEL': 'warning'
}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
for key in input_environment:
self.assertIn(key, env)
self.assertIn(input_environment[key], env.values())
core_executor = 'AIRFLOW__CORE__EXECUTOR'
input_environment = {
core_executor: 'NotLocalExecutor'
}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(env[core_executor], 'LocalExecutor')
def test_get_secrets(self):
# Test when secretRef is None and kube_secrets is not empty
self.kube_config.kube_secrets = {
'AWS_SECRET_KEY': 'airflow-secret=aws_secret_key',
'POSTGRES_PASSWORD': '******'
}
self.kube_config.env_from_secret_ref = None
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
secrets.sort(key=lambda secret: secret.deploy_target)
expected = [
Secret('env', 'AWS_SECRET_KEY', 'airflow-secret',
'aws_secret_key'),
Secret('env', 'POSTGRES_PASSWORD', 'airflow-secret',
'postgres_credentials')
]
self.assertListEqual(expected, secrets)
# Test when secret is not empty and kube_secrets is empty dict
self.kube_config.kube_secrets = {}
self.kube_config.env_from_secret_ref = 'secret_a,secret_b'
worker_config = WorkerConfiguration(self.kube_config)
secrets = worker_config._get_secrets()
expected = [
Secret('env', None, 'secret_a'),
Secret('env', None, 'secret_b')
]
self.assertListEqual(expected, secrets)
def __init__(self, kube_config, task_queue, result_queue, kube_client, worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config)
self.watcher_queue = multiprocessing.Queue()
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def test_make_pod_with_executor_config(self):
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(
affinity=self.affinity_config,
tolerations=self.tolerations_config,
annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", str(uuid.uuid4()),
"test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), 1,
"bash -c 'ls /'", kube_executor_config)
self.assertTrue(pod.affinity['podAntiAffinity'] is not None)
self.assertEqual(
'app', pod.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']['matchExpressions'][0]['key'])
self.assertEqual(2, len(pod.tolerations))
self.assertEqual('prod', pod.tolerations[1]['key'])
def test_make_pod_with_executor_config(self):
worker_config = WorkerConfiguration(self.kube_config)
kube_executor_config = KubernetesExecutorConfig(affinity=self.affinity_config,
tolerations=self.tolerations_config,
annotations=[],
volumes=[],
volume_mounts=[]
)
pod = worker_config.make_pod("default", str(uuid.uuid4()), "test_pod_id", "test_dag_id",
"test_task_id", str(datetime.utcnow()), "bash -c 'ls /'",
kube_executor_config)
self.assertTrue(pod.affinity['podAntiAffinity'] is not None)
self.assertEqual('app',
pod.affinity['podAntiAffinity']
['requiredDuringSchedulingIgnoredDuringExecution'][0]
['labelSelector']
['matchExpressions'][0]
['key'])
self.assertEqual(2, len(pod.tolerations))
self.assertEqual('prod', pod.tolerations[1]['key'])
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)
kube_executor_config = KubernetesExecutorConfig(annotations=[],
volumes=[],
volume_mounts=[])
pod = worker_config.make_pod("default", "sample-uuid", "test_pod_id", "test_dag_id",
"test_task_id", "2019-11-21 11:08:22.920875", 1, "bash -c 'ls /'",
kube_executor_config)
expected_labels = {
'airflow-worker': 'sample-uuid',
'airflow_version': airflow_version.replace('+', '-'),
'dag_id': 'test_dag_id',
'execution_date': '2019-11-21 11:08:22.920875',
'kubernetes_executor': 'True',
'my_label': 'label_id',
'task_id': 'test_task_id',
'try_number': '1'
}
self.assertEqual(pod.labels, expected_labels)
def test_make_pod_git_sync_ssh_with_known_hosts(self):
# Tests the pod created with git-sync SSH authentication option is correct with known hosts
self.kube_config.airflow_configmap = 'airflow-configmap'
self.kube_config.git_ssh_secret_name = 'airflow-secrets'
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
worker_config = WorkerConfiguration(self.kube_config)
init_containers = worker_config._get_init_containers()
git_ssh_known_hosts_file = next((x['value'] for x in init_containers[0]['env']
if x['name'] == 'GIT_SSH_KNOWN_HOSTS_FILE'), None)
volume_mount_ssh_known_hosts_file = next(
(x['mountPath'] for x in init_containers[0]['volumeMounts']
if x['name'] == worker_config.git_sync_ssh_known_hosts_volume_name),
None)
self.assertTrue(git_ssh_known_hosts_file)
self.assertTrue(volume_mount_ssh_known_hosts_file)
self.assertEqual(git_ssh_known_hosts_file,
volume_mount_ssh_known_hosts_file,
'The location where the git known hosts file is mounted'
' needs to be the same as the GIT_SSH_KNOWN_HOSTS_FILE path')
def test_kubernetes_environment_variables_for_init_container(self):
self.kube_config.dags_volume_claim = None
self.kube_config.dags_volume_host = None
self.kube_config.dags_in_image = None
# Tests the kubernetes environment variables get copied into the worker pods
input_environment = {
'ENVIRONMENT': 'prod',
'LOG_LEVEL': 'warning'
}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
for key in input_environment:
self.assertIn(key, env)
self.assertIn(input_environment[key], env.values())
init_containers = worker_config._get_init_containers()
self.assertTrue(init_containers) # check not empty
env = init_containers[0]['env']
self.assertTrue({'name': 'ENVIRONMENT', 'value': 'prod'} in env)
self.assertTrue({'name': 'LOG_LEVEL', 'value': 'warning'} in env)
def test_get_configmaps(self):
# Test when configmap is empty
self.kube_config.env_from_configmap_ref = ''
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_configmaps()
self.assertListEqual([], configmaps)
# test when configmap is not empty
self.kube_config.env_from_configmap_ref = 'configmap_a,configmap_b'
worker_config = WorkerConfiguration(self.kube_config)
configmaps = worker_config._get_configmaps()
self.assertListEqual(['configmap_a', 'configmap_b'], configmaps)
def test_kubernetes_environment_variables(self):
# Tests the kubernetes environment variables get copied into the worker pods
input_environment = {'ENVIRONMENT': 'prod', 'LOG_LEVEL': 'warning'}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
for key in input_environment:
self.assertIn(key, env)
self.assertIn(input_environment[key], env.values())
core_executor = 'AIRFLOW__CORE__EXECUTOR'
input_environment = {core_executor: 'NotLocalExecutor'}
self.kube_config.kube_env_vars = input_environment
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertEqual(env[core_executor], 'LocalExecutor')
def test_get_labels(self):
worker_config = WorkerConfiguration(self.kube_config)
labels = worker_config._get_labels({
'dag_id': 'override_dag_id',
})
self.assertEqual({'my_label': 'label_id', 'dag_id': 'override_dag_id'}, labels)
class AirflowKubernetesScheduler(LoggingMixin):
def __init__(self, kube_config, task_queue, result_queue, kube_client, worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config)
self.watcher_queue = SynchronizedQueue()
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def _make_kube_watcher(self):
resource_version = KubeResourceVersion.get_current_resource_version()
watcher = KubernetesJobWatcher(self.namespace, self.watcher_queue,
resource_version, self.worker_uuid)
watcher.start()
return watcher
def _health_check_kube_watcher(self):
if self.kube_watcher.is_alive():
pass
else:
self.log.error(
'Error while health checking kube watcher process. '
'Process died for unknown reasons')
self.kube_watcher = self._make_kube_watcher()
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
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s", self.kube_config.kube_image)
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, kube_executor_config=kube_executor_config
)
# the watcher will monitor pods, so we do not block.
self.launcher.run_pod_async(pod)
self.log.debug("Kubernetes Job created!")
def delete_pod(self, pod_id):
if self.kube_config.delete_worker_pods:
try:
self.kube_client.delete_namespaced_pod(
pod_id, self.namespace, body=client.V1DeleteOptions())
except ApiException as e:
# If the pod is already deleted
if e.status != 404:
raise
def sync(self):
"""
The sync function checks the status of all currently running kubernetes jobs.
If a job is completed, it's status is placed in the result queue to
be sent back to the scheduler.
:return:
"""
self._health_check_kube_watcher()
while not self.watcher_queue.empty():
self.process_watcher_task()
def process_watcher_task(self):
pod_id, state, labels, resource_version = self.watcher_queue.get()
self.log.info(
'Attempting to finish pod; pod_id: %s; state: %s; labels: %s',
pod_id, state, labels
)
key = self._labels_to_key(labels=labels)
if key:
self.log.debug('finishing job %s - %s (%s)', key, state, pod_id)
self.result_queue.put((key, state, pod_id, resource_version))
@staticmethod
def _strip_unsafe_kubernetes_special_chars(string):
"""
Kubernetes only supports lowercase alphanumeric characters and "-" and "." in
the pod name
However, there are special rules about how "-" and "." can be used so let's
only keep
alphanumeric chars see here for detail:
https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
:param string: The requested Pod name
:return: ``str`` Pod name stripped of any unsafe characters
"""
return ''.join(ch.lower() for ind, ch in enumerate(string) if ch.isalnum())
@staticmethod
def _make_safe_pod_id(safe_dag_id, safe_task_id, safe_uuid):
r"""
Kubernetes pod names must be <= 253 chars and must pass the following regex for
validation
"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
:param safe_dag_id: a dag_id with only alphanumeric characters
:param safe_task_id: a task_id with only alphanumeric characters
:param random_uuid: a uuid
:return: ``str`` valid Pod name of appropriate length
"""
MAX_POD_ID_LEN = 253
safe_key = safe_dag_id + safe_task_id
safe_pod_id = safe_key[:MAX_POD_ID_LEN - len(safe_uuid) - 1] + "-" + safe_uuid
return safe_pod_id
@staticmethod
def _make_safe_label_value(string):
"""
Valid label values must be 63 characters or less and must be empty or begin and
end with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),
dots (.), and alphanumerics between.
If the label value is then greater than 63 chars once made safe, or differs in any
way from the original value sent to this function, then we need to truncate to
53chars, and append it with a unique hash.
"""
MAX_LABEL_LEN = 63
safe_label = re.sub(r'^[^a-z0-9A-Z]*|[^a-zA-Z0-9_\-\.]|[^a-z0-9A-Z]*$', '', string)
if len(safe_label) > MAX_LABEL_LEN or string != safe_label:
safe_hash = hashlib.md5(string.encode()).hexdigest()[:9]
safe_label = safe_label[:MAX_LABEL_LEN - len(safe_hash) - 1] + "-" + safe_hash
return safe_label
@staticmethod
def _create_pod_id(dag_id, task_id):
safe_dag_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
dag_id)
safe_task_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
task_id)
safe_uuid = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
uuid4().hex)
return AirflowKubernetesScheduler._make_safe_pod_id(safe_dag_id, safe_task_id,
safe_uuid)
@staticmethod
def _label_safe_datestring_to_datetime(string):
"""
Kubernetes doesn't permit ":" in labels. ISO datetime format uses ":" but not
"_", let's
replace ":" with "_"
:param string: str
:return: datetime.datetime object
"""
return parser.parse(string.replace('_plus_', '+').replace("_", ":"))
@staticmethod
def _datetime_to_label_safe_datestring(datetime_obj):
"""
Kubernetes doesn't like ":" in labels, since ISO datetime format uses ":" but
not "_" let's
replace ":" with "_"
:param datetime_obj: datetime.datetime object
:return: ISO-like string representing the datetime
"""
return datetime_obj.isoformat().replace(":", "_").replace('+', '_plus_')
def _labels_to_key(self, labels):
try_num = 1
try:
try_num = int(labels.get('try_number', '1'))
except ValueError:
self.log.warn("could not get try_number as an int: %s", labels.get('try_number', '1'))
try:
dag_id = labels['dag_id']
task_id = labels['task_id']
ex_time = self._label_safe_datestring_to_datetime(labels['execution_date'])
except Exception as e:
self.log.warn(
'Error while retrieving labels; labels: %s; exception: %s',
labels, e
)
return None
with create_session() as session:
tasks = (
session
.query(TaskInstance)
.filter_by(execution_date=ex_time).all()
)
self.log.info(
'Checking %s task instances.',
len(tasks)
)
for task in tasks:
if (
self._make_safe_label_value(task.dag_id) == dag_id and
self._make_safe_label_value(task.task_id) == task_id and
task.execution_date == ex_time
):
self.log.info(
'Found matching task %s-%s (%s) with current state of %s',
task.dag_id, task.task_id, task.execution_date, task.state
)
dag_id = task.dag_id
task_id = task.task_id
return (dag_id, task_id, ex_time, try_num)
self.log.warn(
'Failed to find and match task details to a pod; labels: %s',
labels
)
return None
class AirflowKubernetesScheduler(LoggingMixin):
def __init__(self, kube_config, task_queue, result_queue, session,
kube_client, worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config)
self.watcher_queue = multiprocessing.Queue()
self._session = session
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def _make_kube_watcher(self):
resource_version = KubeResourceVersion.get_current_resource_version(self._session)
watcher = KubernetesJobWatcher(self.namespace, self.watcher_queue,
resource_version, self.worker_uuid)
watcher.start()
return watcher
def _health_check_kube_watcher(self):
if self.kube_watcher.is_alive():
pass
else:
self.log.error(
'Error while health checking kube watcher process. '
'Process died for unknown reasons')
self.kube_watcher = self._make_kube_watcher()
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 = key
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s", self.kube_config.kube_image)
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=dag_id, task_id=task_id,
execution_date=self._datetime_to_label_safe_datestring(execution_date),
airflow_command=command, kube_executor_config=kube_executor_config
)
# the watcher will monitor pods, so we do not block.
self.launcher.run_pod_async(pod)
self.log.debug("Kubernetes Job created!")
def delete_pod(self, pod_id):
if self.kube_config.delete_worker_pods:
try:
self.kube_client.delete_namespaced_pod(
pod_id, self.namespace, body=client.V1DeleteOptions())
except ApiException as e:
# If the pod is already deleted
if e.status != 404:
raise
def sync(self):
"""
The sync function checks the status of all currently running kubernetes jobs.
If a job is completed, it's status is placed in the result queue to
be sent back to the scheduler.
:return:
"""
self._health_check_kube_watcher()
while not self.watcher_queue.empty():
self.process_watcher_task()
def process_watcher_task(self):
pod_id, state, labels, resource_version = self.watcher_queue.get()
self.log.info(
'Attempting to finish pod; pod_id: %s; state: %s; labels: %s',
pod_id, state, labels
)
key = self._labels_to_key(labels=labels)
if key:
self.log.debug('finishing job %s - %s (%s)', key, state, pod_id)
self.result_queue.put((key, state, pod_id, resource_version))
@staticmethod
def _strip_unsafe_kubernetes_special_chars(string):
"""
Kubernetes only supports lowercase alphanumeric characters and "-" and "." in
the pod name
However, there are special rules about how "-" and "." can be used so let's
only keep
alphanumeric chars see here for detail:
https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
:param string: The requested Pod name
:return: ``str`` Pod name stripped of any unsafe characters
"""
return ''.join(ch.lower() for ind, ch in enumerate(string) if ch.isalnum())
@staticmethod
def _make_safe_pod_id(safe_dag_id, safe_task_id, safe_uuid):
"""
Kubernetes pod names must be <= 253 chars and must pass the following regex for
validation
"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
:param safe_dag_id: a dag_id with only alphanumeric characters
:param safe_task_id: a task_id with only alphanumeric characters
:param random_uuid: a uuid
:return: ``str`` valid Pod name of appropriate length
"""
MAX_POD_ID_LEN = 253
safe_key = safe_dag_id + safe_task_id
safe_pod_id = safe_key[:MAX_POD_ID_LEN - len(safe_uuid) - 1] + "-" + safe_uuid
return safe_pod_id
@staticmethod
def _create_pod_id(dag_id, task_id):
safe_dag_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
dag_id)
safe_task_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
task_id)
safe_uuid = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
uuid4().hex)
return AirflowKubernetesScheduler._make_safe_pod_id(safe_dag_id, safe_task_id,
safe_uuid)
@staticmethod
def _label_safe_datestring_to_datetime(string):
"""
Kubernetes doesn't permit ":" in labels. ISO datetime format uses ":" but not
"_", let's
replace ":" with "_"
:param string: str
:return: datetime.datetime object
"""
return parser.parse(string.replace('_plus_', '+').replace("_", ":"))
@staticmethod
def _datetime_to_label_safe_datestring(datetime_obj):
"""
Kubernetes doesn't like ":" in labels, since ISO datetime format uses ":" but
not "_" let's
replace ":" with "_"
:param datetime_obj: datetime.datetime object
:return: ISO-like string representing the datetime
"""
return datetime_obj.isoformat().replace(":", "_").replace('+', '_plus_')
def _labels_to_key(self, labels):
try:
return (
labels['dag_id'], labels['task_id'],
self._label_safe_datestring_to_datetime(labels['execution_date']))
except Exception as e:
self.log.warn(
'Error while converting labels to key; labels: %s; exception: %s',
labels, e
)
return None
def test_worker_environment_no_dags_folder(self):
self.kube_config.worker_dags_folder = ''
worker_config = WorkerConfiguration(self.kube_config)
env = worker_config._get_environment()
self.assertNotIn('AIRFLOW__CORE__DAGS_FOLDER', env)
class AirflowKubernetesScheduler(LoggingMixin):
def __init__(self, kube_config, task_queue, result_queue, session,
kube_client, worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(kube_config=self.kube_config)
self.watcher_queue = multiprocessing.Queue()
self._session = session
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def _make_kube_watcher(self):
resource_version = KubeResourceVersion.get_current_resource_version(self._session)
watcher = KubernetesJobWatcher(self.namespace, self.watcher_queue,
resource_version, self.worker_uuid)
watcher.start()
return watcher
def _health_check_kube_watcher(self):
if self.kube_watcher.is_alive():
pass
else:
self.log.error(
'Error while health checking kube watcher process. '
'Process died for unknown reasons')
self.kube_watcher = self._make_kube_watcher()
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
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s", self.kube_config.kube_image)
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=dag_id, task_id=task_id,
execution_date=self._datetime_to_label_safe_datestring(execution_date),
airflow_command=command, kube_executor_config=kube_executor_config
)
# the watcher will monitor pods, so we do not block.
self.launcher.run_pod_async(pod)
self.log.debug("Kubernetes Job created!")
def delete_pod(self, pod_id):
if self.kube_config.delete_worker_pods:
try:
self.kube_client.delete_namespaced_pod(
pod_id, self.namespace, body=client.V1DeleteOptions())
except ApiException as e:
# If the pod is already deleted
if e.status != 404:
raise
def sync(self):
"""
The sync function checks the status of all currently running kubernetes jobs.
If a job is completed, it's status is placed in the result queue to
be sent back to the scheduler.
:return:
"""
self._health_check_kube_watcher()
while not self.watcher_queue.empty():
self.process_watcher_task()
def process_watcher_task(self):
pod_id, state, labels, resource_version = self.watcher_queue.get()
self.log.info(
'Attempting to finish pod; pod_id: %s; state: %s; labels: %s',
pod_id, state, labels
)
key = self._labels_to_key(labels=labels)
if key:
self.log.debug('finishing job %s - %s (%s)', key, state, pod_id)
self.result_queue.put((key, state, pod_id, resource_version))
@staticmethod
def _strip_unsafe_kubernetes_special_chars(string):
"""
Kubernetes only supports lowercase alphanumeric characters and "-" and "." in
the pod name
However, there are special rules about how "-" and "." can be used so let's
only keep
alphanumeric chars see here for detail:
https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
:param string: The requested Pod name
:return: ``str`` Pod name stripped of any unsafe characters
"""
return ''.join(ch.lower() for ind, ch in enumerate(string) if ch.isalnum())
@staticmethod
def _make_safe_pod_id(safe_dag_id, safe_task_id, safe_uuid):
"""
Kubernetes pod names must be <= 253 chars and must pass the following regex for
validation
"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
:param safe_dag_id: a dag_id with only alphanumeric characters
:param safe_task_id: a task_id with only alphanumeric characters
:param random_uuid: a uuid
:return: ``str`` valid Pod name of appropriate length
"""
MAX_POD_ID_LEN = 253
safe_key = safe_dag_id + safe_task_id
safe_pod_id = safe_key[:MAX_POD_ID_LEN - len(safe_uuid) - 1] + "-" + safe_uuid
return safe_pod_id
@staticmethod
def _create_pod_id(dag_id, task_id):
safe_dag_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
dag_id)
safe_task_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
task_id)
safe_uuid = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
uuid4().hex)
return AirflowKubernetesScheduler._make_safe_pod_id(safe_dag_id, safe_task_id,
safe_uuid)
@staticmethod
def _label_safe_datestring_to_datetime(string):
"""
Kubernetes doesn't permit ":" in labels. ISO datetime format uses ":" but not
"_", let's
replace ":" with "_"
:param string: str
:return: datetime.datetime object
"""
return parser.parse(string.replace('_plus_', '+').replace("_", ":"))
@staticmethod
def _datetime_to_label_safe_datestring(datetime_obj):
"""
Kubernetes doesn't like ":" in labels, since ISO datetime format uses ":" but
not "_" let's
replace ":" with "_"
:param datetime_obj: datetime.datetime object
:return: ISO-like string representing the datetime
"""
return datetime_obj.isoformat().replace(":", "_").replace('+', '_plus_')
def _labels_to_key(self, labels):
try:
return (
labels['dag_id'], labels['task_id'],
self._label_safe_datestring_to_datetime(labels['execution_date']),
labels['try_number'])
except Exception as e:
self.log.warn(
'Error while converting labels to key; labels: %s; exception: %s',
labels, e
)
return None
class AirflowKubernetesScheduler(LoggingMixin):
def __init__(self, kube_config, task_queue, result_queue, kube_client,
worker_uuid):
self.log.debug("Creating Kubernetes executor")
self.kube_config = kube_config
self.task_queue = task_queue
self.result_queue = result_queue
self.namespace = self.kube_config.kube_namespace
self.log.debug("Kubernetes using namespace %s", self.namespace)
self.kube_client = kube_client
self.launcher = PodLauncher(kube_client=self.kube_client)
self.worker_configuration = WorkerConfiguration(
kube_config=self.kube_config)
self.watcher_queue = multiprocessing.Queue()
self.worker_uuid = worker_uuid
self.kube_watcher = self._make_kube_watcher()
def _make_kube_watcher(self):
resource_version = KubeResourceVersion.get_current_resource_version()
watcher = KubernetesJobWatcher(self.namespace, self.watcher_queue,
resource_version, self.worker_uuid)
watcher.start()
return watcher
def _health_check_kube_watcher(self):
if self.kube_watcher.is_alive():
pass
else:
self.log.error('Error while health checking kube watcher process. '
'Process died for unknown reasons')
self.kube_watcher = self._make_kube_watcher()
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
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s",
self.kube_config.kube_image)
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,
kube_executor_config=kube_executor_config)
# the watcher will monitor pods, so we do not block.
self.launcher.run_pod_async(pod)
self.log.debug("Kubernetes Job created!")
def delete_pod(self, pod_id):
if self.kube_config.delete_worker_pods:
try:
self.kube_client.delete_namespaced_pod(
pod_id, self.namespace, body=client.V1DeleteOptions())
except ApiException as e:
# If the pod is already deleted
if e.status != 404:
raise
def sync(self):
"""
The sync function checks the status of all currently running kubernetes jobs.
If a job is completed, it's status is placed in the result queue to
be sent back to the scheduler.
:return:
"""
self._health_check_kube_watcher()
while not self.watcher_queue.empty():
self.process_watcher_task()
def process_watcher_task(self):
pod_id, state, labels, resource_version = self.watcher_queue.get()
self.log.info(
'Attempting to finish pod; pod_id: %s; state: %s; labels: %s',
pod_id, state, labels)
key = self._labels_to_key(labels=labels)
if key:
self.log.debug('finishing job %s - %s (%s)', key, state, pod_id)
self.result_queue.put((key, state, pod_id, resource_version))
@staticmethod
def _strip_unsafe_kubernetes_special_chars(string):
"""
Kubernetes only supports lowercase alphanumeric characters and "-" and "." in
the pod name
However, there are special rules about how "-" and "." can be used so let's
only keep
alphanumeric chars see here for detail:
https://kubernetes.io/docs/concepts/overview/working-with-objects/names/
:param string: The requested Pod name
:return: ``str`` Pod name stripped of any unsafe characters
"""
return ''.join(ch.lower() for ind, ch in enumerate(string)
if ch.isalnum())
@staticmethod
def _make_safe_pod_id(safe_dag_id, safe_task_id, safe_uuid):
r"""
Kubernetes pod names must be <= 253 chars and must pass the following regex for
validation
"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
:param safe_dag_id: a dag_id with only alphanumeric characters
:param safe_task_id: a task_id with only alphanumeric characters
:param random_uuid: a uuid
:return: ``str`` valid Pod name of appropriate length
"""
MAX_POD_ID_LEN = 253
safe_key = safe_dag_id + safe_task_id
safe_pod_id = safe_key[:MAX_POD_ID_LEN - len(safe_uuid) -
1] + "-" + safe_uuid
return safe_pod_id
@staticmethod
def _make_safe_label_value(string):
"""
Valid label values must be 63 characters or less and must be empty or begin and
end with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),
dots (.), and alphanumerics between.
If the label value is then greater than 63 chars once made safe, or differs in any
way from the original value sent to this function, then we need to truncate to
53chars, and append it with a unique hash.
"""
MAX_LABEL_LEN = 63
safe_label = re.sub(r'^[^a-z0-9A-Z]*|[^a-zA-Z0-9_\-\.]|[^a-z0-9A-Z]*$',
'', string)
if len(safe_label) > MAX_LABEL_LEN or string != safe_label:
safe_hash = hashlib.md5(string.encode()).hexdigest()[:9]
safe_label = safe_label[:MAX_LABEL_LEN - len(safe_hash) -
1] + "-" + safe_hash
return safe_label
@staticmethod
def _create_pod_id(dag_id, task_id):
safe_dag_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
dag_id)
safe_task_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
task_id)
safe_uuid = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
uuid4().hex)
return AirflowKubernetesScheduler._make_safe_pod_id(
safe_dag_id, safe_task_id, safe_uuid)
@staticmethod
def _label_safe_datestring_to_datetime(string):
"""
Kubernetes doesn't permit ":" in labels. ISO datetime format uses ":" but not
"_", let's
replace ":" with "_"
:param string: str
:return: datetime.datetime object
"""
return parser.parse(string.replace('_plus_', '+').replace("_", ":"))
@staticmethod
def _datetime_to_label_safe_datestring(datetime_obj):
"""
Kubernetes doesn't like ":" in labels, since ISO datetime format uses ":" but
not "_" let's
replace ":" with "_"
:param datetime_obj: datetime.datetime object
:return: ISO-like string representing the datetime
"""
return datetime_obj.isoformat().replace(":",
"_").replace('+', '_plus_')
def _labels_to_key(self, labels):
try_num = 1
try:
try_num = int(labels.get('try_number', '1'))
except ValueError:
self.log.warn("could not get try_number as an int: %s",
labels.get('try_number', '1'))
try:
dag_id = labels['dag_id']
task_id = labels['task_id']
ex_time = self._label_safe_datestring_to_datetime(
labels['execution_date'])
except Exception as e:
self.log.warn(
'Error while retrieving labels; labels: %s; exception: %s',
labels, e)
return None
with create_session() as session:
tasks = (session.query(TaskInstance).filter_by(
execution_date=ex_time).all())
self.log.info('Checking %s task instances.', len(tasks))
for task in tasks:
if (self._make_safe_label_value(task.dag_id) == dag_id and
self._make_safe_label_value(task.task_id) == task_id
and task.execution_date == ex_time):
self.log.info(
'Found matching task %s-%s (%s) with current state of %s',
task.dag_id, task.task_id, task.execution_date,
task.state)
dag_id = task.dag_id
task_id = task.task_id
return (dag_id, task_id, ex_time, try_num)
self.log.warn(
'Failed to find and match task details to a pod; labels: %s',
labels)
return None
