ferruzzi commented on code in PR #61153:
URL: https://github.com/apache/airflow/pull/61153#discussion_r2738489921
##########
airflow-core/src/airflow/executors/base_executor.py:
##########
@@ -227,10 +229,46 @@ def log_task_event(self, *, event: str, extra: str,
ti_key: TaskInstanceKey):
self._task_event_logs.append(Log(event=event, task_instance=ti_key,
extra=extra))
def queue_workload(self, workload: workloads.All, session: Session) ->
None:
- if not isinstance(workload, workloads.ExecuteTask):
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ elif isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ else:
raise ValueError(f"Un-handled workload kind
{type(workload).__name__!r} in {type(self).__name__}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+
+ def _get_workloads_to_schedule(
+ self, open_slots: int
+ ) -> list[tuple[TaskInstanceKey | str, workloads.All]]:
+ """
+ Select and return the next batch of workloads to schedule, respecting
priority policy.
+
+ Priority Policy: Callbacks are scheduled before tasks (callbacks
complete existing work).
+ Callbacks are processed in FIFO order. Tasks are sorted by
priority_weight (higher priority first).
+
+ :param open_slots: Number of available execution slots
+ """
+ workloads_to_schedule: list[tuple[TaskInstanceKey | str,
workloads.All]] = []
+
+ if self.queued_callbacks:
+ for key, workload in self.queued_callbacks.items():
+ if len(workloads_to_schedule) >= open_slots:
+ break
+ workloads_to_schedule.append((key, workload))
+
+ remaining_slots = open_slots - len(workloads_to_schedule)
+ if remaining_slots and self.queued_tasks:
Review Comment:
Totally fair, it felt a little more clear/explicit but maybe I just
overthinking it. I'll make the change.
--
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]
