Enforce pool timeouts for map() chunks instead of crashing the TimeoutHandler

CHANGES.txt

@@ -1,3 +1,11 @@
+4.4.0 - unreleased
+------------------
+- Enforce pool soft and hard timeouts for map() chunks, and only time out and
+  kill the workers that still run a chunk (#473). `TimeoutHandler.on_soft_timeout`
+  and `on_hard_timeout` take an optional `pid`, and `MapResult` takes optional
+  `soft_timeout` and `timeout` keywords.
+
+
 4.3.1 - 2026-10-05
 ------------------
 - Reject bool for Pool numeric args (processes, timeouts, limits) - #460

billiard/pool.py

@@ -721,24 +721,34 @@ def _process_by_pid(self, pid):
             if proc.pid == pid
         ), (None, None))
 
-    def on_soft_timeout(self, job):
+    def on_soft_timeout(self, job, pid=None):
+        if pid is None:
+            pid = job._worker_pid
         debug('soft time limit exceeded for %r', job)
-        process, _index = self._process_by_pid(job._worker_pid)
+        process, _index = self._process_by_pid(pid)
         if not process:
             return
 
         # Run timeout callback
         job.handle_timeout(soft=True)
 
         try:
-            _kill(job._worker_pid, SIG_SOFT_TIMEOUT)
+            _kill(pid, SIG_SOFT_TIMEOUT)
         except OSError as exc:
             if get_errno(exc) != errno.ESRCH:
                 raise
 
-    def on_hard_timeout(self, job):
+    def on_hard_timeout(self, job, pid=None):
         if job.ready():
+            if isinstance(job, MapResult) and pid is not None:
+                # An earlier chunk already failed this map(); the other
+                # chunks that ran out of time still hold their workers.
+                process, _index = self._process_by_pid(pid)
+                if process and not self._runs_other_job(job, pid):
+                    self._trywaitkill(process)
             return
+        if pid is None:
+            pid = job._worker_pid
         debug('hard time limit exceeded for %r', job)
         # Remove from cache and set return value to an exception
         try:
@@ -749,14 +759,23 @@ def on_hard_timeout(self, job):
             pass
 
         # Remove from _pool
-        process, _index = self._process_by_pid(job._worker_pid)
+        process, _index = self._process_by_pid(pid)
 
         # Run timeout callback
         job.handle_timeout(soft=False)
 
         if process:
             self._trywaitkill(process)
 
+    def _runs_other_job(self, job, pid):
+        # A failed map() has left the cache and no longer sees results, so
+        # a worker that finished its chunk meanwhile may be running another
+        # task by now, and that task must not be killed.
+        return any(
+            pid in other.worker_pids()
+            for other in list(self.cache.values()) if other is not job
+        )
+
     def _trywaitkill(self, worker):
         debug('timeout: sending TERM to %s', worker._name)
         try:
@@ -800,22 +819,40 @@ def _timed_out(start, timeout):
             cache = copy.copy(self.cache)
 
             # Remove dirty items not in cache anymore
+            # (map chunks are tracked as ``(job_id, chunk)``).
             if dirty:
-                dirty = set(k for k in dirty if k in cache)
+                dirty = set(
+                    k for k in dirty
+                    if (k[0] if isinstance(k, tuple) else k) in cache
+                )
 
             for i, job in cache.items():
-                ack_time = job._time_accepted
+                if isinstance(job, IMapIterator):
+                    # imap() results do not record when their tasks
+                    # were accepted, so there is nothing to time.
+                    continue
                 soft_timeout = job._soft_timeout
                 if soft_timeout is None:
                     soft_timeout = t_soft
                 hard_timeout = job._timeout
                 if hard_timeout is None:
                     hard_timeout = t_hard
-                if _timed_out(ack_time, hard_timeout):
-                    on_hard_timeout(job)
-                elif i not in dirty and _timed_out(ack_time, soft_timeout):
-                    on_soft_timeout(job)
-                    dirty.add(i)
+                if isinstance(job, MapResult):
+                    # Every chunk of a map() is a separate task with its
+                    # own worker and start time.
+                    tasks = [
+                        ((i, chunk), ack_time, pid)
+                        for chunk, ack_time, pid in job._running_chunks()
+                    ]
+                else:
+                    tasks = [(i, job._time_accepted, job._worker_pid)]
+                for key, ack_time, pid in tasks:
+                    if _timed_out(ack_time, hard_timeout):
+                        on_hard_timeout(job, pid)
+                    elif key not in dirty and _timed_out(ack_time,
+                                                         soft_timeout):
+                        on_soft_timeout(job, pid)
+                        dirty.add(key)
             yield
 
     def body(self):
@@ -1673,7 +1710,9 @@ def _map_async(self, func, iterable, mapper, chunksize=None,
 
         task_batches = Pool._get_tasks(func, iterable, chunksize)
         result = MapResult(self._cache, chunksize, len(iterable), callback,
-                           error_callback=error_callback)
+                           error_callback=error_callback,
+                           soft_timeout=self.soft_timeout,
+                           timeout=self.timeout)
         self._taskqueue.put((((TASK, (result._job, i, mapper, (x,), {}))
                               for i, x in enumerate(task_batches)), None))
         return result
@@ -1973,16 +2012,19 @@ def _ack(self, i, time_accepted, pid, synqW_fd):
 
 class MapResult(ApplyResult):
 
-    def __init__(self, cache, chunksize, length, callback, error_callback):
+    def __init__(self, cache, chunksize, length, callback, error_callback,
+                 soft_timeout=None, timeout=None):
         ApplyResult.__init__(
             self, cache, callback, error_callback=error_callback,
+            soft_timeout=soft_timeout, timeout=timeout,
         )
         self._success = True
         self._length = length
         self._value = [None] * length
         self._accepted = [False] * length
         self._worker_pid = [None] * length
         self._time_accepted = [None] * length
