General cleanup of things around KubernetesExecutor (#29148)

jedcunningham · web-flow · commit bfbe2cb6fc39 · 2023-01-25T08:45:30.000-07:00
The biggest change here is using `State` or `TaskInstanceState` where
appropriate. Even though their strings are the same, it provides useful
context, so we should use the correct enum.

One could argue `State` isn't the right thing to use, but at least it's
consistently wrong now :)
diff --git a/airflow/executors/kubernetes_executor.py b/airflow/executors/kubernetes_executor.py
@@ -51,18 +51,18 @@
 from airflow.utils.event_scheduler import EventScheduler
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import provide_session
-from airflow.utils.state import State
+from airflow.utils.state import State, TaskInstanceState
 
 ALL_NAMESPACES = "ALL_NAMESPACES"
 POD_EXECUTOR_DONE_KEY = "airflow_executor_done"
 
 # TaskInstance key, command, configuration, pod_template_file
 KubernetesJobType = Tuple[TaskInstanceKey, CommandType, Any, Optional[str]]
 
-# key, state, pod_id, namespace, resource_version
+# key, pod state, pod_id, namespace, resource_version
 KubernetesResultsType = Tuple[TaskInstanceKey, Optional[str], str, str, str]
 
-# pod_id, namespace, state, annotations, resource_version
+# pod_id, namespace, pod state, annotations, resource_version
 KubernetesWatchType = Tuple[str, str, Optional[str], Dict[str, str], str]
 
 
@@ -498,8 +498,8 @@ def clear_not_launched_queued_tasks(self, session=None) -> None:
         """
         Clear tasks that were not yet launched, but were previously queued.
 
-        Tasks can end up in a "Queued" state through when a rescheduled/deferred
-        operator comes back up for execution (with the same try_number) before the
+        Tasks can end up in a "Queued" state 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).
 
         It's also possible when an executor abruptly shuts down (leaving a non-empty
@@ -514,7 +514,7 @@ def clear_not_launched_queued_tasks(self, session=None) -> None:
 
         self.log.debug("Clearing tasks that have not been launched")
         query = session.query(TaskInstance).filter(
-            TaskInstance.state == State.QUEUED, TaskInstance.queued_by_job_id == self.job_id
+            TaskInstance.state == TaskInstanceState.QUEUED, TaskInstance.queued_by_job_id == self.job_id
         )
         if self.kubernetes_queue:
             query = query.filter(TaskInstance.queue == self.kubernetes_queue)
@@ -566,7 +566,7 @@ def clear_not_launched_queued_tasks(self, session=None) -> None:
                 TaskInstance.task_id == ti.task_id,
                 TaskInstance.run_id == ti.run_id,
                 TaskInstance.map_index == ti.map_index,
-            ).update({TaskInstance.state: State.SCHEDULED})
+            ).update({TaskInstance.state: TaskInstanceState.SCHEDULED})
 
     def start(self) -> None:
         """Starts the executor."""
@@ -621,7 +621,7 @@ def execute_async(
             pod_template_file = executor_config.get("pod_template_file", None)
         else:
             pod_template_file = None
-        self.event_buffer[key] = (State.QUEUED, self.scheduler_job_id)
+        self.event_buffer[key] = (TaskInstanceState.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
@@ -691,7 +691,7 @@ def sync(self) -> None:
                     if e.status in (400, 422):
                         self.log.error("Pod creation failed with reason %r. Failing task", e.reason)
                         key, _, _, _ = task
-                        self.change_state(key, State.FAILED, e)
+                        self.change_state(key, TaskInstanceState.FAILED, e)
                     else:
                         self.log.warning(
                             "ApiException when attempting to run task, re-queueing. Reason: %r. Message: %s",
diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
@@ -614,7 +614,7 @@ def _process_executor_events(self, session: Session) -> int:
                 state,
                 ti_key.try_number,
             )
-            if state in (TaskInstanceState.FAILED, TaskInstanceState.SUCCESS, TaskInstanceState.QUEUED):
+            if state in (State.FAILED, State.SUCCESS, State.QUEUED):
                 tis_with_right_state.append(ti_key)
 
         # Return if no finished tasks
@@ -637,7 +637,7 @@ def _process_executor_events(self, session: Session) -> int:
             buffer_key = ti.key.with_try_number(try_number)
             state, info = event_buffer.pop(buffer_key)
 
-            if state == TaskInstanceState.QUEUED:
+            if state == State.QUEUED:
                 ti.external_executor_id = info
                 self.log.info("Setting external_id for %s to %s", ti, info)
                 continue
diff --git a/airflow/kubernetes/pod_generator.py b/airflow/kubernetes/pod_generator.py
@@ -110,7 +110,7 @@ class PodGenerator:
     Any configuration that is container specific gets applied to
     the first container in the list of containers.
 
-    :param pod: The fully specified pod. Mutually exclusive with `path_or_string`
+    :param pod: The fully specified pod. Mutually exclusive with `pod_template_file`
     :param pod_template_file: Path to YAML file. Mutually exclusive with `pod`
     :param extract_xcom: Whether to bring up a container for xcom
     """
