FrankYang0529 commented on code in PR #75086:
URL: https://github.com/apache/airflow/pull/75086#discussion_r4345272155
##########
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:
The test case `test_process_workloads_undoes_dispatch_when_put_fails` only
covers a key with no earlier dispatch, so the test never reaches the `remaining
> 0` branch of `_undo_dispatch()`. How about changing the test case like
following?
```py
@pytest.mark.parametrize(
("earlier_dispatches", "still_running"),
[
pytest.param(0, False, id="no-earlier-dispatch"),
pytest.param(1, True, id="earlier-dispatch-still-running"),
],
)
def test_process_workloads_undoes_dispatch_when_put_fails(self,
earlier_dispatches, still_running):
"""A failed put undoes only its own dispatch, so an earlier dispatch
of the same key stays running."""
executor = LocalExecutor(parallelism=1)
executor.workers = {}
executor.activity_queue = mock.MagicMock(spec=SimpleQueueType)
executor.activity_queue.put.side_effect = ValueError("unpicklable")
executor._read_results = mock.MagicMock(spec=executor._read_results)
executor._unread_messages = multiprocessing.Value("I", 0)
executor._check_workers =
mock.MagicMock(spec=executor._check_workers)
workload = _make_task_workload()
key = executor.get_workload_key(workload)
executor.executor_queues.setdefault(workload.type, {})[key] =
workload
if earlier_dispatches:
# A resumed attempt reuses the key while the earlier dispatch
has not reported its result yet.
executor.running.add(key)
executor._dispatch_counts[key] = earlier_dispatches
with pytest.raises(ValueError, match="unpicklable"):
executor._process_workloads([workload])
assert (key in executor.running) is still_running
assert executor._dispatch_counts.get(key, 0) == earlier_dispatches
assert executor.executor_queues[workload.type][key] is workload
assert executor._unread_messages.value == 0
```
--
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]
