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 run_pod_async(self, pod, **kwargs): pod_mutation_hook(pod) req = self.kube_req_factory.create(pod) self.log.debug('Pod Creation Request: \n%s', json.dumps(req, indent=2)) try: resp = self._client.create_namespaced_pod(body=req, namespace=pod.namespace, **kwargs) self.log.debug('Pod Creation Response: %s', resp) except ApiException: self.log.exception('Exception when attempting to create Namespaced Pod.') raise return resp
def test_pod_mutation_hook(self): """ Tests that pods are mutated by the pod_mutation_hook function in airflow_local_settings. """ with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"): from airflow import settings settings.import_local_settings() # pylint: ignore pod = Pod(image="ubuntu", envs={}, cmds=['echo "1"']) settings.pod_mutation_hook(pod) assert pod.namespace == 'airflow-tests'
def test_pod_mutation_hook(self): """ Tests that pods are mutated by the pod_mutation_hook function in airflow_local_settings. """ with SettingsContext(SETTINGS_FILE_POD_MUTATION_HOOK, "airflow_local_settings"): from airflow import settings settings.import_local_settings() pod = MagicMock() settings.pod_mutation_hook(pod) assert pod.namespace == 'airflow-tests'
def run_pod_async(self, pod: V1Pod, **kwargs): pod_mutation_hook(pod) sanitized_pod = self._client.api_client.sanitize_for_serialization(pod) json_pod = json.dumps(sanitized_pod, indent=2) self.log.debug('Pod Creation Request: \n%s', json_pod) try: resp = self._client.create_namespaced_pod(body=sanitized_pod, namespace=pod.metadata.namespace, **kwargs) self.log.debug('Pod Creation Response: %s', resp) except Exception as e: self.log.exception('Exception when attempting ' 'to create Namespaced Pod: %s', json_pod) raise e return resp
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 _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 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