Add option to specify an initialization function for `'loky'` and `'multiprocessing'` backends by shwina · Pull Request #1525 · joblib/joblib

shwina · 2023-11-20T14:16:21Z

Closes #381

This PR adds the ability to specify an initialization function that is run once per worker process when using the 'loky' and 'multiprocessing' backends.

Usage:

Parallel(initializer=my_func, initargs=(x, y, z))
# or
with parallel_config(initializer=my_func, initargs=(x, y, z)):
    ...

This reverts commit aa8fc42.

joblib/test/test_parallel.py

shwina · 2024-02-06T18:41:01Z

Hi @ogrisel - wondering if you have any thoughts on how to test the implementation here correctly?

bdice · 2024-06-24T18:09:51Z

Hi @ogrisel, this feature request has come up a couple times recently, including in a StackOverflow question: https://stackoverflow.com/questions/78642680/using-load-ext-cudf-pandas-throws-attributeerror

Would it be possible to work with a joblib developer to identify how to move forward with tests?

SJ-Innovation · 2024-07-15T09:27:06Z

Hi, just wondering is theres been any behind-the-scenes motion on this, as this would be a REALLY handy piece of functionality to have for pre-seeding workers with large, constant, but procedurally generated/loaded datasets at the start of a large number of jobs.

…for-loky-and-mp

codecov · 2025-02-27T23:43:46Z

Codecov Report

Attention: Patch coverage is 96.20253% with 3 lines in your changes missing coverage. Please review.

Project coverage is 95.50%. Comparing base (ad13c66) to head (0d9412d).
Report is 23 commits behind head on main.

Files with missing lines	Patch %	Lines
joblib/_parallel_backends.py	90.90%	1 Missing ⚠️
joblib/parallel.py	85.71%	1 Missing ⚠️
joblib/test/test_parallel.py	98.36%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1525      +/-   ##
==========================================
- Coverage   95.54%   95.50%   -0.04%     
==========================================
  Files          46       46              
  Lines        7741     7838      +97     
==========================================
+ Hits         7396     7486      +90     
- Misses        345      352       +7

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

joblib/test/test_parallel.py

joblib/_parallel_backends.py

ogrisel

Thanks for finalizing this PR. It appears that we have a problem with what we test, though.

CHANGES.rst

joblib/parallel.py

joblib/test/test_parallel.py

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>

ogrisel

LGTM. Hopefully the new pid-based assert statement will not randomly fail in case some workers take too long to start up compared to others.

If that happen we will have to rewrite those new tests to make them more robust to such race conditions.

ogrisel · 2025-04-07T16:00:08Z

Let's wait for the CI to go through and maybe trigger it a few consecutive times before merging.

tomMoral · 2025-04-07T19:53:00Z

This does not seem stable because the first worker fetches the second job between the time where the initializer has been run and the time where the second worker would take the job.

~~We can either say we fix this test to make it more complex and fix the race condition (putting a counter instead of True/False) or we test that the number of pid is large enough.~~
I fixed it to be more stable I think. I ran the tests three times, hopefully it is good now :)

lesteve · 2025-04-10T09:51:54Z

What is the expected behaviour when initializer raises an error? Maybe we can add test for it?

Quickly trying with loky you get TerminatedWorkerError and if you look before the stack-trace you get the stack-trace with the initializer error:

import math
from joblib import Parallel, delayed

def initializer(): 1/0

Parallel(n_jobs=2, initializer=initializer)(delayed(math.sqrt)(i*i) for i in range(10))

Details

Exception in initializer:
Traceback (most recent call last):
  File "/home/lesteve/dev/joblib/joblib/externals/loky/process_executor.py", line 409, in _process_worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
ZeroDivisionError: division by zero
Exception in initializer:
Traceback (most recent call last):
  File "/home/lesteve/dev/joblib/joblib/externals/loky/process_executor.py", line 409, in _process_worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
ZeroDivisionError: division by zero
---------------------------------------------------------------------------
TerminatedWorkerError                     Traceback (most recent call last)
Cell In[21], line 6
      2 from joblib import Parallel, delayed
      4 def initializer(): 1/0