+        self._chunks_done = set()
         self._chunksize = chunksize
         if chunksize <= 0:
             self._number_left = 0
@@ -1995,6 +2037,7 @@ def _set(self, i, success_result):
         success, result = success_result
         if success:
             self._value[i * self._chunksize:(i + 1) * self._chunksize] = result
+            self._chunks_done.add(i)
             self._number_left -= 1
             if self._number_left == 0:
                 if self._callback:
@@ -2024,8 +2067,21 @@ def _ack(self, i, time_accepted, pid, *args):
     def accepted(self):
         return all(self._accepted)
 
+    def _running_chunks(self):
+        # Yield ``(chunk, time_accepted, pid)`` for every chunk that has
+        # been accepted by a worker but has not returned a result yet.
+        if self._chunksize <= 0:
+            return
+        for start in range(0, self._length, self._chunksize):
+            chunk = start // self._chunksize
+            time_accepted = self._time_accepted[start]
+            if time_accepted is not None and chunk not in self._chunks_done:
+                yield chunk, time_accepted, self._worker_pid[start]
+
     def worker_pids(self):
-        return [pid for pid in self._worker_pid if pid]
+        # Only chunks that are still running: a worker that has returned
+        # its chunk no longer holds anything for this map().
+        return [pid for _, _, pid in self._running_chunks() if pid]
 
 #
 # Class whose instances are returned by `Pool.imap()`

t/unit/test_pool.py

@@ -94,6 +94,87 @@ def test_timeout_handler_iterates_with_cache(self):
         # If I call to handle the timeouts I expect no exception
         next(timeout_handler.handle_timeouts())
 
+    def test_timeout_handler_times_map_chunks(self):
+        # A map() result keeps one accept time and pid per item, and
+        # imap() results keep none at all; neither may crash the handler.
+        cache = {}
+        result = billiard.pool.MapResult(cache, 2, 4, None, None, timeout=1)
+        billiard.pool.IMapIterator(cache)
+        long_ago = time.monotonic() - 10
+        result._ack(0, long_ago, 111)
+        result._ack(1, long_ago, 222)
+        result._set(0, (True, [0, 1]))
+        timeout_handler = billiard.pool.TimeoutHandler([], cache, None, None)
+        timed_out = []
+        timeout_handler.on_hard_timeout = (
+            lambda job, pid=None: timed_out.append((job, pid)))
+        next(timeout_handler.handle_timeouts())
+        # Only the chunk that is still running has timed out.
+        assert timed_out == [(result, 222)]
+
+    def test_map_worker_pids_only_running_chunks(self):
+        result = billiard.pool.MapResult({}, 1, 3, None, None)
+        now = time.monotonic()
+        result._ack(0, now, 111)
+        result._ack(1, now, 222)
+        result._set(0, (True, [0]))
+        assert result.worker_pids() == [222]
+
+    def test_map_hard_timeout_kills_every_timed_out_chunk(self):
+        cache = {}
+        result = billiard.pool.MapResult(cache, 1, 2, None, None, timeout=1)
+        long_ago = time.monotonic() - 10
+        result._ack(0, long_ago, 111)
+        result._ack(1, long_ago, 222)
+        procs = [Mock(pid=111), Mock(pid=222)]
+        timeout_handler = billiard.pool.TimeoutHandler(
+            procs, cache, None, None)
+        timeout_handler._trywaitkill = Mock()
+        next(timeout_handler.handle_timeouts())
+        # The first chunk fails the map, the second must still be killed.
+        assert result.ready()
+        killed = [c.args[0].pid
+                  for c in timeout_handler._trywaitkill.call_args_list]
+        assert killed == [111, 222]
+
+    def test_map_hard_timeout_spares_worker_that_moved_on(self):
+        cache = {}
+        result = billiard.pool.MapResult(cache, 1, 2, None, None, timeout=1)
+        long_ago = time.monotonic() - 10
+        result._ack(0, long_ago, 111)
+        result._ack(1, long_ago, 222)
+        procs = [Mock(pid=111), Mock(pid=222)]
+        timeout_handler = billiard.pool.TimeoutHandler(
+            procs, cache, None, None)
+        killed = []
+
+        def trywaitkill(worker):
+            killed.append(worker.pid)
+            # While the first worker is being killed, the second chunk
+            # finishes and its worker accepts the next task.
+            if worker.pid == 111:
+                other = billiard.pool.ApplyResult(cache, None)
+                other._ack(0, time.monotonic(), 222, None)
+
+        timeout_handler._trywaitkill = trywaitkill
+        next(timeout_handler.handle_timeouts())
+        assert result.ready()
+        assert killed == [111]
+
+    def test_map_hard_timeout(self):
+        pool = billiard.pool.Pool(2, timeout=1)
+        try:
+            # The pool timeout applies to each chunk, not the whole map.
+            assert pool.map(time.sleep, [0.4] * 8, chunksize=1) == [None] * 8
+            result = pool.map_async(time.sleep, [0, 30], chunksize=1)
+            with pytest.raises(Exception) as excinfo:
+                result.get(timeout=10)
+            exc = getattr(excinfo.value, 'exc', excinfo.value)
+            assert isinstance(exc, TimeLimitExceeded)
+        finally:
+            pool.terminate()
+            pool.join()
+
     def test_timeout_handler_trywaitkill_windows(self, monkeypatch):
         monkeypatch.setattr(sys, "platform", "win32")
         from unittest.mock import MagicMock