Repository: joblib/joblib

Issue:
Parallel(): timeout does not trigger with return_as="generator_unordered"
Jobs don't timeout when you specify `return_as="generator_unordered"`, i.e. when you want to receive partial results instead of having all jobs fail. I'm on Windows 10, Python 3.9.19, joblib 1.4.2, installed via conda.

```py
import multiprocessing
import time
from joblib import Parallel, delayed


def inf_loop():
    while True:
        time.sleep(1)


try:
    ret = Parallel(n_jobs=2, timeout=3)(delayed(inf_loop)() for _ in range(3))
except multiprocessing.context.TimeoutError as e:
    print(f"Caught {e}")
    print("Continuing")

try:
    for _ in Parallel(n_jobs=2, timeout=3, return_as="generator_unordered")(delayed(inf_loop)() for _ in range(3)):
        print("One done")
except multiprocessing.context.TimeoutError as e:
    print(f"Caught {e}")
    print("Continuing")
```
This code prints the first "Caught" and "Continuing", but not the second and never terminates.k

Discussion:
@lesteve: I can reproduce on Linux, Python 3.11, joblib 1.4.2, this seems like a bug indeed that `return_as='generator_unordered'` does not work with `timeout`.

The snippet works fine with `return_as='generator'` and `timeout`.@fcharras: Indeed, the timeout logic has issues. I'm preparing a PR.

Solution patch:
diff --git a/README.rst b/README.rst
index 661bae4..824b8ef 100644
--- a/README.rst
+++ b/README.rst
@@ -34,15 +34,13 @@ of the latest code: https://github.com/joblib/joblib/archive/refs/heads/main.zip
 Installing
 ==========
 
-You can use `pip` to install joblib::
+You can use `pip` to install joblib from any directory::
 
     pip install joblib
 
-from any directory or::
+or install it in editable mode from the source directory::
 
-    python setup.py install
-
-from the source directory.
+    pip install -e .
 
 Dependencies
 ============
@@ -69,10 +67,9 @@ Workflow to contribute
 To contribute to joblib, first create an account on `github
 <https://github.com/>`_. Once this is done, fork the `joblib repository
 <https://github.com/joblib/joblib>`_ to have your own repository,
-clone it using 'git clone' on the computers where you want to work. Make
-your changes in your clone, push them to your github account, test them
-on several computers, and when you are happy with them, send a pull
-request to the main repository.
+clone it using ``git clone``. Make your changes in a branch of your clone, push
+them to your github account, test them locally, and when you are happy with
+them, send a pull request to the main repository.
 
 Running the test suite
 ======================
diff --git a/joblib/parallel.py b/joblib/parallel.py
index 17be0a8..04f7f52 100644
--- a/joblib/parallel.py
+++ b/joblib/parallel.py
@@ -712,6 +712,7 @@ class BatchCompletionCallBack(object):
         self.batch_size = batch_size
         self.parallel = parallel
         self.parallel_call_id = parallel._call_id
+        self._completion_timeout_counter = None
 
         # Internals to keep track of the status and outcome of the task.
 
@@ -793,7 +794,7 @@ class BatchCompletionCallBack(object):
         # The computation are running and the status is pending.
         # Check that we did not wait for this jobs more than `timeout`.
         now = time.time()
-        if not hasattr(self, "_completion_timeout_counter"):
+        if self._completion_timeout_counter is None:
             self._completion_timeout_counter = now
 
         if (now - self._completion_timeout_counter) > timeout:
@@ -835,11 +836,6 @@ class BatchCompletionCallBack(object):
             # a new batch if needed.
             job_succeeded = self._retrieve_result(out)
 
-            if not self.parallel.return_ordered:
-                # Append the job to the queue in the order of completion
-                # instead of submission.
-                self.parallel._jobs.append(self)
-
         if job_succeeded:
             self._dispatch_new()
 
@@ -905,6 +901,14 @@ class BatchCompletionCallBack(object):
             self.parallel._exception = True
             self.parallel._aborting = True
 
+        if self.parallel.return_ordered:
+            return
+
+        with self.parallel._lock:
+            # For `return_as=generator_unordered`, append the job to the queue
+            # in the order of completion instead of submission.
+            self.parallel._jobs.append(self)
+
 
 ###############################################################################
 def register_parallel_backend(name, factory, make_default=False):
@@ -1346,6 +1350,7 @@ class Parallel(Logger):
             # with the async callback thread of our the pool.
             self._lock = threading.RLock()
             self._jobs = collections.deque()
+            self._jobs_set = set()
             self._pending_outputs = list()
             self._ready_batches = queue.Queue()
             self._reducer_callback = None
@@ -1423,8 +1428,7 @@ class Parallel(Logger):
             dispatch_timestamp, batch_size, self
         )
 
-        if self.return_ordered:
-            self._jobs.append(batch_tracker)
+        self._register_new_job(batch_tracker)
 
         # If return_ordered is False, the batch_tracker is not stored in the
         # jobs queue at the time of submission. Instead, it will be appended to
@@ -1434,6 +1438,12 @@ class Parallel(Logger):
         job = self._backend.apply_async(batch, callback=batch_tracker)
         batch_tracker.register_job(job)
 
