potiuk commented on code in PR #75086:
URL: https://github.com/apache/airflow/pull/75086#discussion_r4435490114
##########
airflow-core/src/airflow/executors/local_executor.py:
##########
@@ -337,17 +338,74 @@ def _terminate_worker_process(self, proc:
multiprocessing.Process) -> None:
def _process_workloads(self, workload_list):
for workload in workload_list:
key = self.get_workload_key(workload)
- self._read_results()
- self.activity_queue.put(workload)
removed = self.executor_queues[workload.type].pop(key, None)
if not removed:
raise KeyError(f"Workload {key} was not found in any queue")
+ # Register before the put: results read while the put is in flight
are dropped for keys
+ # that are not running yet.
self.running.add(key)
self._dispatch_counts[key] = self._dispatch_counts.get(key, 0) + 1
- with self._unread_messages:
- self._unread_messages.value += len(workload_list)
+ # Counted before the put so a worker taking the message never
decrements below zero, and so
+ # _check_workers sees it as outstanding while the put is blocked.
+ with self._unread_messages:
+ self._unread_messages.value += 1
+ self._read_results()
+ if (put_error := self._put_workload(workload)) is not None:
+ self._undo_dispatch(key, workload)
+ raise put_error
self._check_workers()
+ def _undo_dispatch(self, key: WorkloadKey, workload: ExecutorWorkload) ->
None:
+ """Reverse the bookkeeping of a workload whose put failed, leaving it
queued as it was before."""
+ self.executor_queues[workload.type][key] = workload
+ remaining = self._dispatch_counts.pop(key, 1) - 1
+ if remaining > 0:
+ self._dispatch_counts[key] = remaining
Review Comment:
Thanks — applied as suggested in 1d250da0c6: the test is now parametrized
with `no-earlier-dispatch` and `earlier-dispatch-still-running`, so the
`remaining > 0` branch of `_undo_dispatch()` is covered. (Branch also rebased
onto current `main`.)
---
Drafted-by: Claude Code (Claude Opus 5); reviewed by @potiuk before posting
--
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]