----> 6 Parallel(n_jobs=2, initializer=initializer)(delayed(math.sqrt)(i*i) for i in range(10))

File ~/dev/joblib/joblib/parallel.py:2071, in Parallel.__call__(self, iterable)
   2065 # The first item from the output is blank, but it makes the interpreter
   2066 # progress until it enters the Try/Except block of the generator and
   2067 # reaches the first `yield` statement. This starts the asynchronous
   2068 # dispatch of the tasks to the workers.
   2069 next(output)
-> 2071 return output if self.return_generator else list(output)

File ~/dev/joblib/joblib/parallel.py:1681, in Parallel._get_outputs(self, iterator, pre_dispatch)
   1678     yield
   1680     with self._backend.retrieval_context():
-> 1681         yield from self._retrieve()
   1683 except GeneratorExit:
   1684     # The generator has been garbage collected before being fully
   1685     # consumed. This aborts the remaining tasks if possible and warn
   1686     # the user if necessary.
   1687     self._exception = True

File ~/dev/joblib/joblib/parallel.py:1783, in Parallel._retrieve(self)
   1777 while self._wait_retrieval():
   1778     # If the callback thread of a worker has signaled that its task
   1779     # triggered an exception, or if the retrieval loop has raised an
   1780     # exception (e.g. `GeneratorExit`), exit the loop and surface the
   1781     # worker traceback.
   1782     if self._aborting:
-> 1783         self._raise_error_fast()
   1784         break
   1786     nb_jobs = len(self._jobs)

File ~/dev/joblib/joblib/parallel.py:1858, in Parallel._raise_error_fast(self)
   1854 # If this error job exists, immediately raise the error by
   1855 # calling get_result. This job might not exists if abort has been
   1856 # called directly or if the generator is gc'ed.
   1857 if error_job is not None:
-> 1858     error_job.get_result(self.timeout)

File ~/dev/joblib/joblib/parallel.py:757, in BatchCompletionCallBack.get_result(self, timeout)
    751 backend = self.parallel._backend
    753 if backend.supports_retrieve_callback:
    754     # We assume that the result has already been retrieved by the
    755     # callback thread, and is stored internally. It's just waiting to
    756     # be returned.
--> 757     return self._return_or_raise()
    759 # For other backends, the main thread needs to run the retrieval step.
    760 try:

File ~/dev/joblib/joblib/parallel.py:772, in BatchCompletionCallBack._return_or_raise(self)
    770 try:
    771     if self.status == TASK_ERROR:
--> 772         raise self._result
    773     return self._result
    774 finally:

TerminatedWorkerError: A worker process managed by the executor was unexpectedly terminated. This could be caused by a segmentation fault while calling the function or by an excessive memory usage causing the Operating System to kill the worker.

The exit codes of the workers are {EXIT(0)}

With multiprocessing you get an infinite loop of error messages until you do Control-C (the behaviour is similar when using multiprocessing.Pool.map):

import math
from joblib import Parallel, delayed

def initializer(): 1/0

Parallel(n_jobs=2, backend='multiprocessing', initializer=initializer)(delayed(math.sqrt)(i*i) for i in range(10))

Details

Process ForkPoolWorker-16075:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16078:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16076:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16079:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-27-fd537094daee>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16081:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-27-fd537094daee>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16080:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
Process ForkPoolWorker-16082:
Traceback (most recent call last):
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/lesteve/micromamba/lib/python3.12/multiprocessing/pool.py", line 109, in worker
    initializer(*initargs)
  File "<ipython-input-21-a74918bfdf5d>", line 4, in initializer
    def initializer(): 1/0
                       ~^~
ZeroDivisionError: division by zero
.
.
.

tomMoral · 2025-04-11T11:32:56Z

For the behavior in case of errors, I think this is backend-specific and I would like to say joblib is not responsible 😅
This PR allows passing extra args to the backend, which are mapped to the specifics of the backend, so I would not test more in depth.

I propose to merge the PR as is, and we can always improve testing in a follow up PR if we feel this is necesary.

