def test_execution_date_serialize_deserialize(self):
datetime_obj = datetime.now()
serialized_datetime = \
AirflowKubernetesScheduler._datetime_to_label_safe_datestring(
datetime_obj)
new_datetime_obj = AirflowKubernetesScheduler._label_safe_datestring_to_datetime(
serialized_datetime)
self.assertEqual(datetime_obj, new_datetime_obj)
def test_make_safe_label_value(self):
for dag_id, task_id in self._cases():
safe_dag_id = AirflowKubernetesScheduler._make_safe_label_value(
dag_id)
self.assertTrue(self._is_safe_label_value(safe_dag_id))
safe_task_id = AirflowKubernetesScheduler._make_safe_label_value(
task_id)
self.assertTrue(self._is_safe_label_value(safe_task_id))
id = "my_dag_id"
self.assertEqual(
id, AirflowKubernetesScheduler._make_safe_label_value(id))
id = "my_dag_id_" + "a" * 64
self.assertEqual(
"my_dag_id_" + "a" * 43 + "-0ce114c45",
AirflowKubernetesScheduler._make_safe_label_value(id))
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 test_create_pod_id(self):
for dag_id, task_id in self._cases():
pod_name = AirflowKubernetesScheduler._create_pod_id(
dag_id, task_id)
self.assertTrue(self._is_valid_pod_id(pod_name))
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))
