Add option to specify an initialization function for 'loky' and 'multiprocessing' backends

CHANGES.rst

@@ -34,6 +34,10 @@ Parallel:
 - Pretty printing of ``Parallel`` execution progress when the number of tasks is
   known. https://github.com/joblib/joblib/pull/1608
 
+- Make it possible to pass extra arguments to the ``LokyBackend`` and
+  ``MultiprocessingBackend``, enabling the use of ``initializer``.
+  https://github.com/joblib/joblib/pull/1525
+
 - Refactor and document the custom parallel backend API.
   https://github.com/joblib/joblib/pull/1667
 

doc/parallel.rst

@@ -129,8 +129,20 @@ The context manager allow to set various backend parameters:
   of the backend. See :ref:`auto_memmapping_doc` for more details.
 
 Extra arguments passed in the :func:`~joblib.parallel_config` context are
-passed to the backend constructor, giving the possibility to set up the
-backend with additional parameters.
+passed to the backend constructor, allowing additional parameters to be set up:
+
+- ``LokyBackend``
+    - ``initializer``, ``initargs``: setup function and its arguments to call
+      in each worker process.
+    - ``idle_worker_timeout``: timeout in seconds for a worker to wait
+      for a new task before being clean up. The default is `300 s`.
+- ``MultprocessingBackend``
+    - ``initializer``, ``initargs``: setup function and its arguments to call
+      in each worker process.
+    - ``maxtasksperchild``: maximum number of tasks a worker can
+      execute before being replace with a fresh one.
+    - ``context``: specify the start method to use for creating new worker
+      processes.
 
 In addition to the builtin joblib backends, there are several cluster-specific
 backends you can use:

joblib/_parallel_backends.py

@@ -48,10 +48,13 @@
 
     nesting_level = None
 
-    def __init__(self, nesting_level=None, inner_max_num_threads=None, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self, nesting_level=None, inner_max_num_threads=None, **backend_kwargs
+    ):
+        super().__init__()
         self.nesting_level = nesting_level
         self.inner_max_num_threads = inner_max_num_threads