tomMoral · 2025-04-11T11:35:47Z

Thanks a lot @shwina, for the proposal and for your patience !!!

shwina added 6 commits November 18, 2023 06:36

Add an option to specify initializer function for processes

d9122f6

Add tests

e2832a8

Update docs

cb2478a

Don't be too prescriptive about how to use initializer

f86bfab

Move initializer definition and use batch_size=1

aa8fc42

Revert "Move initializer definition and use batch_size=1"

3cb1440

This reverts commit aa8fc42.

shwina commented Nov 20, 2023

View reviewed changes

joblib/test/test_parallel.py Outdated Show resolved Hide resolved

tomMoral added 3 commits February 28, 2025 00:14

Merge branch 'main' of github.com:joblib/joblib into add-initializer-…

a06e405

…for-loky-and-mp

ENH allow initializer implicitely

3e33540

ENH allow initializer in multiprocessing+CLN bad merge

1f9ab6c

tomMoral marked this pull request as ready for review February 27, 2025 23:37

tomMoral added 2 commits February 28, 2025 00:37

FIX linter

c095bcf

FIX code spell and typo

943110a

tomMoral added 3 commits February 28, 2025 01:00

FIX bad merge

6de81c4

FIX bad merge

49991e5

FIX make test more robust

8e398db

tomMoral approved these changes Feb 28, 2025

View reviewed changes

joblib/test/test_parallel.py Outdated Show resolved Hide resolved

joblib/_parallel_backends.py Show resolved Hide resolved

tomMoral added 8 commits February 28, 2025 01:28

DOC add entry in changelog

509437f

Merge branch 'main' into add-initializer-for-loky-and-mp

5e18921

FIX one test

9f57de3

FIX linter

380e853

TST make test_initializer more robust

a9ee170

TST make test_initializer_reuse more robust too

8d13fc6

FIX linter

733893a

TST quickend test_initializer

72bee0d

tomMoral mentioned this pull request Apr 3, 2025

No simple way to pass initializer for process #381

Closed

ENH further improve loky backend args

b746448

ogrisel reviewed Apr 4, 2025

View reviewed changes

CHANGES.rst Outdated Show resolved Hide resolved

joblib/parallel.py Outdated Show resolved Hide resolved

joblib/test/test_parallel.py Outdated Show resolved Hide resolved

tomMoral and others added 3 commits April 7, 2025 11:40

Update joblib/test/test_parallel.py

dbb5ca8

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>

Apply suggestions from code review

5bb8f66

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>

TST make test_initializer_reuse_* check worker reuse

9039d82

ogrisel approved these changes Apr 7, 2025

View reviewed changes

FIx linter

54897ba

tomMoral added 3 commits April 7, 2025 22:25

FIX more stable initializer tests

6ad8287

CI trigger

1107003

CI trigger

0d9412d

lesteve mentioned this pull request Apr 10, 2025

[BUG] Error message Issue: Joblib loky Backend Fails to Inherit Module Modifications Outside if __name__ == "__main__" #1705

Open

tomMoral merged commit 3aa56cc into joblib:main Apr 11, 2025
29 checks passed

tomMoral mentioned this pull request Apr 30, 2025

parallel_config(backend='loky', idle_worker_timeout=1.0) fails #1717

Closed

Conversation

shwina commented Nov 20, 2023

Uh oh!

Uh oh!

shwina commented Feb 6, 2024

Uh oh!

bdice commented Jun 24, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SJ-Innovation commented Jul 15, 2024

Uh oh!

codecov bot commented Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

ogrisel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ogrisel left a comment

Choose a reason for hiding this comment

Uh oh!

ogrisel commented Apr 7, 2025

Uh oh!

tomMoral commented Apr 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lesteve commented Apr 10, 2025

Uh oh!

tomMoral commented Apr 11, 2025

Uh oh!

Uh oh!

tomMoral commented Apr 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

bdice commented Jun 24, 2024 •

edited

Loading

codecov bot commented Feb 27, 2025 •

edited

Loading

tomMoral commented Apr 7, 2025 •

edited

Loading