Repository: joblib/joblib
Pull request #1463 (2023-06-26 08:10:26)

FEA Implement generator unordered parameter
Closes https://github.com/joblib/joblib/issues/1449 

The PR builds upon https://github.com/joblib/joblib/pull/1458 

It looks straightforward, the diff in `parallel.py` outside of documentation is about 10 lines added/removed.

More testing to be added, and an extension to the tutorial, that highlights how RAM usage can be further decreased if the tasks have significantly different processing times maybe ?


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
Patch coverage: **`66.66`**% and project coverage change: **`-26.07`** :warning:
> Comparison is base [(`b469022`)](https://app.codecov.io/gh/joblib/joblib/commit/b46902217ebf071cddd5076826f8fcce06aecbec?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) 94.93% compared to head [(`7bb322f`)](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) 68.86%.

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


```diff
@@             Coverage Diff             @@
##           master    #1463       +/-   ##
===========================================
- Coverage   94.93%   68.86%   -26.07%     
===========================================
  Files          44       44               
  Lines        7439     7247      -192     
===========================================
- Hits         7062     4991     -2071     
- Misses        377     2256     +1879     
```


| [Impacted Files](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | Coverage Δ | |
|---|---|---|
| [joblib/\_parallel\_backends.py](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19wYXJhbGxlbF9iYWNrZW5kcy5weQ==) | `41.26% <ø> (-52.20%)` | :arrow_down: |
| [joblib/parallel.py](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3BhcmFsbGVsLnB5) | `54.54% <42.85%> (-42.70%)` | :arrow_down: |
| [joblib/test/test\_parallel.py](https://app.codecov.io/gh/joblib/joblib/pull/1463?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9wYXJhbGxlbC5weQ==) | `53.77% <100.00%> (-42.43%)` | :arrow_down: |

... and [32 files with indirect coverage changes](https://app.codecov.io/gh/joblib/joblib/pull/1463/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/joblib/joblib/pull/1463?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: I don't think the fail on pypy is related to the PR and it's benign
ogrisel: There are two warnings making the HTML rendering fail on the CI:

```
/home/docs/checkouts/readthedocs.org/user_builds/joblib/envs/1463/lib/python3.7/site-packages/joblib/parallel.py:docstring of joblib.parallel.Parallel:46: WARNING: Block quote ends without a blank line; unexpected unindent.
/home/docs/checkouts/readthedocs.org/user_builds/joblib/envs/1463/lib/python3.7/site-packages/joblib/parallel.py:docstring of joblib.parallel.Parallel:45: WARNING: Block quote ends without a blank line; unexpected unindent.
```

could you please fix those to make it possible to review the rendered HTML for the examples expanded in this PR?
fcharras: The example is running now.  The pypy test fails consistently on this PR though, something new in this PR probably has been delaying backend abortion to the point that this test fails now :thinking: 

fcharras: NB: current pipeline issues are not related to the latest commit, it seems that the pipeline setup is down because of dependencies issues.
ogrisel: @fcharras I think we do not need the `_jobs_unordered` set. Let's see if this the tests pass on the CI with that suggested change in #1500.
ogrisel: Once the above comments are addressed, LGTM.
ogrisel: > Regarding [fce3475](https://github.com/joblib/joblib/commit/fce34755f7cb8bb9ee62dd7c31266c706bd47fcf) I thought about that too but in this case after a new task is fired the main process doesn't keep any reference to it, it could just be garbage collected and then strange behavior might happen, and said behavior could depend on the backend ?

It's not possible for a backend to call callback successfully without keeping a reference to the callback object. If it does not call the callback, then the backend would be broken for any the value of `return_as`.
fcharras: But couldn't the opposite be true: _because_ some given backend does not keep a reference to any task object internally, then the callback might not be called if the user-level reference is deleted ? basically a backend could use this for early-interruption of tasks  
ogrisel: Please also document the new feature in the changelog, targeting the next release.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index f6933bf..48c9fba 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,6 +7,15 @@ In development
 - Ensure that errors in the task generator given to Parallel's call
   are raised in the results consumming thread.
   https://github.com/joblib/joblib/pull/1491
+  
+- Adjust codebase to NumPy 2.0 by changing ``np.NaN`` to ``np.nan``
+  and importing ``byte_bounds`` from ``np.lib.array_utils``.
+  https://github.com/joblib/joblib/pull/1501
+
+- The parameter ``return_as`` in ``joblib.Parallel`` can now be set to
+  ``generator_unordered``. In this case the results will be returned in the
+  order of task completion rather than the order of submission.
+  https://github.com/joblib/joblib/pull/1463
 
 Release 1.3.2 -- 2023/08/08
 ---------------------------
diff --git a/benchmarks/bench_pickle.py b/benchmarks/bench_pickle.py
index 084d15a..671d57d 100755
--- a/benchmarks/bench_pickle.py
+++ b/benchmarks/bench_pickle.py
@@ -51,7 +51,7 @@ def delete_obj(obj):
 def memory_used(func, *args, **kwargs):
     """Compute memory usage of func."""
     if memory_usage is None:
-        return np.NaN
+        return np.nan
 
     gc.collect()
     mem_use = memory_usage((func, args, kwargs), interval=.001)
diff --git a/examples/parallel_generator.py b/examples/parallel_generator.py
index 286518f..e41ff1f 100644
--- a/examples/parallel_generator.py
+++ b/examples/parallel_generator.py
@@ -14,9 +14,10 @@ Using ``return_as='generator'`` allows to progressively consume the outputs
 as they arrive and keeps the memory at an acceptable level.
 
 In this case, the output of the `Parallel` call is a generator that yields the
-results in the order the tasks have been submitted with. Future releases are
-also planned to support the ``return_as="unordered_generator"`` parameter to
-have the generator yield results as soon as available.
+results in the order the tasks have been submitted with. If the order of the
+tasks does not matter (for instance if they are consumed by a commutative
+aggregation function), then using ``return_as='generator_unordered'`` can be
+even more efficient.
 
 """
 
@@ -156,6 +157,7 @@ print(f"Peak memory usage: {peak:.2f}MB")
 # the results.
 
 import matplotlib.pyplot as plt
+plt.figure(0)
 plt.semilogy(
     np.maximum.accumulate(monitor.memory_buffer),
     label='return_as="list"'
@@ -176,3 +178,114 @@ plt.show()
 # 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.
+
+
+##############################################################################
+# Further memory efficiency for commutative aggregation
+##############################################################################
+
+##############################################################################
+# There is still room for improving the relief on memory allocation we get
+# using ``return_as="generator"``. Indeed, notice how the generator of the
+# previous example respects the order the tasks have been submitted with. This
+# behavior can cause a build up in memory of results waiting to be consumed,
+# in case some tasks finished before other tasks despite being submitted
+# later. The corresponding results will be kept in memory until the slower
+# tasks submitted earlier are done and have been iterated over.
+#
+# In case the downstream consumer of the results is reliant on the assumption
+# that the results are yielded in the same order that the tasks were submitted,
+# it can't be helped. But in our example, since the `+` operator is
+# commutative, the function ``accumulator_sum`` does not need the generator to
+# return the results with any particular order. In this case it's safe to use
+# the option ``return_as="generator_unordered"``, so that the results are
+# returned as soon as a task is completed, ignoring the order of task
+# submission.
+#
+# Beware that the downstream consumer of the results must not expect them be
+# returned with any deterministic or predictable order at all, since the
+# progress of the tasks can depend on the availability of the workers, which
+# can be affected by external events, such as system load, implementation
+# details in the backend, etc.
+
+
+##############################################################################
+# To better highlight improvements in memory usage when using the parameter
+# ``return_as="generator_unordered"``, let's explcitly add delay in some of
+# the submitted tasks.
+
+
+def return_big_object_delayed(i):
+    if (i + 20) % 60:
+        time.sleep(0.1)
+    else:
+        time.sleep(5)
+    return i * np.ones((10000, 200), dtype=np.float64)
+
+
+##############################################################################
+# Let's check memory usage when using ``return_as="generator"``...
+
+monitor_delayed_gen = MemoryMonitor()
+print("Create result generator on delayed tasks with return_as='generator'...")
+res = Parallel(n_jobs=2, return_as="generator")(
+    delayed(return_big_object_delayed)(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_delayed_gen.join()
+peak = max(monitor_delayed_gen.memory_buffer) / 1e6
+print(f"Peak memory usage: {peak:.2f}MB")
+
+##############################################################################
+# If we use ``return_as="generator_unordered"``, ``res`` will not enforce any
+# order when returning the results, and will simply enable iterating on the
+# results as soon as it's available. The peak memory usage is now controlled
+# to an even lower level, since that results can be consumed immediately
+# rather than being delayed by the compute of slower tasks that have been
+# submitted earlier.
+
+monitor_delayed_gen_unordered = MemoryMonitor()
+print(
+  "Create result generator on delayed tasks with "
+  "return_as='generator_unordered'..."
+)
+res = Parallel(n_jobs=2, return_as="generator_unordered")(
+    delayed(return_big_object_delayed)(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_delayed_gen_unordered.join()
+peak = max(monitor_delayed_gen_unordered.memory_buffer) / 1e6
+print(f"Peak memory usage: {peak:.2f}MB")
+
+
+##############################################################################
+# Notice how the plot for ``'return_as="generator'`` now shows a high memory
+# usage plateau when slow jobs cause a congestion of intermediate results
+# waiting in RAM before in-order aggregation. This high memory usage is never
+# observed when using ``'return_as="generator_unordered"``.
+
+plt.figure(1)
+plt.semilogy(
+    np.maximum.accumulate(monitor_delayed_gen.memory_buffer),
+    label='return_as="generator"'
+)
+plt.semilogy(
+    np.maximum.accumulate(monitor_delayed_gen_unordered.memory_buffer),
+    label='return_as="generator_unordered"'
+)
+plt.xlabel("Time")
+plt.xticks([], [])
+plt.ylabel("Memory usage")
+plt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])
+plt.legend()
+plt.show()
diff --git a/joblib/_memmapping_reducer.py b/joblib/_memmapping_reducer.py
index 283399c..13f5c4a 100644
--- a/joblib/_memmapping_reducer.py
+++ b/joblib/_memmapping_reducer.py
@@ -266,8 +266,13 @@ def _reduce_memmap_backed(a, m):
     # offset that comes from the striding differences between a and m
     util.debug('[MEMMAP REDUCE] reducing a memmap-backed array '
                '(shape, {}, pid: {})'.format(a.shape, os.getpid()))
-    a_start, a_end = np.byte_bounds(a)
-    m_start = np.byte_bounds(m)[0]
+    try:
+        from numpy.lib.array_utils import byte_bounds
+    except (ModuleNotFoundError, ImportError):
+        # Backward-compat for numpy < 2.0
+        from numpy import byte_bounds
+    a_start, a_end = byte_bounds(a)
+    m_start = byte_bounds(m)[0]
     offset = a_start - m_start
 
     # offset from the backing memmap
diff --git a/joblib/parallel.py b/joblib/parallel.py
index d8a32a3..af842b4 100644
--- a/joblib/parallel.py
+++ b/joblib/parallel.py
@@ -810,13 +810,18 @@ 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()
 
     def _dispatch_new(self):
         """Schedule the next batch of tasks to be processed."""
 
-        # This steps ensure that auto-baching works as expected.
+        # This steps ensure that auto-batching works as expected.
         this_batch_duration = time.time() - self.dispatch_timestamp
         self.parallel._backend.batch_completed(self.batch_size,
                                                this_batch_duration)
@@ -976,15 +981,17 @@ 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_config` context manager.
-        return_as: str in {'list', 'generator'}, default: 'list'
+        return_as: str in {'list', 'generator', 'generator_unordered'},
+            default: 'list'
             If 'list', calls to this instance will return a list, only when
             all results have been processed and retrieved.
             If 'generator', it will return a generator that yields the results
             as soon as they are available, in the order the tasks have been
             submitted with.
-            Future releases are planned to also support 'generator_unordered',
-            in which case the generator immediately yields available results
-            independently of the submission order.
+            If 'generator_unordered', the generator will immediately yield
+            available results independently of the submission order. The output
+            order is not deterministic in this case because it depends on the
+            concurrency of the workers.
         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_config` context manager.
@@ -1205,13 +1212,15 @@ class Parallel(Logger):
         self.timeout = timeout
         self.pre_dispatch = pre_dispatch
 
-        if return_as not in {"list", "generator"}:
+        if return_as not in {"list", "generator", "generator_unordered"}:
             raise ValueError(
                 'Expected `return_as` parameter to be a string equal to "list"'
-                f' or "generator", but got {return_as} instead'
+                f',"generator" or "generator_unordered", but got {return_as} '
+                "instead."
             )
         self.return_as = return_as
         self.return_generator = return_as != "list"
+        self.return_ordered = return_as != "generator_unordered"
 
         # Check if we are under a parallel_config or parallel_backend
         # context manager and use the config from the context manager
@@ -1379,7 +1388,14 @@ class Parallel(Logger):
         batch_tracker = BatchCompletionCallBack(
             dispatch_timestamp, batch_size, self
         )
-        self._jobs.append(batch_tracker)
+
+        if self.return_ordered:
+            self._jobs.append(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
+        # the queue by itself as soon as the callback is triggered to be able
+        # to return the results in the order of completion.
 
         job = self._backend.apply_async(batch, callback=batch_tracker)
         batch_tracker.register_job(job)


Test patch:
diff --git a/joblib/test/test_parallel.py b/joblib/test/test_parallel.py
index 2a8403b..642fcbc 100644
--- a/joblib/test/test_parallel.py
+++ b/joblib/test/test_parallel.py
@@ -1253,6 +1253,38 @@ def test_parallel_return_order_with_return_as_generator_parameter(n_jobs):
     assert all(v == r for v, r in zip(input_list, result))
 
 
+def _sqrt_with_delay(e, delay):
+    if delay:
+        sleep(30)
+    return sqrt(e)
+
+
+@pytest.mark.parametrize('n_jobs', [2, 4])
+# NB: for this test to work, the backend must be allowed to process tasks
+# concurrently, so at least two jobs with a non-sequential backend are
+# mandatory.
+@with_multiprocessing
+@parametrize('backend', set(RETURN_GENERATOR_BACKENDS) - {"sequential"})
+def test_parallel_unordered_generator_returns_fastest_first(backend, n_jobs):
+    # This test submits 10 tasks, but the second task is super slow. This test
+    # checks that the 9 other tasks return before the slow task is done, when
+    # `return_as` parameter is set to `'generator_unordered'`
+    result = Parallel(n_jobs=n_jobs, return_as="generator_unordered",
+                      backend=backend)(
+        delayed(_sqrt_with_delay)(i**2, (i == 1)) for i in range(10))
+
+    quickly_returned = sorted(next(result) for _ in range(9))
+
+    expected_quickly_returned = [0] + list(range(2, 10))
+
+    assert all(
+        v == r for v, r in zip(expected_quickly_returned, quickly_returned)
+    )
+
+    del result
+    force_gc_pypy()
+
+
 @parametrize('backend', ALL_VALID_BACKENDS)
 @parametrize('n_jobs', [1, 2, -2, -1])
 def test_abort_backend(n_jobs, backend):
@@ -1273,12 +1305,13 @@ def get_large_object(arg):
 
 @with_numpy
 @parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('return_as', ["generator", "generator_unordered"])
 @parametrize('n_jobs', [1, 2, -2, -1])
-def test_deadlock_with_generator(backend, n_jobs):
+def test_deadlock_with_generator(backend, return_as, 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_as="generator") as parallel:
+                  return_as=return_as) as parallel:
         result = parallel(delayed(get_large_object)(i) for i in range(10))
         next(result)
         next(result)
@@ -1289,14 +1322,15 @@ def test_deadlock_with_generator(backend, n_jobs):
 
 
 @parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('return_as', ["generator", "generator_unordered"])
 @parametrize('n_jobs', [1, 2, -2, -1])
-def test_multiple_generator_call(backend, n_jobs):
+def test_multiple_generator_call(backend, return_as, 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_as="generator")
+        parallel = Parallel(n_jobs, backend=backend, return_as=return_as)
         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
@@ -1314,13 +1348,14 @@ def test_multiple_generator_call(backend, n_jobs):
 
 
 @parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('return_as', ["generator", "generator_unordered"])
 @parametrize('n_jobs', [1, 2, -2, -1])
-def test_multiple_generator_call_managed(backend, n_jobs):
+def test_multiple_generator_call_managed(backend, return_as, 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_as="generator") as parallel:
+                  return_as=return_as) as parallel:
         g = parallel(delayed(sleep)(10) for _ in range(10))  # noqa: F841
         t_start = time.time()
         with raises(RuntimeError,
@@ -1340,16 +1375,26 @@ def test_multiple_generator_call_managed(backend, n_jobs):
 
 
 @parametrize('backend', RETURN_GENERATOR_BACKENDS)
+@parametrize('return_as_1', ["generator", "generator_unordered"])
+@parametrize('return_as_2', ["generator", "generator_unordered"])
 @parametrize('n_jobs', [1, 2, -2, -1])
-def test_multiple_generator_call_separated(backend, n_jobs):
+def test_multiple_generator_call_separated(
+        backend, return_as_1, return_as_2, n_jobs
+):
     # Check that for separated Parallel, both tasks are correctly returned.
-    g = Parallel(n_jobs, backend=backend, return_as="generator")(
+    g = Parallel(n_jobs, backend=backend, return_as=return_as_1)(
         delayed(sqrt)(i ** 2) for i in range(10)
     )
-    g2 = Parallel(n_jobs, backend=backend, return_as="generator")(
+    g2 = Parallel(n_jobs, backend=backend, return_as=return_as_2)(
         delayed(sqrt)(i ** 2) for i in range(10, 20)
     )
 
+    if return_as_1 == "generator_unordered":
+        g = sorted(g)
+
+    if return_as_2 == "generator_unordered":
+        g2 = sorted(g2)
+
     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)))
 
@@ -1359,14 +1404,18 @@ def test_multiple_generator_call_separated(backend, n_jobs):
     ('threading', False),
     ('sequential', False),
 ])
-def test_multiple_generator_call_separated_gc(backend, error):
+@parametrize('return_as_1', ["generator", "generator_unordered"])
+@parametrize('return_as_2', ["generator", "generator_unordered"])
+def test_multiple_generator_call_separated_gc(
+        backend, return_as_1, return_as_2, 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_as="generator")
+    parallel = Parallel(2, backend=backend, return_as=return_as_1)
     g = parallel(delayed(sleep)(10) for i in range(10))
     g_wr = weakref.finalize(g, lambda: print("Generator collected"))
     ctx = (
@@ -1379,13 +1428,17 @@ def test_multiple_generator_call_separated_gc(backend, error):
         # 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_as="generator")(
+        g = Parallel(2, backend=backend, return_as=return_as_2)(
             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()
+
+        if return_as_2 == "generator_unordered":
+            g = sorted(g)
+
         assert all(res == i for res, i in zip(g, range(10, 20)))
 
     assert time.time() - t_start < 5