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 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: 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)) # We keep a temporary local record that we've handled this so we don't # try and remove it from the QUEUED state while we process it self.last_handled[key] = time.time()
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 execute_async(self, key, command, queue=None, executor_config=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) self.task_queue.put((key, command, kube_executor_config))
def test_from_obj(self): 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 = 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' }], }], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{ 'hostPath': { 'path': '/tmp/' }, 'name': 'example-kubernetes-test-volume' }], } }, result)
def execute_async(self, key: TaskInstanceKeyType, 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) assert self.task_queue, NOT_STARTED_MESSAGE self.task_queue.put((key, command, kube_executor_config))
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 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))
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_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 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 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)