+        self.backend_kwargs = backend_kwargs
 
     MAX_NUM_THREADS_VARS = [
         "OMP_NUM_THREADS",
@@ -152,12 +155,7 @@
             return out.get()
 
     def configure(
-        self,
-        n_jobs=1,
-        parallel=None,
-        prefer=None,
-        require=None,
-        **backend_args,
+        self, n_jobs=1, parallel=None, prefer=None, require=None, **backend_kwargs
     ):
         """Reconfigure the backend and return the number of workers.
 
@@ -360,7 +358,7 @@
             self.configure(
                 n_jobs=self.parallel.n_jobs,
                 parallel=self.parallel,
-                **self.parallel._backend_args,
+                **self.parallel._backend_kwargs,
             )
 
 
@@ -489,7 +487,7 @@
     uses_threads = True
     supports_sharedmem = True
 
-    def configure(self, n_jobs=1, parallel=None, **backend_args):
+    def configure(self, n_jobs=1, parallel=None, **backend_kwargs):
         """Build a process or thread pool and return the number of workers"""
         n_jobs = self.effective_n_jobs(n_jobs)
         if n_jobs == 1:
@@ -574,16 +572,26 @@
         return super(MultiprocessingBackend, self).effective_n_jobs(n_jobs)
 
     def configure(
-        self, n_jobs=1, parallel=None, prefer=None, require=None, **memmappingpool_args
+        self,
+        n_jobs=1,
+        parallel=None,
+        prefer=None,
+        require=None,
+        **memmapping_pool_kwargs,
     ):
         """Build a process or thread pool and return the number of workers"""
         n_jobs = self.effective_n_jobs(n_jobs)
         if n_jobs == 1:
             raise FallbackToBackend(SequentialBackend(nesting_level=self.nesting_level))
 
+        memmapping_pool_kwargs = {
+            **self.backend_kwargs,
+            **memmapping_pool_kwargs,
+        }
+
         # Make sure to free as much memory as possible before forking
         gc.collect()
-        self._pool = MemmappingPool(n_jobs, **memmappingpool_args)
+        self._pool = MemmappingPool(n_jobs, **memmapping_pool_kwargs)
         self.parallel = parallel
         return n_jobs
 
@@ -605,20 +613,35 @@
         parallel=None,
         prefer=None,
         require=None,
-        idle_worker_timeout=300,
-        **memmappingexecutor_args,
+        idle_worker_timeout=None,
+        **memmapping_executor_kwargs,
     ):
         """Build a process executor and return the number of workers"""
         n_jobs = self.effective_n_jobs(n_jobs)
         if n_jobs == 1:
             raise FallbackToBackend(SequentialBackend(nesting_level=self.nesting_level))
 
+        memmapping_executor_kwargs = {
+            **self.backend_kwargs,
+            **memmapping_executor_kwargs,
+        }
+
+        # Prohibit the use of 'timeout' in the LokyBackend, as 'idle_worker_timeout'
+        # better describes the backend's behavior.
+        if "timeout" in memmapping_executor_kwargs:
+            raise ValueError(
+                "The 'timeout' parameter is not supported by the LokyBackend. "
+                "Please use the `idle_worker_timeout` parameter instead."
+            )
+        if idle_worker_timeout is None:
+            idle_worker_timeout = self.backend_kwargs.get("idle_worker_timeout", 300)
+
         self._workers = get_memmapping_executor(
             n_jobs,
             timeout=idle_worker_timeout,
             env=self._prepare_worker_env(n_jobs=n_jobs),
             context_id=parallel._id,
-            **memmappingexecutor_args,
+            **memmapping_executor_kwargs,
         )
         self.parallel = parallel
         return n_jobs

joblib/parallel.py

@@ -175,7 +175,6 @@
         # context manager or the context manager did not set a backend.
         # create the default backend instance now.
         backend = BACKENDS[DEFAULT_BACKEND](nesting_level=0)
-
         explicit_backend = False
 
     # Try to use the backend set by the user with the context manager.
@@ -300,7 +299,7 @@
           overridden with ``TMP``, ``TMPDIR`` or ``TEMP`` environment
           variables, typically ``/tmp`` under Unix operating systems.
 
-    max_nbytes int, str, or None, optional, default='1M'
+    max_nbytes: int, str, or None, optional, default='1M'
         Threshold on the size of arrays passed to the workers that
         triggers automated memory mapping in temp_folder. Can be an int
         in Bytes, or a human-readable string, e.g., '1M' for 1 megabyte.
@@ -542,8 +541,7 @@
 
     See Also
     --------
-    joblib.parallel_config: context manager to change the backend
-        configuration.
+    joblib.parallel_config: context manager to change the backend configuration.
     """
 
     def __init__(
@@ -1084,6 +1082,8 @@
         disable memmapping, other modes defined in the numpy.memmap doc:
         https://numpy.org/doc/stable/reference/generated/numpy.memmap.html
         Also, see 'max_nbytes' parameter documentation for more details.
+    backend_kwargs: dict, optional
+        Additional parameters to pass to the backend `configure` method.
 
     Notes
     -----
@@ -1222,6 +1222,7 @@
         mmap_mode=default_parallel_config["mmap_mode"],
         prefer=default_parallel_config["prefer"],
         require=default_parallel_config["require"],
+        **backend_kwargs,
     ):
         # Initiate parent Logger class state
         super().__init__()
@@ -1253,28 +1254,31 @@
         # Check if we are under a parallel_config or parallel_backend
         # context manager and use the config from the context manager
         # for arguments that are not explicitly set.
-        self._backend_args = {
-            k: _get_config_param(param, context_config, k)
-            for param, k in [
-                (max_nbytes, "max_nbytes"),
-                (temp_folder, "temp_folder"),
-                (mmap_mode, "mmap_mode"),
-                (prefer, "prefer"),
-                (require, "require"),
-                (verbose, "verbose"),
-            ]
+        self._backend_kwargs = {
+            **backend_kwargs,
+            **{
+                k: _get_config_param(param, context_config, k)
+                for param, k in [
+                    (max_nbytes, "max_nbytes"),
+                    (temp_folder, "temp_folder"),
+                    (mmap_mode, "mmap_mode"),
+                    (prefer, "prefer"),
+                    (require, "require"),
+                    (verbose, "verbose"),
+                ]
+            },
         }
 
-        if isinstance(self._backend_args["max_nbytes"], str):
-            self._backend_args["max_nbytes"] = memstr_to_bytes(
-                self._backend_args["max_nbytes"]
+        if isinstance(self._backend_kwargs["max_nbytes"], str):
+            self._backend_kwargs["max_nbytes"] = memstr_to_bytes(
+                self._backend_kwargs["max_nbytes"]
             )
-        self._backend_args["verbose"] = max(0, self._backend_args["verbose"] - 50)
+        self._backend_kwargs["verbose"] = max(0, self._backend_kwargs["verbose"] - 50)
 
         if DEFAULT_MP_CONTEXT is not None:
-            self._backend_args["context"] = DEFAULT_MP_CONTEXT
+            self._backend_kwargs["context"] = DEFAULT_MP_CONTEXT
         elif hasattr(mp, "get_context"):
-            self._backend_args["context"] = mp.get_context()
+            self._backend_kwargs["context"] = mp.get_context()
 
         if backend is default_parallel_config["backend"] or backend is None:
             backend = active_backend
@@ -1289,7 +1293,7 @@
             # Make it possible to pass a custom multiprocessing context as
             # backend to change the start method to forkserver or spawn or
             # preload modules on the forkserver helper process.
-            self._backend_args["context"] = backend
+            self._backend_kwargs["context"] = backend
             backend = MultiprocessingBackend(nesting_level=nesting_level)
 
         elif backend not in BACKENDS and backend in MAYBE_AVAILABLE_BACKENDS:
@@ -1372,7 +1376,7 @@
         """Build a process or thread pool and return the number of workers"""
         try:
             n_jobs = self._backend.configure(
-                n_jobs=self.n_jobs, parallel=self, **self._backend_args
+                n_jobs=self.n_jobs, parallel=self, **self._backend_kwargs
             )
             if self.timeout is not None and not self._backend.supports_timeout:
                 warnings.warn(