kaxil commented on code in PR #63491:
URL: https://github.com/apache/airflow/pull/63491#discussion_r3982155693
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -210,6 +216,46 @@ def jwt_generator(self) -> JWTGenerator:
return generator
+ def __init_subclass__(cls, **kwargs: Any) -> None:
+ super().__init_subclass__(**kwargs)
+ cls._legacy_warned = set()
+ legacy_workload_types: set[WorkloadType] = set()
+ if cls.__dict__.get("supports_callbacks") is True:
+ _warn_deprecated_executor_usage(
+ f"{cls.__name__}: setting `supports_callbacks = True` as a
class attribute is "
+ f"deprecated. Add `WorkloadType.EXECUTE_CALLBACK` to
`supported_workload_types` "
+ f"instead.",
+ )
+ legacy_workload_types.add(WorkloadType.EXECUTE_CALLBACK)
+ if cls.__dict__.get("supports_connection_test") is True:
+ _warn_deprecated_executor_usage(
+ f"{cls.__name__}: setting `supports_connection_test = True` as
a class attribute is "
+ f"deprecated. Add `WorkloadType.TEST_CONNECTION` to
`supported_workload_types` "
+ f"instead.",
+ )
+ legacy_workload_types.add(WorkloadType.TEST_CONNECTION)
+ if legacy_workload_types and "supported_workload_types" not in
cls.__dict__:
+ cls.supported_workload_types = cls.supported_workload_types |
legacy_workload_types
+ legacy_override_replacements = {
Review Comment:
`trigger_connection_tests` is the third method this PR removes from a base
class that shipped it in 3.3.0 and 3.3.1, and it is the only one of the three
missing from this dict, so an out-of-tree override goes dead with no warning
while `trigger_tasks` and `order_queued_tasks_by_priority` both get one. Adding
`"trigger_connection_tests": "trigger_workloads"` covers it in one line.
##########
airflow-core/src/airflow/executors/workloads/base.py:
##########
@@ -32,6 +33,38 @@
from airflow.executors.workloads.types import WorkloadState
+class WorkloadType(str, Enum):
+ """Central registry of executor workload types."""
+
+ EXECUTE_TASK = "ExecuteTask"
+ EXECUTE_CALLBACK = "ExecuteCallback"
+ TEST_CONNECTION = "TestConnection"
+
+
+# Central executor priority registry: tuple is ordered from highest priority
to lowest.
+#
+# Connection tests are short-lived and user-interactive, so they sort ahead of
tasks:
+# otherwise a sustained task backlog would starve them until the reaper times
them out.
+#
+# Adding a new workload type is a four-place change that must stay in sync:
Review Comment:
Two more sites raise on an unrecognised type and are not in this list.
`BaseExecutor.run_workload` is an isinstance chain over exactly these three
schemas ending in `raise ValueError(f"Unknown workload type: ...")`, and
`state_class_for_key` in `workloads/types.py` ends in `raise TypeError`.
`run_workload` is the worker-side entrypoint for local, Celery and edge3, so a
fourth type added by following this checklist would queue, sort and dispatch
cleanly and then fail on every worker.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -263,65 +370,44 @@ def log_task_event(self, *, event: str, extra: str,
ti_key: WorkloadKey):
self._task_event_logs.append(Log(event=event, task_instance=ti_key,
extra=extra))
def queue_workload(self, workload: ExecutorWorkload, session: Session) ->
None:
- if isinstance(workload, workloads.ExecuteTask):
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- elif isinstance(workload, workloads.ExecuteCallback):
- if not self.supports_callbacks:
- raise NotImplementedError(
- f"{type(self).__name__} does not support ExecuteCallback
workloads. "
- f"Set supports_callbacks = True and implement callback
handling in _process_workloads(). "
- f"See LocalExecutor or CeleryExecutor for reference
implementation."
- )
- self.queued_callbacks[workload.key] = workload
- elif isinstance(workload, workloads.TestConnection):
- if not self.supports_connection_test:
- raise NotImplementedError(
- f"{type(self).__name__} does not support TestConnection
workloads. "
- f"Set supports_connection_test = True and implement
connection test handling "
- f"in _process_workloads(). See LocalExecutor for reference
implementation."
- )
- self.queued_connection_tests[workload.key] = workload
- else:
- raise ValueError(
- f"Un-handled workload type {type(workload).__name__!r} in
{type(self).__name__}. "
- f"Workload must be one of: ExecuteTask, ExecuteCallback,
TestConnection."
+ if workload.type not in self.supported_workload_types:
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support {workload.type!r}
workloads. "
Review Comment:
`workload.type` is a `WorkloadType` member, so `!r` renders
`<WorkloadType.TEST_CONNECTION: 'TestConnection'>` and the message ends up
telling the reader to add exactly that to `supported_workload_types`, which is
not something they can type. `workload.type.value` gives the name they need.
The tests having to loosen `match="does not support TestConnection workloads"`
to `match="does not support.*TestConnection"` is this leaking through.
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -252,6 +303,62 @@ def __repr__(self):
_repr += ")"
return _repr
+ @property
+ def queued_tasks(self) -> dict:
+ """Backward-compat property: delegates to
``executor_queues[WorkloadType.EXECUTE_TASK]``."""
+ self._warn_legacy_property(
+ "queued_tasks",
+ "queued_tasks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_TASK] instead.",
+ )
+ return self.executor_queues[WorkloadType.EXECUTE_TASK]
+
+ @queued_tasks.setter
+ def queued_tasks(self, value: dict) -> None:
+ """Backward-compat setter: writes through to
``executor_queues[WorkloadType.EXECUTE_TASK]``."""
+ self._warn_legacy_property(
+ "queued_tasks",
+ "queued_tasks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_TASK] instead.",
+ )
+ self.executor_queues[WorkloadType.EXECUTE_TASK] = value
+
+ @property
+ def queued_callbacks(self) -> dict:
+ """Backward-compat property: delegates to
``executor_queues[WorkloadType.EXECUTE_CALLBACK]``."""
+ self._warn_legacy_property(
+ "queued_callbacks",
+ "queued_callbacks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_CALLBACK] instead.",
+ )
+ return self.executor_queues[WorkloadType.EXECUTE_CALLBACK]
+
+ @queued_callbacks.setter
+ def queued_callbacks(self, value: dict) -> None:
+ """Backward-compat setter: writes through to
``executor_queues[WorkloadType.EXECUTE_CALLBACK]``."""
+ self._warn_legacy_property(
+ "queued_callbacks",
+ "queued_callbacks is deprecated. Use
executor_queues[WorkloadType.EXECUTE_CALLBACK] instead.",
+ )
+ self.executor_queues[WorkloadType.EXECUTE_CALLBACK] = value
+
+ @property
+ def supports_callbacks(self) -> bool:
+ """Backward-compat property: True if EXECUTE_CALLBACK is in
supported_workload_types."""
+ self._warn_legacy_property(
+ "supports_callbacks",
+ "supports_callbacks is deprecated. "
+ "Use WorkloadType.EXECUTE_CALLBACK in supported_workload_types
instead.",
+ )
+ return WorkloadType.EXECUTE_CALLBACK in self.supported_workload_types
+
+ @property
+ def supports_connection_test(self) -> bool:
Review Comment:
`supports_callbacks` and `supports_connection_test` were plain `bool` class
attributes in released 3.3.0 and 3.3.1, so `self.supports_callbacks = True`
inside an executor's `__init__` was legal and is now an `AttributeError` at
instantiation, with `__init_subclass__` only covering the class-body form. This
PR added write-through deprecating setters for `queued_tasks` and
`queued_callbacks` in exactly the same situation, and the newsfragment's "still
honored" is unqualified, so the same setters here would keep that promise true.
The two in-tree sites this already broke are `test_scheduler_job.py:13412` and
`:13646`, both instance assignments the PR had to rewrite.
##########
airflow-core/newsfragments/63491.significant.rst:
##########
@@ -0,0 +1,52 @@
+Unify executor workload queues under ``BaseExecutor.executor_queues``
Review Comment:
The newsfragment is accurate now, but the PR description has drifted from
it. The description still opens with "No behavioral change , scheduling order,
slot accounting ... identically to before", which this file's own "Behavioural
changes" section contradicts; it documents a `queue_key` property that no
longer exists anywhere in the tree; and it says four provider executors dropped
their `queue_workload` overrides when only `KubernetesExecutor`'s was removed,
since ECS, Batch and Lambda kept theirs behind a version gate. Worth refreshing
before merge, since the description is what the next reviewer reads first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
