Repository: joblib/joblib
Pull request #1393 (2023-02-17 10:55:08)

FEA return generator, #588 stripped of unrelated changes for minimal diff review
The main goal of this PR is to make it possible to collect intermediate results without having to wait for all results to be computed first.

See the description of the original PR for the context:

- #588

codecov[bot]: # [Codecov](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
Base: **93.89**% // Head: **69.27**% // Decreases project coverage by **`-24.62%`** :warning:
> Coverage data is based on head [(`daa42e9`)](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) compared to base [(`e4c00de`)](https://codecov.io/gh/joblib/joblib/commit/e4c00defb5b694ba3db9c89f1e3572c50f4ff372?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).
> Patch coverage: 92.68% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master    #1393       +/-   ##
===========================================
- Coverage   93.89%   69.27%   -24.62%     
===========================================
  Files          52       52               
  Lines        7341     7418       +77     
===========================================
- Hits         6893     5139     -1754     
- Misses        448     2279     +1831     
```


| [Impacted Files](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | Coverage Δ | |
|---|---|---|
| [joblib/\_dask.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19kYXNrLnB5) | `24.01% <33.33%> (-68.67%)` | :arrow_down: |
| [joblib/test/common.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvY29tbW9uLnB5) | `69.84% <40.00%> (-7.75%)` | :arrow_down: |
| [joblib/\_memmapping\_reducer.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19tZW1tYXBwaW5nX3JlZHVjZXIucHk=) | `52.91% <85.00%> (-42.97%)` | :arrow_down: |
| [joblib/parallel.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3BhcmFsbGVsLnB5) | `93.04% <92.82%> (-3.00%)` | :arrow_down: |
| [joblib/\_parallel\_backends.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19wYXJhbGxlbF9iYWNrZW5kcy5weQ==) | `87.14% <96.07%> (-4.41%)` | :arrow_down: |
| [joblib/test/test\_parallel.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9wYXJhbGxlbC5weQ==) | `83.17% <96.70%> (-13.88%)` | :arrow_down: |
| [joblib/pool.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Bvb2wucHk=) | `74.40% <100.00%> (-13.41%)` | :arrow_down: |
| [joblib/test/test\_dask.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9kYXNrLnB5) | `3.92% <0.00%> (-87.27%)` | :arrow_down: |
| [joblib/backports.py](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL2JhY2twb3J0cy5weQ==) | `26.53% <0.00%> (-42.86%)` | :arrow_down: |
| ... and [13 more](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/joblib/joblib/pull/1393?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).

fcharras: It looks better now. But the thread bomb mitigation unit test often hangs. Even if it has issues on master it is more unstable here I think ? Starting repeated pipeline runs.
fcharras: Lot of red in the last pipelines but it's only the thread bomb unit test hanging. I can't reproduce. Not sure what to do with it.
fcharras: The thread bomb mitigation unit test [hangs the same on master](https://dev.azure.com/joblib/joblib/_build/results?buildId=2744&view=logs&jobId=8c31ea29-785c-52fc-16ae-5bb169e02d59&j=8c31ea29-785c-52fc-16ae-5bb169e02d59&t=f76d557a-4262-56a5-1178-119687e76288)

I think we can consider merging with the current state.
tomMoral: I ran the proposed benchmark and I don't see any change between the two branches.

![image](https://user-images.githubusercontent.com/3321081/232820157-3bbe69c9-a860-406f-b699-d6d7fd993cd9.png)


I also did a benchmark with small scale gridsearch on synthetic data w/ `sklearn` and looked at the scaling with the number of CPUs. Here are the results (I will put the benchmark_script in a separate PR).
![image](https://user-images.githubusercontent.com/3321081/232832941-5f97207b-93f9-4da8-bddc-f10d002f5744.png)

It seems safe to say that there is not noticeable performance drop, so I propose merging this one :)
ogrisel: The benchmark script is probably not working as intended anymore because the final batch size is always 1 but at least there is not noticeable regression in the PR so it's good.

Patch:
diff --git a/.readthedocs-requirements.txt b/.readthedocs-requirements.txt
index 05dd31d..07fa6e4 100644
--- a/.readthedocs-requirements.txt
+++ b/.readthedocs-requirements.txt
@@ -8,3 +8,4 @@ numpydoc
 pandas
 lz4
 distributed
+psutil
diff --git a/CHANGES.rst b/CHANGES.rst
index 3f1bf8e..90f2b61 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -32,6 +32,14 @@ In development
   previous versions of Joblib.
   https://github.com/joblib/joblib/pull/1374
 
+- Add a ``return_generator`` parameter for ``Parallel``, that allows
+  to consume results asynchronously.
+  https://github.com/joblib/joblib/pull/1393
+
+- Improve the behavior of ``joblib`` for ``n_jobs=1``, with simplified
+  tracebacks and more efficient running time.
+  https://github.com/joblib/joblib/pull/1393
+
 Release 1.2.0
 -------------
 
diff --git a/doc/parallel.rst b/doc/parallel.rst
index 41f7aab..b4c5071 100644
--- a/doc/parallel.rst
+++ b/doc/parallel.rst
@@ -23,6 +23,26 @@ can be spread over 2 CPUs using the following::
     >>> Parallel(n_jobs=2)(delayed(sqrt)(i ** 2) for i in range(10))
     [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
 
+The output can be a generator that yields the results as soon as they're
+available, even if the subsequent tasks aren't completed yet. The order
+of the outputs always matches the order of the inputs::
+
+    >>> from math import sqrt
+    >>> from joblib import Parallel, delayed
+    >>> parallel = Parallel(n_jobs=2, return_generator=True)
+    >>> output_generator = parallel(delayed(sqrt)(i ** 2) for i in range(10))
+    >>> print(type(output_generator))
+    <class 'generator'>
+    >>> print(next(output_generator))
+    0.0
+    >>> print(next(output_generator))
+    1.0
+    >>> print(list(output_generator))
+    [2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
+
+This generator allows to reduce the memory footprint of :class:`joblib.Parallel`
+calls in case the results can benefit from on-the-fly aggregation, as illustrated
+in :ref:`sphx_glr_auto_examples_parallel_generator.py`.
 
 Thread-based parallelism vs process-based parallelism
 =====================================================
diff --git a/examples/parallel_generator.py b/examples/parallel_generator.py
new file mode 100644
index 0000000..0e4dbc3
--- /dev/null
+++ b/examples/parallel_generator.py
@@ -0,0 +1,172 @@
+"""
+========================================
+Returning a generator in joblib.Parallel
+========================================
+
+This example illustrates memory optimization enabled by using
+:class:`joblib.Parallel` to get a generator on the outputs of parallel jobs.
+We first create tasks that return results with large memory footprints.
+If we call :class:`~joblib.Parallel` for several of these tasks directly, we
+observe a high memory usage, as all the results are held in RAM before being
+processed
+
+Using the ``return_generator=True`` option allows to progressively consumes
+the outputs as they arrive and keeps the memory at an acceptable level.
+"""
+
+##############################################################################
+# ``MemoryMonitor`` helper
+##############################################################################
+
+##############################################################################
+# The following class is an helper to monitor the memory of the process and its
+# children in another thread, so we can display it afterward.
+#
+# We will use ``psutil`` to monitor the memory usage in the code. Make sure it
+# is installed with ``pip install psutil`` for this example.
+
+
+import time
+from psutil import Process
+from threading import Thread
+
+
+class MemoryMonitor(Thread):
+    """Monitor the memory usage in MB in a separate thread.
+
+    Note that this class is good enough to highlight the memory profile of
+    Parallel in this example, but is not a general purpose profiler fit for
+    all cases.
+    """
+    def __init__(self):
+        super().__init__()
+        self.stop = False
+        self.memory_buffer = []
+        self.start()
+
+    def get_memory(self):
+        "Get memory of a process and its children."
+        p = Process()
+        memory = p.memory_info().rss
+        for c in p.children():
+            memory += c.memory_info().rss
+        return memory
+
+    def run(self):
+        memory_start = self.get_memory()
+        while not self.stop:
+            self.memory_buffer.append(self.get_memory() - memory_start)
+            time.sleep(0.2)
+
+    def join(self):
+        self.stop = True
+        super().join()
+
+
+##############################################################################
+# Save memory by consuming the outputs of the tasks as fast as possible
+##############################################################################
+
+##############################################################################
+# We create a task whose output takes about 15MB of RAM.
+#
+
+import numpy as np
+
+
+def return_big_object(i):
+    time.sleep(.1)
+    return i * np.ones((10000, 200), dtype=np.float64)
+
+
+##############################################################################
+# We create a reduce step. The input will be a generator on big objects
+# generated in parallel by several instances of :func:`return_big_object`.
+
+def accumulator_sum(generator):
+    result = 0
+    for value in generator:
+        result += value
+        print(".", end="", flush=True)
+    print("")
+    return result
+
+
+##############################################################################
+# We process many of the tasks in parallel. If `return_generator=False`
+# (default), we should expect a usage of more than 2GB in RAM. Indeed, all the
+# results are computed and stored in ``res`` before being processed by
+# `accumulator_sum` and collected by the gc.
+
+from joblib import Parallel, delayed
+
+monitor = MemoryMonitor()
+print("Running tasks with return_generator=False...")
+res = Parallel(n_jobs=2, return_generator=False)(
+    delayed(return_big_object)(i) for i in range(150)
+)
+print("Accumulate results:", end='')
+res = accumulator_sum(res)
+print('All tasks completed and reduced successfully.')
+
+# Report memory usage
+del res  # we clean the result to avoid memory border effects
+monitor.join()
+peak = max(monitor.memory_buffer) / 1e9
+print(f"Peak memory usage: {peak:.2f}GB")
+
+
+##############################################################################
+# If we use ``return_generator=True``, ``res`` is simply a generator with the
+# results that are ready. Here we consume the results as soon as they arrive
+# with the ``accumulator_sum`` and once they have been used, they are collected
+# by the gc. The memory footprint is thus reduced, typically around 300MB.
+
+monitor_gen = MemoryMonitor()
+print("Create result generator with return_generator=True...")
+res = Parallel(n_jobs=2, return_generator=True)(
+    delayed(return_big_object)(i) for i in range(150)
+)
+print("Accumulate results:", end='')
+res = accumulator_sum(res)
+print('All tasks completed and reduced successfully.')
+
+# Report memory usage
+del res  # we clean the result to avoid memory border effects
+monitor_gen.join()
+peak = max(monitor_gen.memory_buffer) / 1e6
+print(f"Peak memory usage: {peak:.2f}MB")
+
+
+##############################################################################
+# We can then report the memory usage accross time of the two runs using the
+# MemoryMonitor.
+#
+# In the first case, as the results accumulate in ``res``, the memory grows
+# linearly and it is freed once the ``accumulator_sum`` function finishes.
+#
+# In the second case, the results are processed by the accumulator as soon as
+# they arrive, and the memory does not need to be able to contain all
+# the results.
+
+import matplotlib.pyplot as plt
+plt.semilogy(
+    np.maximum.accumulate(monitor.memory_buffer),
+    label='return_generator=False'
+)
+plt.semilogy(
+    np.maximum.accumulate(monitor_gen.memory_buffer),
+    label='return_generator=True'
+)
+plt.xlabel("Time")
+plt.xticks([], [])
+plt.ylabel("Memory usage")
+plt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])
+plt.legend()
+plt.show()
+
+##############################################################################
+# It is important to note that with ``return_generator``, the results are
+# still accumulated in RAM after computation. But as we asynchronously process
+# them, they can be freed sooner. However, if the generator is not consumed
+# the memory still grows linearly.
diff --git a/joblib/_dask.py b/joblib/_dask.py
index 4b797c8..c96c7c7 100644
--- a/joblib/_dask.py
+++ b/joblib/_dask.py
@@ -8,8 +8,8 @@ import time
 from uuid import uuid4
 import weakref
 
-from .parallel import AutoBatchingMixin, ParallelBackendBase
 from .parallel import parallel_backend
+from .parallel import AutoBatchingMixin, ParallelBackendBase
 
 try:
     import dask
@@ -124,7 +124,7 @@ class Batch:
         with parallel_backend('dask'):
             for func, args, kwargs in tasks:
                 results.append(func(*args, **kwargs))
-        return results
+            return results
 
     def __repr__(self):
         descr = f"batch_of_{self._funcname}_{self._num_tasks}_calls"
diff --git a/joblib/_parallel_backends.py b/joblib/_parallel_backends.py
index 45ab833..53a7157 100644
--- a/joblib/_parallel_backends.py
+++ b/joblib/_parallel_backends.py
@@ -6,11 +6,10 @@ import gc
 import os
 import warnings
 import threading
-import functools
 import contextlib
 from abc import ABCMeta, abstractmethod
 
-from .my_exceptions import WorkerInterrupt
+
 from ._multiprocessing_helpers import mp
 
 if mp is not None:
@@ -18,17 +17,26 @@ if mp is not None:
     from multiprocessing.pool import ThreadPool
     from .executor import get_memmapping_executor
 
-    # Compat between concurrent.futures and multiprocessing TimeoutError
-    from multiprocessing import TimeoutError
-    from concurrent.futures._base import TimeoutError as CfTimeoutError
+    # Import loky only if multiprocessing is present
     from .externals.loky import process_executor, cpu_count
+    from .externals.loky.process_executor import ShutdownExecutorError
+    from .externals.loky.process_executor import _ExceptionWithTraceback
 
 
 class ParallelBackendBase(metaclass=ABCMeta):
     """Helper abc which defines all methods a ParallelBackend must implement"""
 
-    supports_timeout = False
     supports_inner_max_num_threads = False
+    supports_retrieve_callback = False
+
+    @property
+    def supports_return_generator(self):
+        return self.supports_retrieve_callback
+
+    @property
+    def supports_timeout(self):
+        return self.supports_retrieve_callback
+
     nesting_level = None
 
     def __init__(self, nesting_level=None, inner_max_num_threads=None,
@@ -67,6 +75,14 @@ class ParallelBackendBase(metaclass=ABCMeta):
     def apply_async(self, func, callback=None):
         """Schedule a func to be run"""
 
+    def retrieve_result_callback(self, out):
+        """Called within the callback function passed in apply_async.
+
+        The argument of this function is the argument given to a callback in
+        the considered backend. It is supposed to return the outcome of a task
+        if it succeeded or raise the exception if it failed.
+        """
+
     def configure(self, n_jobs=1, parallel=None, prefer=None, require=None,
                   **backend_args):
         """Reconfigure the backend and return the number of workers.
@@ -195,6 +211,8 @@ class SequentialBackend(ParallelBackendBase):
     """
 
     uses_threads = True
+    supports_timeout = False
+    supports_retrieve_callback = False
     supports_sharedmem = True
 
     def effective_n_jobs(self, n_jobs):
@@ -205,10 +223,10 @@ class SequentialBackend(ParallelBackendBase):
 
     def apply_async(self, func, callback=None):
         """Schedule a func to be run"""
-        result = ImmediateResult(func)
-        if callback:
-            callback(result)
-        return result
+        raise RuntimeError("Should never be called for SequentialBackend.")
+
+    def retrieve_result_callback(self, out):
+        raise RuntimeError("Should never be called for SequentialBackend.")
 
     def get_nested_backend(self):
         # import is not top level to avoid cyclic import errors.
@@ -247,10 +265,32 @@ class PoolManagerMixin(object):
         """Used by apply_async to make it possible to implement lazy init"""
         return self._pool
 
+    @staticmethod
+    def _wrap_func_call(func):
+        """Protect function call and return error with traceback."""
+        try:
+            return func()
+        except BaseException as e:
+            return _ExceptionWithTraceback(e)
+
     def apply_async(self, func, callback=None):
         """Schedule a func to be run"""
+        # Here, we need a wrapper to avoid crashes on KeyboardInterruptErrors.
+        # We also call the callback on error, to make sure the pool does not
+        # wait on crashed jobs.
         return self._get_pool().apply_async(
-            SafeFunction(func), callback=callback)
+            self._wrap_func_call, (func,),
+            callback=callback, error_callback=callback
+        )
+
+    def retrieve_result_callback(self, out):
+        """Mimic concurrent.futures results, raising an error if needed."""
+        if isinstance(out, _ExceptionWithTraceback):
+            rebuild, args = out.__reduce__()
+            out = rebuild(*args)
+        if isinstance(out, BaseException):
+            raise out
+        return out
 
     def abort_everything(self, ensure_ready=True):
         """Shutdown the pool and restart a new one with the same parameters"""
@@ -306,8 +346,9 @@ class AutoBatchingMixin(object):
             self._effective_batch_size = batch_size
             if self.parallel.verbose >= 10:
                 self.parallel._print(
-                    "Batch computation too fast (%.4fs.) "
-                    "Setting batch_size=%d.", (batch_duration, batch_size))
+                    f"Batch computation too fast ({batch_duration}s.) "
+                    f"Setting batch_size={batch_size}."
+                )
         elif (batch_duration > self.MAX_IDEAL_BATCH_DURATION and
               old_batch_size >= 2):
             # The current batch size is too big. If we schedule overly long
@@ -325,8 +366,9 @@ class AutoBatchingMixin(object):
             self._effective_batch_size = batch_size
             if self.parallel.verbose >= 10:
                 self.parallel._print(
-                    "Batch computation too slow (%.4fs.) "
-                    "Setting batch_size=%d.", (batch_duration, batch_size))
+                    f"Batch computation too slow ({batch_duration}s.) "
+                    f"Setting batch_size={batch_size}."
+                )
         else:
             # No batch size adjustment
             batch_size = old_batch_size
@@ -382,7 +424,7 @@ class ThreadingBackend(PoolManagerMixin, ParallelBackendBase):
     ThreadingBackend is used as the default backend for nested calls.
     """
 
-    supports_timeout = True
+    supports_retrieve_callback = True
     uses_threads = True
     supports_sharedmem = True
 
@@ -417,7 +459,8 @@ class MultiprocessingBackend(PoolManagerMixin, AutoBatchingMixin,
     However, does not suffer from the Python Global Interpreter Lock.
     """
 
-    supports_timeout = True
+    supports_retrieve_callback = True
+    supports_return_generator = False
 
     def effective_n_jobs(self, n_jobs):
         """Determine the number of jobs which are going to run in parallel.
@@ -492,7 +535,7 @@ class MultiprocessingBackend(PoolManagerMixin, AutoBatchingMixin,
 class LokyBackend(AutoBatchingMixin, ParallelBackendBase):
     """Managing pool of workers with loky instead of multiprocessing."""
 
-    supports_timeout = True
+    supports_retrieve_callback = True
     supports_inner_max_num_threads = True
 
     def configure(self, n_jobs=1, parallel=None, prefer=None, require=None,
@@ -553,20 +596,22 @@ class LokyBackend(AutoBatchingMixin, ParallelBackendBase):
 
     def apply_async(self, func, callback=None):
         """Schedule a func to be run"""
-        future = self._workers.submit(SafeFunction(func))
-        future.get = functools.partial(self.wrap_future_result, future)
+        future = self._workers.submit(func)
         if callback is not None:
             future.add_done_callback(callback)
         return future
 
-    @staticmethod
-    def wrap_future_result(future, timeout=None):
-        """Wrapper for Future.result to implement the same behaviour as
-        AsyncResults.get from multiprocessing."""
+    def retrieve_result_callback(self, out):
         try:
-            return future.result(timeout=timeout)
-        except CfTimeoutError as e:
-            raise TimeoutError from e
+            return out.result()
+        except ShutdownExecutorError:
+            raise RuntimeError(
+                "The executor underlying Parallel has been shutdown. "
+                "This is likely due to the garbage collection of a previous "
+                "generator from a call to Parallel with return_generator=True."
+                " Make sure the generator is not garbage collected when "
+                "submitting a new job or that it is first properly exhausted."
+            )
 
     def terminate(self):
         if self._workers is not None:
@@ -590,44 +635,6 @@ class LokyBackend(AutoBatchingMixin, ParallelBackendBase):
             self.configure(n_jobs=self.parallel.n_jobs, parallel=self.parallel)
 
 
-class ImmediateResult(object):
-    def __init__(self, batch):
-        # Don't delay the application, to avoid keeping the input
-        # arguments in memory
-        self.results = batch()
-
-    def get(self):
-        return self.results
-
-
-class SafeFunction(object):
-    """Wrapper that handles the serialization of exception tracebacks.
-
-    TODO python2_drop: check whether SafeFunction is still needed since we
-    dropped support for Python 2. If not needed anymore it should be
-    deprecated.
-
-    If an exception is triggered when calling the inner function, a copy of
-    the full traceback is captured to make it possible to serialize
-    it so that it can be rendered in a different Python process.
-
-    """
-    def __init__(self, func):
-        self.func = func
-
-    def __call__(self, *args, **kwargs):
-        try:
-            return self.func(*args, **kwargs)
-        except KeyboardInterrupt as e:
-            # We capture the KeyboardInterrupt and reraise it as
-            # something different, as multiprocessing does not
-            # interrupt processing for a KeyboardInterrupt
-            raise WorkerInterrupt() from e
-        except BaseException:
-            # Rely on Python 3 built-in Remote Traceback reporting
-            raise
-
-
 class FallbackToBackend(Exception):
     """Raised when configuration should fallback to another backend"""
 
diff --git a/joblib/my_exceptions.py b/joblib/my_exceptions.py
deleted file mode 100644
index 3b25209..0000000
--- a/joblib/my_exceptions.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""
-Exceptions
-"""
-# Author: Gael Varoquaux < gael dot varoquaux at normalesup dot org >
-# Copyright: 2010, Gael Varoquaux
-# License: BSD 3 clause
-
-
-class WorkerInterrupt(Exception):
-    """ An exception that is not KeyboardInterrupt to allow subprocesses
-        to be interrupted.
-    """
-
-    pass
diff --git a/joblib/parallel.py b/joblib/parallel.py
index 725f0e7..4608885 100644
--- a/joblib/parallel.py
+++ b/joblib/parallel.py
@@ -11,6 +11,7 @@ import os
 import sys
 from math import sqrt
 import functools
+import collections
 import time
 import threading
 import itertools
@@ -18,6 +19,10 @@ from uuid import uuid4
 from numbers import Integral
 import warnings
 import queue
+import weakref
+from contextlib import nullcontext
+
+from multiprocessing import TimeoutError
 
 from ._multiprocessing_helpers import mp
 
@@ -34,6 +39,9 @@ from ._parallel_backends import AutoBatchingMixin  # noqa
 from ._parallel_backends import ParallelBackendBase  # noqa
 
 
+IS_PYPY = hasattr(sys, "pypy_version_info")
+
+
 BACKENDS = {
     'threading': ThreadingBackend,
     'sequential': SequentialBackend,
@@ -59,6 +67,7 @@ if mp is not None:
 
 DEFAULT_THREAD_BACKEND = 'threading'
 
+
 # Thread local value that can be overridden by the ``parallel_backend`` context
 # manager
 _backend = threading.local()
@@ -301,7 +310,7 @@ class BatchedCalls(object):
     def __reduce__(self):
         if self._reducer_callback is not None:
             self._reducer_callback()
-        # no need pickle the callback.
+        # no need to pickle the callback.
         return (
             BatchedCalls,
             (self.items, (self._backend, self._n_jobs), None,
@@ -312,6 +321,12 @@ class BatchedCalls(object):
         return self._size
 
 
+# Possible exit status for a task
+TASK_DONE = "Done"
+TASK_ERROR = "Error"
+TASK_PENDING = "Pending"
+
+
 ###############################################################################
 # CPU count that works also when multiprocessing has been disabled via
 # the JOBLIB_MULTIPROCESSING environment variable
@@ -368,32 +383,216 @@ def delayed(function):
 
 ###############################################################################
 class BatchCompletionCallBack(object):
-    """Callback used by joblib.Parallel's multiprocessing backend.
+    """Callback to keep track of completed results and schedule the next tasks.
 
     This callable is executed by the parent process whenever a worker process
-    has returned the results of a batch of tasks.
+    has completed a batch of tasks.
 
     It is used for progress reporting, to update estimate of the batch
     processing duration and to schedule the next batch of tasks to be
     processed.
 
+    It is assumed that this callback will always be triggered by the backend
+    right after the end of a task, in case of success as well as in case of
+    failure.
     """
+
+    ##########################################################################
+    #                   METHODS CALLED BY THE MAIN THREAD                    #
+    ##########################################################################
     def __init__(self, dispatch_timestamp, batch_size, parallel):
         self.dispatch_timestamp = dispatch_timestamp
         self.batch_size = batch_size
         self.parallel = parallel
+        self.parallel_call_id = parallel._call_id
+
+        # Internals to keep track of the status and outcome of the task.
+
+        # Used to hold a reference to the future-like object returned by the
+        # backend after launching this task
+        # This will be set later when calling `register_job`, as it is only
+        # created once the task has been submitted.
+        self.job = None
+
+        if not parallel._backend.supports_retrieve_callback:
+            # The status is only used for asynchronous result retrieval in the
+            # callback.
+            self.status = None
+        else:
+            # The initial status for the job is TASK_PENDING.
+            # Once it is done, it will be either TASK_DONE, or TASK_ERROR.
+            self.status = TASK_PENDING
+
+    def register_job(self, job):
+        """Register the object returned by `apply_async`."""
+        self.job = job
+
+    def get_result(self, timeout):
+        """Returns the raw result of the task that was submitted.
+
+        If the task raised an exception rather than returning, this same
+        exception will be raised instead.
+
+        If the backend supports the retrieval callback, it is assumed that this
+        method is only called after the result has been registered. It is
+        ensured by checking that `self.status(timeout)` does not return
+        TASK_PENDING. In this case, `get_result` directly returns the
+        registered result (or raise the registered exception).
+
+        For other backends, there are no such assumptions, but `get_result`
+        still needs to synchronously retrieve the result before it can
+        return it or raise. It will block at most `self.timeout` seconds
+        waiting for retrieval to complete, after that it raises a TimeoutError.
+        """
+
+        backend = self.parallel._backend
+
+        if backend.supports_retrieve_callback:
+            # We assume that the result has already been retrieved by the
+            # callback thread, and is stored internally. It's just waiting to
+            # be returned.
+            return self._return_or_raise()
+
+        # For other backends, the main thread needs to run the retrieval step.
+        try:
+            if backend.supports_timeout:
+                result = self.job.get(timeout=timeout)
+            else:
+                result = self.job.get()
+            outcome = dict(result=result, status=TASK_DONE)
+        except BaseException as e:
+            outcome = dict(result=e, status=TASK_ERROR)
+        self._register_outcome(outcome)
 
+        return self._return_or_raise()
+
+    def _return_or_raise(self):
+        try:
+            if self.status == TASK_ERROR:
+                raise self._result
+            return self._result
+        finally:
+            del self._result
+
+    def get_status(self, timeout):
+        """Get the status of the task.
+
+        This function also checks if the timeout has been reached and register
+        the TimeoutError outcome when it is the case.
+        """
+        if timeout is None or self.status != TASK_PENDING:
+            return self.status
+
+        # 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"):
+            self._completion_timeout_counter = now
+
+        if (now - self._completion_timeout_counter) > timeout:
+            outcome = dict(result=TimeoutError(), status=TASK_ERROR)
+            self._register_outcome(outcome)
+
+        return self.status
+
+    ##########################################################################
+    #                     METHODS CALLED BY CALLBACK THREADS                 #
+    ##########################################################################
     def __call__(self, out):
-        self.parallel.n_completed_tasks += self.batch_size
-        this_batch_duration = time.time() - self.dispatch_timestamp
+        """Function called by the callback thread after a job is completed."""
+
+        # If the backend doesn't support callback retrievals, the next batch of
+        # tasks is dispatched regardless. The result will be retrieved by the
+        # main thread when calling `get_result`.
+        if not self.parallel._backend.supports_retrieve_callback:
+            self._dispatch_new()
+            return
+
+        # If the backend supports retrieving the result in the callback, it
+        # registers the task outcome (TASK_ERROR or TASK_DONE), and schedules
+        # the next batch if needed.
+        with self.parallel._lock:
+            # Edge case where while the task was processing, the `parallel`
+            # instance has been reset and a new call has been issued, but the
+            # worker managed to complete the task and trigger this callback
+            # call just before being aborted by the reset.
+            if self.parallel._call_id != self.parallel_call_id:
+                return
+
+            # When aborting, stop as fast as possible and do not retrieve the
+            # result as it won't be returned by the Parallel call.
+            if self.parallel._aborting:
+                return
+
+            # Retrieves the result of the task in the main process and dispatch
+            # a new batch if needed.
+            job_succeeded = self._retrieve_result(out)
+
+        if job_succeeded:
+            self._dispatch_new()
 
+    def _dispatch_new(self):
+        """Schedule the next batch of tasks to be processed."""
+
+        # This steps ensure that auto-baching works as expected.
+        this_batch_duration = time.time() - self.dispatch_timestamp
         self.parallel._backend.batch_completed(self.batch_size,
                                                this_batch_duration)
-        self.parallel.print_progress()
+
+        # Schedule the next batch of tasks.
         with self.parallel._lock:
+            self.parallel.n_completed_tasks += self.batch_size
+            self.parallel.print_progress()
             if self.parallel._original_iterator is not None:
                 self.parallel.dispatch_next()
 
+    def _retrieve_result(self, out):
+        """Fetch and register the outcome of a task.
+
+        Return True if the task succeeded, False otherwise.
+        This function is only called by backends that support retrieving
+        the task result in the callback thread.
+        """
+        try:
+            result = self.parallel._backend.retrieve_result_callback(out)
+            outcome = dict(status=TASK_DONE, result=result)
+        except BaseException as e:
+            # Avoid keeping references to parallel in the error.
+            e.__traceback__ = None
+            outcome = dict(result=e, status=TASK_ERROR)
+
+        self._register_outcome(outcome)
+        return outcome['status'] != TASK_ERROR
+
+    ##########################################################################
+    #            This method can be called either in the main thread         #
+    #                        or in the callback thread.                      #
+    ##########################################################################
+    def _register_outcome(self, outcome):
+        """Register the outcome of a task.
+
+        This method can be called only once, future calls will be ignored.
+        """
+        # Covers the edge case where the main thread tries to register a
+        # `TimeoutError` while the callback thread tries to register a result
+        # at the same time.
+        with self.parallel._lock:
+            if self.status not in (TASK_PENDING, None):
+                return
+            self.status = outcome["status"]
+
+        self._result = outcome["result"]
+
+        # Once the result and the status are extracted, the last reference to
+        # the job can be deleted.
+        self.job = None
+
+        # As soon as an error as been spotted, early stopping flags are sent to
+        # the `parallel` instance.
+        if self.status == TASK_ERROR:
+            self.parallel._exception = True
+            self.parallel._aborting = True
+
 
 ###############################################################################
 def register_parallel_backend(name, factory, make_default=False):
@@ -440,6 +639,9 @@ def effective_n_jobs(n_jobs=-1):
     .. versionadded:: 0.10
 
     """
+    if n_jobs == 1:
+        return 1
+
     backend, backend_n_jobs = get_active_backend()
     if n_jobs is None:
         n_jobs = backend_n_jobs
@@ -458,14 +660,15 @@ class Parallel(Logger):
             The maximum number of concurrently running jobs, such as the number
             of Python worker processes when backend="multiprocessing"
             or the size of the thread-pool when backend="threading".
-            If -1 all CPUs are used. If 1 is given, no parallel computing code
-            is used at all, which is useful for debugging. For n_jobs below -1,
-            (n_cpus + 1 + n_jobs) are used. Thus for n_jobs = -2, all
-            CPUs but one are used.
+            If -1 all CPUs are used.
+            If 1 is given, no parallel computing code is used at all, and the
+            behavior amounts to a simple python `for` loop. This mode is not
+            compatible with `timeout`.
+            For n_jobs below -1, (n_cpus + 1 + n_jobs) are used. Thus for
+            n_jobs = -2, all CPUs but one are used.
             None is a marker for 'unset' that will be interpreted as n_jobs=1
-            (sequential execution) unless the call is performed under a
-            :func:`~parallel_backend` context manager that sets another value
-            for n_jobs.
+            unless the call is performed under a :func:`~parallel_backend`
+            context manager that sets another value for ``n_jobs``.
         backend: str, ParallelBackendBase instance or None, default: 'loky'
             Specify the parallelization backend implementation.
             Supported backends are:
@@ -493,6 +696,11 @@ class Parallel(Logger):
             soft hints (prefer) or hard constraints (require) so as to make it
             possible for library users to change the backend from the outside
             using the :func:`~parallel_backend` context manager.
+        return_generator: bool
+            If True, calls to this instance will return a generator, yielding
+            the results as soon as they are available, in the original order.
+            Note that the intended usage is to run one call at a time. Multiple
+            calls to the same Parallel object will result in a ``RuntimeError``
         prefer: str in {'processes', 'threads'} or None, default: None
             Soft hint to choose the default backend if no specific backend
             was selected with the :func:`~parallel_backend` context manager.
@@ -677,11 +885,11 @@ class Parallel(Logger):
         [Parallel(n_jobs=2)]: Done 6 out of 6 | elapsed:  0.0s remaining: 0.0s
         [Parallel(n_jobs=2)]: Done 6 out of 6 | elapsed:  0.0s finished
 
-    '''
-    def __init__(self, n_jobs=None, backend=None, verbose=0, timeout=None,
-                 pre_dispatch='2 * n_jobs', batch_size='auto',
-                 temp_folder=None, max_nbytes='1M', mmap_mode='r',
-                 prefer=None, require=None):
+    '''  # noqa: E501
+    def __init__(self, n_jobs=None, backend=None, return_generator=False,
+                 verbose=0, timeout=None, pre_dispatch='2 * n_jobs',
+                 batch_size='auto', temp_folder=None, max_nbytes='1M',
+                 mmap_mode='r', prefer=None, require=None):
         active_backend, context_n_jobs = get_active_backend(
             prefer=prefer, require=require, verbose=verbose)
         nesting_level = active_backend.nesting_level
@@ -697,9 +905,7 @@ class Parallel(Logger):
         self.verbose = verbose
         self.timeout = timeout
         self.pre_dispatch = pre_dispatch
-        self._ready_batches = queue.Queue()
-        self._id = uuid4().hex
-        self._reducer_callback = None
+        self.return_generator = return_generator
 
         if isinstance(max_nbytes, str):
             max_nbytes = memstr_to_bytes(max_nbytes)
@@ -763,23 +969,38 @@ class Parallel(Logger):
                 "batch_size must be 'auto' or a positive integer, got: %r"
                 % batch_size)
 
+        if not isinstance(backend, SequentialBackend):
+            if return_generator and not backend.supports_return_generator:
+                raise ValueError(
+                    "Backend {} does not support "
+                    "return_generator=True".format(backend)
+                )
+            # This lock is used to coordinate the main thread of this process
+            # with the async callback thread of our the pool.
+            self._lock = threading.RLock()
+            self._jobs = collections.deque()
+            self._pending_outputs = list()
+            self._ready_batches = queue.Queue()
+            self._reducer_callback = None
+
+        # Internal variables
         self._backend = backend
-        self._output = None
-        self._jobs = list()
+        self._running = False
         self._managed_backend = False
-
-        # This lock is used coordinate the main thread of this process with
-        # the async callback thread of our the pool.
-        self._lock = threading.RLock()
+        self._id = uuid4().hex
+        self._call_ref = None
 
     def __enter__(self):
         self._managed_backend = True
+        self._calling = False
         self._initialize_backend()
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self._terminate_backend()
         self._managed_backend = False
+        if self.return_generator and self._calling:
+            self._abort()
+        self._terminate_and_reset()
 
     def _initialize_backend(self):
         """Build a process or thread pool and return the number of workers"""
@@ -806,8 +1027,11 @@ class Parallel(Logger):
             return self._backend.effective_n_jobs(self.n_jobs)
         return 1
 
-    def _terminate_backend(self):
-        if self._backend is not None:
+    def _terminate_and_reset(self):
+        if hasattr(self._backend, 'stop_call') and self._calling:
+            self._backend.stop_call()
+        self._calling = False
+        if not self._managed_backend:
             self._backend.terminate()
 
     def _dispatch(self, batch):
@@ -821,19 +1045,20 @@ class Parallel(Logger):
         if self._aborting:
             return
 
-        self.n_dispatched_tasks += len(batch)
+        batch_size = len(batch)
+
+        self.n_dispatched_tasks += batch_size
         self.n_dispatched_batches += 1
 
         dispatch_timestamp = time.time()
-        cb = BatchCompletionCallBack(dispatch_timestamp, len(batch), self)
-        with self._lock:
-            job_idx = len(self._jobs)
-            job = self._backend.apply_async(batch, callback=cb)
-            # A job can complete so quickly than its callback is
-            # called before we get here, causing self._jobs to
-            # grow. To ensure correct results ordering, .insert is
-            # used (rather than .append) in the following line
-            self._jobs.insert(job_idx, job)
+
+        batch_tracker = BatchCompletionCallBack(
+            dispatch_timestamp, batch_size, self
+        )
+        self._jobs.append(batch_tracker)
+
+        job = self._backend.apply_async(batch, callback=batch_tracker)
+        batch_tracker.register_job(job)
 
     def dispatch_next(self):
         """Dispatch more data for parallel processing
@@ -857,11 +1082,11 @@ class Parallel(Logger):
         lock so calling this function should be thread safe.
 
         """
-        if self.batch_size == 'auto':
-            batch_size = self._backend.compute_batch_size()
-        else:
-            # Fixed batch size strategy
-            batch_size = self.batch_size
+
+        if self._aborting:
+            return False
+
+        batch_size = self._get_batch_size()
 
         with self._lock:
             # to ensure an even distribution of the workolad between workers,
@@ -885,8 +1110,8 @@ class Parallel(Logger):
                 islice = list(itertools.islice(iterator, big_batch_size))
                 if len(islice) == 0:
                     return False
-                elif (iterator is self._original_iterator
-                      and len(islice) < big_batch_size):
+                elif (iterator is self._original_iterator and
+                      len(islice) < big_batch_size):
                     # We reached the end of the original iterator (unless
                     # iterator is the ``pre_dispatch``-long initial slice of
                     # the original iterator) -- decrease the batch size to
@@ -913,7 +1138,15 @@ class Parallel(Logger):
                 self._dispatch(tasks)
                 return True
 
-    def _print(self, msg, msg_args):
+    def _get_batch_size(self):
+        """Returns the effective batch size for dispatch"""
+        if self.batch_size == 'auto':
+            return self._backend.compute_batch_size()
+        else:
+            # Fixed batch size strategy
+            return self.batch_size
+
+    def _print(self, msg):
         """Display the message on stout or stderr depending on verbosity"""
         # XXX: Not using the logger framework: need to
         # learn to use logger better.
@@ -923,28 +1156,45 @@ class Parallel(Logger):
             writer = sys.stderr.write
         else:
             writer = sys.stdout.write
-        msg = msg % msg_args
-        writer('[%s]: %s\n' % (self, msg))
+        writer(f"[{self}]: {msg}\n")
+
+    def _is_completed(self):
+        """Check if all tasks have been completed"""
+        return self.n_completed_tasks == self.n_dispatched_tasks and not (
+            self._iterating or self._aborting
+        )
 
     def print_progress(self):
         """Display the process of the parallel execution only a fraction
            of time, controlled by self.verbose.
         """
+
         if not self.verbose:
             return
+
         elapsed_time = time.time() - self._start_time
 
+        if self._is_completed():
+            # Make sure that we get a last message telling us we are done
+            self._print(
+                f"Done {self.n_completed_tasks:3d} out of "
+                f"{self.n_completed_tasks:3d} | elapsed: "
+                f"{short_format_time(elapsed_time)} finished"
+            )
+            return
+
         # Original job iterator becomes None once it has been fully
         # consumed : at this point we know the total number of jobs and we are
         # able to display an estimation of the remaining time based on already
         # completed jobs. Otherwise, we simply display the number of completed
         # tasks.
-        if self._original_iterator is not None:
+        elif self._original_iterator is not None:
             if _verbosity_filter(self.n_dispatched_batches, self.verbose):
                 return
-            self._print('Done %3i tasks      | elapsed: %s',
-                        (self.n_completed_tasks,
-                         short_format_time(elapsed_time), ))
+            self._print(
+                f"Done {self.n_completed_tasks:3d} tasks      | elapsed: "
+                f"{short_format_time(elapsed_time)}"
+            )
         else:
             index = self.n_completed_tasks
             # We are finished dispatching
@@ -962,65 +1212,334 @@ class Parallel(Logger):
             remaining_time = (elapsed_time / index) * \
                              (self.n_dispatched_tasks - index * 1.0)
             # only display status if remaining time is greater or equal to 0
-            self._print('Done %3i out of %3i | elapsed: %s remaining: %s',
-                        (index,
-                         total_tasks,
-                         short_format_time(elapsed_time),
-                         short_format_time(remaining_time),
-                         ))
-
-    def retrieve(self):
-        self._output = list()
-        while self._iterating or len(self._jobs) > 0:
-            if len(self._jobs) == 0:
-                # Wait for an async callback to dispatch new jobs
+            self._print(
+                f"Done {index:3d} out of {total_tasks:3d} | elapsed: "
+                f"{short_format_time(elapsed_time)} remaining: "
+                f"{short_format_time(remaining_time)}"
+            )
+
+    def _abort(self):
+        # Stop dispatching new jobs in the async callback thread
+        self._aborting = True
+
+        # If the backend allows it, cancel or kill remaining running
+        # tasks without waiting for the results as we will raise
+        # the exception we got back to the caller instead of returning
+        # any result.
+        backend = self._backend
+        if (not self._aborted and hasattr(backend, 'abort_everything')):
+            # If the backend is managed externally we need to make sure
+            # to leave it in a working state to allow for future jobs
+            # scheduling.
+            ensure_ready = self._managed_backend
+            backend.abort_everything(ensure_ready=ensure_ready)
+        self._aborted = True
+
+    def _start(self, iterator, pre_dispatch):
+        # Only set self._iterating to True if at least a batch
+        # was dispatched. In particular this covers the edge
+        # case of Parallel used with an exhausted iterator. If
+        # self._original_iterator is None, then this means either
+        # that pre_dispatch == "all", n_jobs == 1 or that the first batch
+        # was very quick and its callback already dispatched all the
+        # remaining jobs.
+        self._iterating = False
+        if self.dispatch_one_batch(iterator):
+            self._iterating = self._original_iterator is not None
+
+        while self.dispatch_one_batch(iterator):
+            pass
+
+        if pre_dispatch == "all":
+            # The iterable was consumed all at once by the above for loop.
+            # No need to wait for async callbacks to trigger to
+            # consumption.
+            self._iterating = False
+
+    def _get_outputs(self, iterator, pre_dispatch):
+        """Iterator returning the tasks' output as soon as they are ready."""
+        current_thread_id = threading.get_ident()
+        pypy_detach_generator_exit = False
+        try:
+            self._start(iterator, pre_dispatch)
+            # first yield returns None, for internal use only. This ensures
+            # that we enter the try/except block and start dispatching the
+            # tasks.
+            yield
+
+            with self._backend.retrieval_context():
+                yield from self._retrieve()
+
+        except GeneratorExit:
+            # The generator has been garbage collected before being fully
+            # consumned. This aborts the remaining tasks if possible and warn
+            # the user if necessary.
+            self._exception = True
+
+            # In CPython, the gc and the error happen in the main thread.
+            # However, this is not the case for other interpreters such as in
+            # PyPy, where this block can be executed in arbitrary threads. This
+            # can lead to hang when a process tries to join itself. As a
+            # workaround, we detach the execution of the aborting code to a
+            # dedicated thread. We then need to make sure the rest of the
+            # function does not call `_terminate_and_reset` in finally.
+            if IS_PYPY and current_thread_id != threading.get_ident():
+                pypy_detach_generator_exit = True
+                _parallel = self
+
+                class _PypyGeneratorExitThread(threading.Thread):
+                    def run(self):
+                        _parallel._abort()
+                        if _parallel.return_generator:
+                            _parallel._warn_exit_early()
+                        _parallel._terminate_and_reset()
+
+                _PypyGeneratorExitThread(
+                    name="PypyGeneratorExitThread"
+                ).start()
+                return
+
+            # Otherwise, we are in CPython or in PyPy in the main thread and
+            # we can safely abort the execution and warn the user.
+            self._abort()
+            if self.return_generator:
+                self._warn_exit_early()
+
+            raise
+
+        # Note: we catch any BaseException instead of just Exception instances
+        # to also include KeyboardInterrupt
+        except BaseException:
+            self._exception = True
+            self._abort()
+            raise
+        finally:
+            # Store the unconsumed tasks and terminate the workers if necessary
+            _remaining_outputs = ([] if self._exception else self._jobs)
+            self._jobs = collections.deque()
+            self._running = False
+            if not pypy_detach_generator_exit:
+                self._terminate_and_reset()
+
+        while len(_remaining_outputs) > 0:
+            batched_results = _remaining_outputs.popleft()
+            batched_results = batched_results.get_result(self.timeout)
+            for result in batched_results:
+                yield result
+
+    def _wait_retrieval(self):
+        """Return True if we need to continue retriving some tasks."""
+
+        # If the input load is still being iterated over, it means that tasks
+        # are still on the dispatch wait list and their results will need to
+        # be retrieved later on.
+        if self._iterating:
+            return True
+
+        # If some of the dispatched tasks are still being processed by the
+        # workers, wait for the compute to finish before starting retrieval
+        if self.n_completed_tasks < self.n_dispatched_tasks:
+            return True
+
+        # For backends that does not support retrieving asynchronously the
+        # result to the main process, all results must be carefully retrieved
+        # in the _retrieve loop in the main thread while the backend is alive.
+        # For other backends, the actual retrieval is done asynchronously in
+        # the callback thread, and we can terminate the backend before the
+        # `self._jobs` result list has been emptied. The remaining results
+        # will be collected in the `finally` step of the generator.
+        if not self._backend.supports_retrieve_callback:
+            if len(self._jobs) > 0:
+                return True
+
+        return False
+
+    def _retrieve(self):
+        while self._wait_retrieval():
+
+            # If the callback thread of a worker has signaled that its task
+            # triggerd an exception, or if the retrieval loop has raised an
+            # exception (e.g. `GeneratorExit`), exit the loop and surface the
+            # worker traceback.
+            if self._aborting:
+                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)):
                 time.sleep(0.01)
                 continue
+
             # 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
+            # we empty it and Python list are not thread-safe by
+            # default hence the use of the lock
             with self._lock:
-                job = self._jobs.pop(0)
+                batched_results = self._jobs.popleft()
 
-            try:
-                if getattr(self._backend, 'supports_timeout', False):
-                    self._output.extend(job.get(timeout=self.timeout))
-                else:
-                    self._output.extend(job.get())
-
-            except BaseException:
-                # Note: we catch any BaseException instead of just Exception
-                # instances to also include KeyboardInterrupt.
-
-                # Stop dispatching any new job in the async callback thread
-                self._aborting = True
-
-                # If the backend allows it, cancel or kill remaining running
-                # tasks without waiting for the results as we will raise
-                # the exception we got back to the caller instead of returning
-                # any result.
-                backend = self._backend
-                if (backend is not None and
-                        hasattr(backend, 'abort_everything')):
-                    # If the backend is managed externally we need to make sure
-                    # to leave it in a working state to allow for future jobs
-                    # scheduling.
-                    ensure_ready = self._managed_backend
-                    backend.abort_everything(ensure_ready=ensure_ready)
-                raise
+            # Flatten the batched results to output one output at a time
+            batched_results = batched_results.get_result(self.timeout)
+            for result in batched_results:
+                self._nb_consumed += 1
+                yield result
 
-    def __call__(self, iterable):
-        if self._jobs:
-            raise ValueError('This Parallel instance is already running')
-        # A flag used to abort the dispatching of jobs in case an
-        # exception is found
+    def _raise_error_fast(self):
+        """If we are aborting, raise if a job caused an error."""
+
+        # Find the first job whose status is TASK_ERROR if it exists.
+        with self._lock:
+            error_job = next((job for job in self._jobs
+                              if job.status == TASK_ERROR), None)
+
+        # If this error job exists, immediatly raise the error by
+        # calling get_result. This job might not exists if abort has been
+        # called directly or if the generator is gc'ed.
+        if error_job is not None:
+            error_job.get_result(self.timeout)
+
+    def _warn_exit_early(self):
+        """Warn the user if the generator is gc'ed before being consumned."""
+        ready_outputs = self.n_completed_tasks - self._nb_consumed
+        is_completed = self._is_completed()
+        msg = ""
+        if ready_outputs:
+            msg += (
+                f"{ready_outputs} tasks have been successfully executed "
+                " but not used."
+            )
+            if not is_completed:
+                msg += " Additionally, "
+
+        if not is_completed:
+            msg += (
+                f"{self.n_dispatched_tasks - self.n_completed_tasks} tasks "
+                "which were still being processed by the workers have been "
+                "cancelled."
+            )
+
+        if msg:
+            msg += (
+                " You could benefit from adjusting the input task "
+                "iterator to limit unnecessary computation time."
+            )
+
+            warnings.warn(msg)
+
+    def _get_sequential_output(self, iterable):
+        """Separate loop for sequential output.
+
+        This simplifies the traceback in case of errors and reduces the
+        overhead of calling sequential tasks with `joblib`.
+        """
+        try:
+            self._iterating = True
+            self._original_iterator = iterable
+            batch_size = self._get_batch_size()
+
+            if batch_size != 1:
+                it = iter(iterable)
+                iterable_batched = iter(
+                    lambda: tuple(itertools.islice(it, batch_size)), ()
+                )
+                iterable = (
+                    task for batch in iterable_batched for task in batch
+                )
+
+            # first yield returns None, for internal use only. This ensures
+            # that we enter the try/except block and setup the generator.
+            yield None
+
+            # Sequentially call the tasks and yield the results.
+            for func, args, kwargs in iterable:
+                self.n_dispatched_batches += 1
+                self.n_dispatched_tasks += 1
+                res = func(*args, **kwargs)
+                self.n_completed_tasks += 1
+                self.print_progress()
+                yield res
+                self._nb_consumed += 1
+        except BaseException:
+            self._exception = True
+            self._aborting = True
+            self._aborted = True
+            raise
+        finally:
+            self.print_progress()
+            self._running = False
+            self._iterating = False
+            self._original_iterator = None
+
+    def _reset_run_tracking(self):
+        """Reset the counters and flags used to track the execution."""
+
+        # Makes sur the parallel instance was not previously running in a
+        # thread-safe way.
+        with getattr(self, '_lock', nullcontext()):
+            if self._running:
+                msg = 'This Parallel instance is already running !'
+                if self.return_generator is True:
+                    msg += (
+                        " Before submitting new tasks, you must wait for the "
+                        "completion of all the previous tasks, or clean all "
+                        "references to the output generator."
+                    )
+                raise RuntimeError(msg)
+            self._running = True
+
+        # Counter to keep track of the task dispatched and completed.
+        self.n_dispatched_batches = 0
+        self.n_dispatched_tasks = 0
+        self.n_completed_tasks = 0
+
+        # Following count is incremented by one each time the user iterates
+        # on the output generator, it is used to prepare an informative
+        # warning message in case the generator is deleted before all the
+        # dispatched tasks have been consumed.
+        self._nb_consumed = 0
+
+        # Following flags are used to synchronize the threads in case one of
+        # the tasks error-out to ensure that all workers abort fast and that
+        # the backend terminates properly.
+
+        # Set to True as soon as a worker signals that a task errors-out
+        self._exception = False
+        # Set to True in case of early termination following an incident
         self._aborting = False
+        # Set to True after abortion is complete
+        self._aborted = False
+
+    def __call__(self, iterable):
+        """Main function to dispatch parallel tasks."""
+
+        self._reset_run_tracking()
+        self._start_time = time.time()
 
         if not self._managed_backend:
             n_jobs = self._initialize_backend()
         else:
             n_jobs = self._effective_n_jobs()
 
+        if n_jobs == 1:
+            # If n_jobs==1, run the computation sequentially and return
+            # immediatly to avoid overheads.
+            output = self._get_sequential_output(iterable)
+            next(output)
+            return output if self.return_generator else list(output)
+
+        # Let's create an ID that uniquely identifies the current call. If the
+        # call is interrupted early and that the same instance is immediately
+        # re-used, this id will be used to prevent workers that were
+        # concurrently finalizing a task from the previous call to run the
+        # callback.
+        with self._lock:
+            self._call_id = uuid4().hex
+
+        # self._effective_n_jobs should be called in the Parallel.__call__
+        # thread only -- store its value in an attribute for further queries.
+        self._cached_effective_n_jobs = n_jobs
+
         if isinstance(self._backend, LokyBackend):
             # For the loky backend, we add a callback executed when reducing
             # BatchCalls, that makes the loky executor use a temporary folder
@@ -1050,14 +1569,19 @@ class Parallel(Logger):
         if n_jobs == 0:
             raise RuntimeError("%s has no active worker." % backend_name)
 
-        self._print("Using backend %s with %d concurrent workers.",
-                    (backend_name, n_jobs))
+        self._print(
+            f"Using backend {backend_name} with {n_jobs} concurrent workers."
+        )
         if hasattr(self._backend, 'start_call'):
             self._backend.start_call()
+
+        # Following flag prevents double calls to `backend.stop_call`.
+        self._calling = True
+
         iterator = iter(iterable)
         pre_dispatch = self.pre_dispatch
 
-        if pre_dispatch == 'all' or n_jobs == 1:
+        if pre_dispatch == 'all':
             # prevent further dispatch via multiprocessing callback thread
             self._original_iterator = None
             self._pre_dispatch_amount = 0
@@ -1076,53 +1600,22 @@ class Parallel(Logger):
             # TODO: this iterator should be batch_size * n_jobs
             iterator = itertools.islice(iterator, self._pre_dispatch_amount)
 
-        self._start_time = time.time()
-        self.n_dispatched_batches = 0
-        self.n_dispatched_tasks = 0
-        self.n_completed_tasks = 0
         # Use a caching dict for callables that are pickled with cloudpickle to
         # improve performances. This cache is used only in the case of
-        # functions that are defined in the __main__ module, functions that are
-        # defined locally (inside another function) and lambda expressions.
+        # functions that are defined in the __main__ module, functions that
+        # are defined locally (inside another function) and lambda expressions.
         self._pickle_cache = dict()
-        try:
-            # Only set self._iterating to True if at least a batch
-            # was dispatched. In particular this covers the edge
-            # case of Parallel used with an exhausted iterator. If
-            # self._original_iterator is None, then this means either
-            # that pre_dispatch == "all", n_jobs == 1 or that the first batch
-            # was very quick and its callback already dispatched all the
-            # remaining jobs.
-            self._iterating = False
-            if self.dispatch_one_batch(iterator):
-                self._iterating = self._original_iterator is not None
 
-            while self.dispatch_one_batch(iterator):
-                pass
+        output = self._get_outputs(iterator, pre_dispatch)
+        self._call_ref = weakref.ref(output)
 
-            if pre_dispatch == "all" or n_jobs == 1:
-                # The iterable was consumed all at once by the above for loop.
-                # No need to wait for async callbacks to trigger to
-                # consumption.
-                self._iterating = False
+        # The first item from the output is blank, but it makes the interpreter
+        # progress until it enters the Try/Except block of the generator and
+        # reach the first `yield` statement. This starts the aynchronous
+        # dispatch of the tasks to the workers.
+        next(output)
 
-            with self._backend.retrieval_context():
-                self.retrieve()
-            # Make sure that we get a last message telling us we are done
-            elapsed_time = time.time() - self._start_time
-            self._print('Done %3i out of %3i | elapsed: %s finished',
-                        (len(self._output), len(self._output),
-                         short_format_time(elapsed_time)))
-        finally:
-            if hasattr(self._backend, 'stop_call'):
-                self._backend.stop_call()
-            if not self._managed_backend:
-                self._terminate_backend()
-            self._jobs = list()
-            self._pickle_cache = None
-        output = self._output
-        self._output = None
-        return output
+        return output if self.return_generator else list(output)
 
     def __repr__(self):
         return '%s(n_jobs=%s)' % (self.__class__.__name__, self.n_jobs)


Test patch:
diff --git a/joblib/test/common.py b/joblib/test/common.py
index ee5e398..b0ca0c6 100644
--- a/joblib/test/common.py
+++ b/joblib/test/common.py
@@ -3,6 +3,7 @@ Small utilities for testing.
 """
 import os
 import gc
+import sys
 
 from joblib._multiprocessing_helpers import mp
 from joblib.testing import SkipTest, skipif
@@ -12,6 +13,8 @@ try:
 except ImportError:
     lz4 = None
 
+IS_PYPY = hasattr(sys, "pypy_version_info")
+
 # A decorator to run tests only when numpy is available
 try:
     import numpy as np
@@ -55,6 +58,18 @@ except ImportError:
     memory_usage = memory_used = None
 
 
+def force_gc_pypy():
+    # The gc in pypy can be delayed. Force it to test the behavior when it
+    # will eventually be collected.
+    if IS_PYPY:
+        # Run gc.collect() twice to make sure the weakref is collected, as
+        # mentionned in the pypy doc:
+        # https://doc.pypy.org/en/latest/config/objspace.usemodules._weakref.html
+        import gc
+        gc.collect()
+        gc.collect()
+
+
 with_multiprocessing = skipif(
     mp is None, reason='Needs multiprocessing to run.')
 
diff --git a/joblib/test/test_parallel.py b/joblib/test/test_parallel.py
index 81ab7aa..57b554b 100644
--- a/joblib/test/test_parallel.py
+++ b/joblib/test/test_parallel.py
@@ -10,12 +10,14 @@ import os
 import sys
 import time
 import mmap
+import weakref
 import warnings
 import threading
 from traceback import format_exception
 from math import sqrt
 from time import sleep
 from pickle import PicklingError
+from contextlib import nullcontext
 from multiprocessing import TimeoutError
 
 import pytest
@@ -28,6 +30,7 @@ from joblib._multiprocessing_helpers import mp
 
 from joblib.test.common import np, with_numpy
 from joblib.test.common import with_multiprocessing
+from joblib.test.common import IS_PYPY, force_gc_pypy
 from joblib.testing import (parametrize, raises, check_subprocess_call,
                             skipif, warns)
 
@@ -57,16 +60,17 @@ from joblib._parallel_backends import ThreadingBackend
 from joblib._parallel_backends import MultiprocessingBackend
 from joblib._parallel_backends import ParallelBackendBase
 from joblib._parallel_backends import LokyBackend
-from joblib._parallel_backends import SafeFunction
 
 from joblib.parallel import Parallel, delayed
 from joblib.parallel import register_parallel_backend, parallel_backend
 from joblib.parallel import effective_n_jobs, cpu_count
 
 from joblib.parallel import mp, BACKENDS, DEFAULT_BACKEND, EXTERNAL_BACKENDS
-from joblib.my_exceptions import WorkerInterrupt
 
 
+RETURN_GENERATOR_BACKENDS = BACKENDS.copy()
+RETURN_GENERATOR_BACKENDS.pop("multiprocessing", None)
+
 ALL_VALID_BACKENDS = [None] + sorted(BACKENDS.keys())
 # Add instances of backend classes deriving from ParallelBackendBase
 ALL_VALID_BACKENDS += [BACKENDS[backend_str]() for backend_str in BACKENDS]
@@ -246,6 +250,9 @@ def test_nested_parallel_warnings(parent_backend, child_backend, expected):
     # warning handling is not thread safe. One thread might see multiple
     # warning or no warning at all.
     if parent_backend == "threading":
+        if IS_PYPY and not any(res):
+            # Related to joblib#1426, should be removed once it is solved.
+            pytest.xfail(reason="This test often fails in PyPy.")
         assert any(res)
     else:
         assert all(res)
@@ -356,7 +363,7 @@ def test_parallel_pickling():
 @parametrize('backend', PARALLEL_BACKENDS)
 def test_parallel_timeout_success(backend):
     # Check that timeout isn't thrown when function is fast enough
-    assert len(Parallel(n_jobs=2, backend=backend, timeout=10)(
+    assert len(Parallel(n_jobs=2, backend=backend, timeout=30)(
         delayed(sleep)(0.001) for x in range(10))) == 10
 
 
@@ -379,7 +386,8 @@ def test_error_capture(backend):
             Parallel(n_jobs=2, backend=backend)(
                 [delayed(division)(x, y)
                     for x, y in zip((0, 1), (1, 0))])
-        with raises(WorkerInterrupt):
+
+        with raises(KeyboardInterrupt):
             Parallel(n_jobs=2, backend=backend)(
                 [delayed(interrupt_raiser)(x) for x in (1, 0)])
 
@@ -403,7 +411,7 @@ def test_error_capture(backend):
                     parallel(delayed(f)(x, y=1) for x in range(10)))
 
             original_workers = get_workers(parallel._backend)
-            with raises(WorkerInterrupt):
+            with raises(KeyboardInterrupt):
                 parallel([delayed(interrupt_raiser)(x) for x in (1, 0)])
 
             # The pool should still be available despite the exception
@@ -413,7 +421,10 @@ def test_error_capture(backend):
             assert get_workers(parallel._backend) is not original_workers
 
             assert ([f(x, y=1) for x in range(10)] ==
-                    parallel(delayed(f)(x, y=1) for x in range(10)))
+                    parallel(delayed(f)(x, y=1) for x in range(10))), (
+                parallel._iterating, parallel.n_completed_tasks,
+                parallel.n_dispatched_tasks, parallel._aborting
+            )
 
         # Check that the inner pool has been terminated when exiting the
         # context manager
@@ -818,11 +829,13 @@ def test_backend_nesting_level(outer_backend, inner_backend):
 
 
 @with_multiprocessing
-def test_retrieval_context():
+@parametrize('with_retrieve_callback', [True, False])
+def test_retrieval_context(with_retrieve_callback):
     import contextlib
 
     class MyBackend(ThreadingBackend):
         i = 0
+        supports_retrieve_callback = with_retrieve_callback
 
         @contextlib.contextmanager
         def retrieval_context(self):
@@ -845,16 +858,6 @@ def test_retrieval_context():
 ###############################################################################
 # Test helpers
 
-def test_safe_function():
-    safe_division = SafeFunction(division)
-    with raises(ZeroDivisionError):
-        safe_division(1, 0)
-
-    safe_interrupt = SafeFunction(interrupt_raiser)
-    with raises(WorkerInterrupt):
-        safe_interrupt('x')
-
-
 @parametrize('batch_size', [0, -1, 1.42])
 def test_invalid_batch_size(batch_size):
     with raises(ValueError):
@@ -1116,6 +1119,12 @@ def test_parallel_with_exhausted_iterator():
     assert Parallel(n_jobs=2)(exhausted_iterator) == []
 
 
+def _cleanup_worker():
+    """Helper function to force gc in each worker."""
+    force_gc_pypy()
+    time.sleep(.1)
+
+
 def check_memmap(a):
     if not isinstance(a, np.memmap):
         raise TypeError('Expected np.memmap instance, got %r',
@@ -1170,7 +1179,7 @@ def test_memmap_with_big_offset(tmpdir):
 
 def test_warning_about_timeout_not_supported_by_backend():
     with warnings.catch_warnings(record=True) as warninfo:
-        Parallel(timeout=1)(delayed(square)(i) for i in range(50))
+        Parallel(n_jobs=1, timeout=1)(delayed(square)(i) for i in range(50))
     assert len(warninfo) == 1
     w = warninfo[0]
     assert isinstance(w.message, UserWarning)
@@ -1180,6 +1189,27 @@ def test_warning_about_timeout_not_supported_by_backend():
         "will not be used.")
 
 
+def set_list_value(input_list, index, value):
+    input_list[index] = value
+    return value
+
+
+@pytest.mark.parametrize('n_jobs', [1, 2, 4])
+def test_parallel_return_generator_order(n_jobs):
+    # This test inserts values in a list in some expected order
+    # in sequential computing, and then checks that this order has been
+    # respected by Parallel output generator.
+    input_list = [0] * 5
+    result = Parallel(n_jobs=n_jobs, return_generator=True,
+                      backend='threading')(
+        delayed(set_list_value)(input_list, i, i) for i in range(5))
+
+    # Ensure that all the tasks are completed before checking the result
+    result = list(result)
+
+    assert all(v == r for v, r in zip(input_list, result))
+
+
 @parametrize('backend', ALL_VALID_BACKENDS)
 @parametrize('n_jobs', [1, 2, -2, -1])
 def test_abort_backend(n_jobs, backend):
@@ -1192,6 +1222,144 @@ def test_abort_backend(n_jobs, backend):
     assert dt < 20
 
 
+def get_large_object(arg):
+    result = np.ones(int(5 * 1e5), dtype=bool)
+    result[0] = False
+    return result
+
+
+@with_numpy
+@parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('n_jobs', [1, 2, -2, -1])
+def test_deadlock_with_generator(backend, n_jobs):
+    # Non-regression test for a race condition in the backends when the pickler
+    # is delayed by a large object.
+    with Parallel(n_jobs=n_jobs, backend=backend,
+                  return_generator=True) as parallel:
+        result = parallel(delayed(get_large_object)(i) for i in range(10))
+        next(result)
+        next(result)
+        del result
+        # The gc in pypy can be delayed. Force it to make sure this test does
+        # not cause timeout on the CI.
+        force_gc_pypy()
+
+
+@parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('n_jobs', [1, 2, -2, -1])
+def test_multiple_generator_call(backend, n_jobs):
+    # Non-regression test that ensures the dispatch of the tasks starts
+    # immediately when Parallel.__call__ is called. This test relies on the
+    # assumption that only one generator can be submitted at a time.
+    with raises(RuntimeError,
+                match="This Parallel instance is already running"):
+        parallel = Parallel(n_jobs, backend=backend, return_generator=True)
+        g = parallel(delayed(sleep)(1) for _ in range(10))  # noqa: F841
+        t_start = time.time()
+        gen2 = parallel(delayed(id)(i) for i in range(100))  # noqa: F841
+
+    # Make sure that the error is raised quickly
+    assert time.time() - t_start < 2, (
+        "The error should be raised immediatly when submitting a new task "
+        "but it took more than 2s."
+    )
+
+    del g
+    # The gc in pypy can be delayed. Force it to make sure this test does not
+    # cause timeout on the CI.
+    force_gc_pypy()
+
+
+@parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('n_jobs', [1, 2, -2, -1])
+def test_multiple_generator_call_managed(backend, n_jobs):
+    # Non-regression test that ensures the dispatch of the tasks starts
+    # immediately when Parallel.__call__ is called. This test relies on the
+    # assumption that only one generator can be submitted at a time.
+    with Parallel(n_jobs, backend=backend,
+                  return_generator=True) as parallel:
+        g = parallel(delayed(sleep)(10) for _ in range(10))  # noqa: F841
+        t_start = time.time()
+        with raises(RuntimeError,
+                    match="This Parallel instance is already running"):
+            g2 = parallel(delayed(id)(i) for i in range(100))  # noqa: F841
+
+        # Make sure that the error is raised quickly
+        assert time.time() - t_start < 2, (
+            "The error should be raised immediatly when submitting a new task "
+            "but it took more than 2s."
+        )
+
+    # The gc in pypy can be delayed. Force it to make sure this test does not
+    # cause timeout on the CI.
+    del g
+    force_gc_pypy()
+
+
+@parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('n_jobs', [1, 2, -2, -1])
+def test_multiple_generator_call_separated(backend, n_jobs):
+    # Check that for separated Parallel, both tasks are correctly returned.
+    g = Parallel(n_jobs, backend=backend, return_generator=True)(
+        delayed(sqrt)(i ** 2) for i in range(10)
+    )
+    g2 = Parallel(n_jobs, backend=backend, return_generator=True)(
+        delayed(sqrt)(i ** 2) for i in range(10, 20)
+    )
+
+    assert all(res == i for res, i in zip(g, range(10)))
+    assert all(res == i for res, i in zip(g2, range(10, 20)))
+
+
+@parametrize('backend, error', [
+    ('loky', True),
+    ('threading', False),
+    ('sequential', False),
+])
+def test_multiple_generator_call_separated_gc(backend, error):
+
+    if (backend == 'loky') and (mp is None):
+        pytest.skip("Requires multiprocessing")
+
+    # Check that in loky, only one call can be run at a time with
+    # a single executor.
+    parallel = Parallel(2, backend=backend, return_generator=True)
+    g = parallel(delayed(sleep)(10) for i in range(10))
+    g_wr = weakref.finalize(g, lambda: print("Generator collected"))
+    ctx = (
+        raises(RuntimeError, match="The executor underlying Parallel")
+        if error else nullcontext()
+    )
+    with ctx:
+        # For loky, this call will raise an error as the gc of the previous
+        # generator will shutdown the shared executor.
+        # For the other backends, as the worker pools are not shared between
+        # the two calls, this should proceed correctly.
+        t_start = time.time()
+        g = Parallel(2, backend=backend, return_generator=True)(
+            delayed(sqrt)(i ** 2) for i in range(10, 20)
+        )
+
+        # The gc in pypy can be delayed. Force it to test the behavior when it
+        # will eventually be collected.
+        force_gc_pypy()
+        assert all(res == i for res, i in zip(g, range(10, 20)))
+
+    assert time.time() - t_start < 5
+
+    # Make sure that the computation are stopped for the gc'ed generator
+    retry = 0
+    while g_wr.alive and retry < 3:
+        retry += 1
+        time.sleep(.5)
+    assert time.time() - t_start < 5
+
+    if parallel._effective_n_jobs() != 1:
+        # check that the first parallel object is aborting (the final _aborted
+        # state might be delayed).
+        assert parallel._aborting
+
+
 @with_numpy
 @with_multiprocessing
 @parametrize('backend', PROCESS_BACKENDS)
@@ -1209,6 +1377,13 @@ def test_memmapping_leaks(backend, tmpdir):
         # The memmap folder should not be clean in the context scope
         assert len(os.listdir(tmpdir)) > 0
 
+        # Cleaning of the memmap folder is triggered by the garbage
+        # collection. With pypy the garbage collection has been observed to be
+        # delayed, sometimes up until the shutdown of the interpreter. This
+        # cleanup job executed in the worker ensures that it's triggered
+        # immediately.
+        p(delayed(_cleanup_worker)() for _ in range(2))
+
     # Make sure that the shared memory is cleaned at the end when we exit
     # the context
     for _ in range(100):
@@ -1221,6 +1396,7 @@ def test_memmapping_leaks(backend, tmpdir):
     # Make sure that the shared memory is cleaned at the end of a call
     p = Parallel(n_jobs=2, max_nbytes=1, backend=backend)
     p(delayed(check_memmap)(a) for a in [np.random.random(10)] * 2)
+    p(delayed(_cleanup_worker)() for _ in range(2))
 
     for _ in range(100):
         if not os.listdir(tmpdir):
@@ -1481,7 +1657,7 @@ def test_thread_bomb_mitigation(backend):
         # Local import because loky may not be importable for lack of
         # multiprocessing
         from joblib.externals.loky.process_executor import TerminatedWorkerError # noqa
-        if isinstance(exc, TerminatedWorkerError):
+        if isinstance(exc, (TerminatedWorkerError, PicklingError)):
             # The recursion exception can itself cause an error when
             # pickling it to be send back to the parent process. In this
             # case the worker crashes but the original traceback is still
diff --git a/joblib/test/test_testing.py b/joblib/test/test_testing.py
index cc94463..e8095aa 100644
--- a/joblib/test/test_testing.py
+++ b/joblib/test/test_testing.py
@@ -60,7 +60,7 @@ def test_check_subprocess_call_timeout():
         'sys.stderr.flush()',
         # We need to sleep for at least 2 * timeout seconds in case the SIGKILL
         # is triggered.
-        'time.sleep(2.1)',
+        'time.sleep(10)',
         'print("process should have be killed before")',
         'sys.stdout.flush()'])