+    def _register_new_job(self, batch_tracker):
+        if self.return_ordered:
+            self._jobs.append(batch_tracker)
+        else:
+            self._jobs_set.add(batch_tracker)
+
     def dispatch_next(self):
         """Dispatch more data for parallel processing
 
@@ -1497,7 +1507,7 @@ class Parallel(Logger):
                     batch_tracker = BatchCompletionCallBack(
                         0, batch_size, self
                     )
-                    self._jobs.append(batch_tracker)
+                    self._register_new_job(batch_tracker)
                     batch_tracker._register_outcome(dict(
                         result=e, status=TASK_ERROR
                     ))
@@ -1734,6 +1744,7 @@ class Parallel(Logger):
             # Store the unconsumed tasks and terminate the workers if necessary
             _remaining_outputs = ([] if self._exception else self._jobs)
             self._jobs = collections.deque()
+            self._jobs_set = set()
             self._running = False
             if not detach_generator_exit:
                 self._terminate_and_reset()
@@ -1772,6 +1783,7 @@ class Parallel(Logger):
         return False
 
     def _retrieve(self):
+        timeout_control_job = None
         while self._wait_retrieval():
 
             # If the callback thread of a worker has signaled that its task
@@ -1782,19 +1794,58 @@ class Parallel(Logger):
                 self._raise_error_fast()
                 break
 
-            # If the next job is not ready for retrieval yet, we just wait for
-            # async callbacks to progress.
-            if ((len(self._jobs) == 0) or
-                (self._jobs[0].get_status(
-                    timeout=self.timeout) == TASK_PENDING)):
+            nb_jobs = len(self._jobs)
+            # Now wait for a job to be ready for retrieval.
+            if self.return_ordered:
+                # Case ordered: wait for completion (or error) of the next job
+                # that have been dispatched and not retrieved yet. If no job
+                # have been dispatched yet, wait for dispatch.
+                # We assume that the time to wait for the next job to be
+                # dispatched is always low, so that the timeout
+                # control only have to be done on the amount of time the next
+                # dispatched job is pending.
+                if ((nb_jobs == 0) or
+                    (self._jobs[0].get_status(
+                        timeout=self.timeout) == TASK_PENDING)):
+                    time.sleep(0.01)
+                    continue
+
+            elif (nb_jobs == 0):
+                # Case unordered: jobs are added to the list of jobs to
+                # retrieve `self._jobs` only once completed or in error, which
+                # is too late to enable timeout control in the same way than in
+                # the previous case.
+                # Instead, if no job is ready to be retrieved yet, we
+                # arbitrarily pick a dispatched job, and the timeout control is
+                # done such that an error is raised if this control job
+                # timeouts before any other dispatched job has completed and
+                # been added to `self._jobs` to be retrieved.
+                if timeout_control_job is None:
+                    timeout_control_job = next(iter(self._jobs_set), None)
+
+                # NB: it can be None if no job has been dispatched yet.
+                if timeout_control_job is not None:
+                    timeout_control_job.get_status(timeout=self.timeout)
+
                 time.sleep(0.01)
                 continue
 
+            elif timeout_control_job is not None:
+                # Case unordered, when `nb_jobs > 0`:
+                # It means that a job is ready to be retrieved, so no timeout
+                # will occur during this iteration.
+                # Before proceeding to retrieval of the next ready job, reset
+                # the timeout control state to prepare the next iteration.
+                timeout_control_job._completion_timeout_counter = None
+                timeout_control_job = None
+
             # We need to be careful: the job list can be filling up as
             # we empty it and Python list are not thread-safe by
             # default hence the use of the lock
             with self._lock:
                 batched_results = self._jobs.popleft()
+                if not self.return_ordered:
+                    self._jobs_set.remove(batched_results)
 
             # Flatten the batched results to output one output at a time
             batched_results = batched_results.get_result(self.timeout)


Test patch:
diff --git a/joblib/test/test_parallel.py b/joblib/test/test_parallel.py
index d48aa3c..43538a5 100644
--- a/joblib/test/test_parallel.py
+++ b/joblib/test/test_parallel.py
@@ -444,6 +444,33 @@ def test_parallel_timeout_fail(backend):
             delayed(sleep)(10) for x in range(10))
 
 
+@with_multiprocessing
+@parametrize('backend', set(RETURN_GENERATOR_BACKENDS) - {"sequential"})
+@parametrize('return_as', ["generator", "generator_unordered"])
+def test_parallel_timeout_fail_with_generator(backend, return_as):
+    # Check that timeout properly fails when function is too slow with
+    # return_as=generator
+    with raises(TimeoutError):
+        list(
+            Parallel(
+                n_jobs=2,
+                backend=backend,
+                return_as=return_as,
+                timeout=0.1
+            )(delayed(sleep)(10) for x in range(10))
+        )
+
+    # Fast tasks and high timeout should not raise
+    list(
+        Parallel(
+            n_jobs=2,
+            backend=backend,
+            return_as=return_as,
+            timeout=10
+        )(delayed(sleep)(0.01) for x in range(10))
+    )
+
+
 @with_multiprocessing
 @parametrize('backend', PROCESS_BACKENDS)
 def test_error_capture(backend):