This commit is contained in:
Iliyan Angelov
2025-12-01 06:50:10 +02:00
parent 91f51bc6fe
commit 62c1fe5951
4682 changed files with 544807 additions and 31208 deletions

View File

@@ -0,0 +1,45 @@
r"""The :mod:`loky` module manages a pool of worker that can be re-used across time.
It provides a robust and dynamic implementation os the
:class:`ProcessPoolExecutor` and a function :func:`get_reusable_executor` which
hide the pool management under the hood.
"""
from concurrent.futures import (
ALL_COMPLETED,
FIRST_COMPLETED,
FIRST_EXCEPTION,
CancelledError,
Executor,
TimeoutError,
as_completed,
wait,
)
from ._base import Future
from .backend.context import cpu_count
from .backend.reduction import set_loky_pickler
from .reusable_executor import get_reusable_executor
from .cloudpickle_wrapper import wrap_non_picklable_objects
from .process_executor import BrokenProcessPool, ProcessPoolExecutor
__all__ = [
"get_reusable_executor",
"cpu_count",
"wait",
"as_completed",
"Future",
"Executor",
"ProcessPoolExecutor",
"BrokenProcessPool",
"CancelledError",
"TimeoutError",
"FIRST_COMPLETED",
"FIRST_EXCEPTION",
"ALL_COMPLETED",
"wrap_non_picklable_objects",
"set_loky_pickler",
]
__version__ = "3.5.6"

View File

@@ -0,0 +1,28 @@
###############################################################################
# Modification of concurrent.futures.Future
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from concurrent/futures/_base.py (17/02/2017)
# * Do not use yield from
# * Use old super syntax
#
# Copyright 2009 Brian Quinlan. All Rights Reserved.
# Licensed to PSF under a Contributor Agreement.
from concurrent.futures import Future as _BaseFuture
from concurrent.futures._base import LOGGER
# To make loky._base.Future instances awaitable by concurrent.futures.wait,
# derive our custom Future class from _BaseFuture. _invoke_callback is the only
# modification made to this class in loky.
# TODO investigate why using `concurrent.futures.Future` directly does not
# always work in our test suite.
class Future(_BaseFuture):
def _invoke_callbacks(self):
for callback in self._done_callbacks:
try:
callback(self)
except BaseException:
LOGGER.exception(f"exception calling callback for {self!r}")

View File

@@ -0,0 +1,14 @@
import os
from multiprocessing import synchronize
from .context import get_context
def _make_name():
return f"/loky-{os.getpid()}-{next(synchronize.SemLock._rand)}"
# monkey patch the name creation for multiprocessing
synchronize.SemLock._make_name = staticmethod(_make_name)
__all__ = ["get_context"]

View File

@@ -0,0 +1,67 @@
###############################################################################
# Extra reducers for Unix based system and connections objects
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/reduction.py (17/02/2017)
# * Add adapted reduction for LokyProcesses and socket/Connection
#
import os
import socket
import _socket
from multiprocessing.connection import Connection
from multiprocessing.context import get_spawning_popen
from .reduction import register
HAVE_SEND_HANDLE = (
hasattr(socket, "CMSG_LEN")
and hasattr(socket, "SCM_RIGHTS")
and hasattr(socket.socket, "sendmsg")
)
def _mk_inheritable(fd):
os.set_inheritable(fd, True)
return fd
def DupFd(fd):
"""Return a wrapper for an fd."""
popen_obj = get_spawning_popen()
if popen_obj is not None:
return popen_obj.DupFd(popen_obj.duplicate_for_child(fd))
elif HAVE_SEND_HANDLE:
from multiprocessing import resource_sharer
return resource_sharer.DupFd(fd)
else:
raise TypeError(
"Cannot pickle connection object. This object can only be "
"passed when spawning a new process"
)
def _reduce_socket(s):
df = DupFd(s.fileno())
return _rebuild_socket, (df, s.family, s.type, s.proto)
def _rebuild_socket(df, family, type, proto):
fd = df.detach()
return socket.fromfd(fd, family, type, proto)
def rebuild_connection(df, readable, writable):
fd = df.detach()
return Connection(fd, readable, writable)
def reduce_connection(conn):
df = DupFd(conn.fileno())
return rebuild_connection, (df, conn.readable, conn.writable)
register(socket.socket, _reduce_socket)
register(_socket.socket, _reduce_socket)
register(Connection, reduce_connection)

View File

@@ -0,0 +1,18 @@
###############################################################################
# Extra reducers for Windows system and connections objects
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/reduction.py (17/02/2017)
# * Add adapted reduction for LokyProcesses and socket/PipeConnection
#
import socket
from multiprocessing import connection
from multiprocessing.reduction import _reduce_socket
from .reduction import register
# register reduction for win32 communication objects
register(socket.socket, _reduce_socket)
register(connection.Connection, connection.reduce_connection)
register(connection.PipeConnection, connection.reduce_pipe_connection)

View File

@@ -0,0 +1,405 @@
###############################################################################
# Basic context management with LokyContext
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/context.py
# * Create a context ensuring loky uses only objects that are compatible
# * Add LokyContext to the list of context of multiprocessing so loky can be
# used with multiprocessing.set_start_method
# * Implement a CFS-aware amd physical-core aware cpu_count function.
#
import os
import sys
import math
import subprocess
import traceback
import warnings
import multiprocessing as mp
from multiprocessing import get_context as mp_get_context
from multiprocessing.context import BaseContext
from concurrent.futures.process import _MAX_WINDOWS_WORKERS
from .process import LokyProcess, LokyInitMainProcess
# Apparently, on older Python versions, loky cannot work 61 workers on Windows
# but instead 60: ¯\_(ツ)_/¯
if sys.version_info < (3, 10):
_MAX_WINDOWS_WORKERS = _MAX_WINDOWS_WORKERS - 1
START_METHODS = ["loky", "loky_init_main", "spawn"]
if sys.platform != "win32":
START_METHODS += ["fork", "forkserver"]
_DEFAULT_START_METHOD = None
# Cache for the number of physical cores to avoid repeating subprocess calls.
# It should not change during the lifetime of the program.
physical_cores_cache = None
def get_context(method=None):
# Try to overload the default context
method = method or _DEFAULT_START_METHOD or "loky"
if method == "fork":
# If 'fork' is explicitly requested, warn user about potential issues.
warnings.warn(
"`fork` start method should not be used with "
"`loky` as it does not respect POSIX. Try using "
"`spawn` or `loky` instead.",
UserWarning,
)
try:
return mp_get_context(method)
except ValueError:
raise ValueError(
f"Unknown context '{method}'. Value should be in "
f"{START_METHODS}."
)
def set_start_method(method, force=False):
global _DEFAULT_START_METHOD
if _DEFAULT_START_METHOD is not None and not force:
raise RuntimeError("context has already been set")
assert method is None or method in START_METHODS, (
f"'{method}' is not a valid start_method. It should be in "
f"{START_METHODS}"
)
_DEFAULT_START_METHOD = method
def get_start_method():
return _DEFAULT_START_METHOD
def cpu_count(only_physical_cores=False):
"""Return the number of CPUs the current process can use.
The returned number of CPUs accounts for:
* the number of CPUs in the system, as given by
``multiprocessing.cpu_count``;
* the CPU affinity settings of the current process
(available on some Unix systems);
* Cgroup CPU bandwidth limit (available on Linux only, typically
set by docker and similar container orchestration systems);
* the value of the LOKY_MAX_CPU_COUNT environment variable if defined.
and is given as the minimum of these constraints.
If ``only_physical_cores`` is True, return the number of physical cores
instead of the number of logical cores (hyperthreading / SMT). Note that
this option is not enforced if the number of usable cores is controlled in
any other way such as: process affinity, Cgroup restricted CPU bandwidth
or the LOKY_MAX_CPU_COUNT environment variable. If the number of physical
cores is not found, return the number of logical cores.
Note that on Windows, the returned number of CPUs cannot exceed 61 (or 60 for
Python < 3.10), see:
https://bugs.python.org/issue26903.
It is also always larger or equal to 1.
"""
# Note: os.cpu_count() is allowed to return None in its docstring
os_cpu_count = os.cpu_count() or 1
if sys.platform == "win32":
# On Windows, attempting to use more than 61 CPUs would result in a
# OS-level error. See https://bugs.python.org/issue26903. According to
# https://learn.microsoft.com/en-us/windows/win32/procthread/processor-groups
# it might be possible to go beyond with a lot of extra work but this
# does not look easy.
os_cpu_count = min(os_cpu_count, _MAX_WINDOWS_WORKERS)
cpu_count_user = _cpu_count_user(os_cpu_count)
aggregate_cpu_count = max(min(os_cpu_count, cpu_count_user), 1)
if not only_physical_cores:
return aggregate_cpu_count
if cpu_count_user < os_cpu_count:
# Respect user setting
return max(cpu_count_user, 1)
cpu_count_physical, exception = _count_physical_cores()
if cpu_count_physical != "not found":
return cpu_count_physical
# Fallback to default behavior
if exception is not None:
# warns only the first time
warnings.warn(
"Could not find the number of physical cores for the "
f"following reason:\n{exception}\n"
"Returning the number of logical cores instead. You can "
"silence this warning by setting LOKY_MAX_CPU_COUNT to "
"the number of cores you want to use."
)
traceback.print_tb(exception.__traceback__)
return aggregate_cpu_count
def _cpu_count_cgroup(os_cpu_count):
# Cgroup CPU bandwidth limit available in Linux since 2.6 kernel
cpu_max_fname = "/sys/fs/cgroup/cpu.max"
cfs_quota_fname = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
cfs_period_fname = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
if os.path.exists(cpu_max_fname):
# cgroup v2
# https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html
with open(cpu_max_fname) as fh:
cpu_quota_us, cpu_period_us = fh.read().strip().split()
elif os.path.exists(cfs_quota_fname) and os.path.exists(cfs_period_fname):
# cgroup v1
# https://www.kernel.org/doc/html/latest/scheduler/sched-bwc.html#management
with open(cfs_quota_fname) as fh:
cpu_quota_us = fh.read().strip()
with open(cfs_period_fname) as fh:
cpu_period_us = fh.read().strip()
else:
# No Cgroup CPU bandwidth limit (e.g. non-Linux platform)
cpu_quota_us = "max"
cpu_period_us = 100_000 # unused, for consistency with default values
if cpu_quota_us == "max":
# No active Cgroup quota on a Cgroup-capable platform
return os_cpu_count
else:
cpu_quota_us = int(cpu_quota_us)
cpu_period_us = int(cpu_period_us)
if cpu_quota_us > 0 and cpu_period_us > 0:
return math.ceil(cpu_quota_us / cpu_period_us)
else: # pragma: no cover
# Setting a negative cpu_quota_us value is a valid way to disable
# cgroup CPU bandwith limits
return os_cpu_count
def _cpu_count_affinity(os_cpu_count):
# Number of available CPUs given affinity settings
if hasattr(os, "sched_getaffinity"):
try:
return len(os.sched_getaffinity(0))
except NotImplementedError:
pass
# On some platforms, os.sched_getaffinity does not exist or raises
# NotImplementedError, let's try with the psutil if installed.
try:
import psutil
p = psutil.Process()
if hasattr(p, "cpu_affinity"):
return len(p.cpu_affinity())
except ImportError: # pragma: no cover
if (
sys.platform == "linux"
and os.environ.get("LOKY_MAX_CPU_COUNT") is None
):
# Some platforms don't implement os.sched_getaffinity on Linux which
# can cause severe oversubscription problems. Better warn the
# user in this particularly pathological case which can wreck
# havoc, typically on CI workers.
warnings.warn(
"Failed to inspect CPU affinity constraints on this system. "
"Please install psutil or explictly set LOKY_MAX_CPU_COUNT."
)
# This can happen for platforms that do not implement any kind of CPU
# infinity such as macOS-based platforms.
return os_cpu_count
def _cpu_count_user(os_cpu_count):
"""Number of user defined available CPUs"""
cpu_count_affinity = _cpu_count_affinity(os_cpu_count)
cpu_count_cgroup = _cpu_count_cgroup(os_cpu_count)
# User defined soft-limit passed as a loky specific environment variable.
cpu_count_loky = int(os.environ.get("LOKY_MAX_CPU_COUNT", os_cpu_count))
return min(cpu_count_affinity, cpu_count_cgroup, cpu_count_loky)
def _count_physical_cores():
"""Return a tuple (number of physical cores, exception)
If the number of physical cores is found, exception is set to None.
If it has not been found, return ("not found", exception).
The number of physical cores is cached to avoid repeating subprocess calls.
"""
exception = None
# First check if the value is cached
global physical_cores_cache
if physical_cores_cache is not None:
return physical_cores_cache, exception
# Not cached yet, find it
try:
if sys.platform == "linux":
cpu_count_physical = _count_physical_cores_linux()
elif sys.platform == "win32":
cpu_count_physical = _count_physical_cores_win32()
elif sys.platform == "darwin":
cpu_count_physical = _count_physical_cores_darwin()
else:
raise NotImplementedError(f"unsupported platform: {sys.platform}")
# if cpu_count_physical < 1, we did not find a valid value
if cpu_count_physical < 1:
raise ValueError(f"found {cpu_count_physical} physical cores < 1")
except Exception as e:
exception = e
cpu_count_physical = "not found"
# Put the result in cache
physical_cores_cache = cpu_count_physical
return cpu_count_physical, exception
def _count_physical_cores_linux():
try:
cpu_info = subprocess.run(
"lscpu --parse=core".split(), capture_output=True, text=True
)
cpu_info = cpu_info.stdout.splitlines()
cpu_info = {line for line in cpu_info if not line.startswith("#")}
return len(cpu_info)
except:
pass # fallback to /proc/cpuinfo
cpu_info = subprocess.run(
"cat /proc/cpuinfo".split(), capture_output=True, text=True
)
cpu_info = cpu_info.stdout.splitlines()
cpu_info = {line for line in cpu_info if line.startswith("core id")}
return len(cpu_info)
def _count_physical_cores_win32():
try:
cmd = "-Command (Get-CimInstance -ClassName Win32_Processor).NumberOfCores"
cpu_info = subprocess.run(
f"powershell.exe {cmd}".split(),
capture_output=True,
text=True,
)
cpu_info = cpu_info.stdout.splitlines()
return int(cpu_info[0])
except:
pass # fallback to wmic (older Windows versions; deprecated now)
cpu_info = subprocess.run(
"wmic CPU Get NumberOfCores /Format:csv".split(),
capture_output=True,
text=True,
)
cpu_info = cpu_info.stdout.splitlines()
cpu_info = [
l.split(",")[1] for l in cpu_info if (l and l != "Node,NumberOfCores")
]
return sum(map(int, cpu_info))
def _count_physical_cores_darwin():
cpu_info = subprocess.run(
"sysctl -n hw.physicalcpu".split(),
capture_output=True,
text=True,
)
cpu_info = cpu_info.stdout
return int(cpu_info)
class LokyContext(BaseContext):
"""Context relying on the LokyProcess."""
_name = "loky"
Process = LokyProcess
cpu_count = staticmethod(cpu_count)
def Queue(self, maxsize=0, reducers=None):
"""Returns a queue object"""
from .queues import Queue
return Queue(maxsize, reducers=reducers, ctx=self.get_context())
def SimpleQueue(self, reducers=None):
"""Returns a queue object"""
from .queues import SimpleQueue
return SimpleQueue(reducers=reducers, ctx=self.get_context())
if sys.platform != "win32":
"""For Unix platform, use our custom implementation of synchronize
ensuring that we use the loky.backend.resource_tracker to clean-up
the semaphores in case of a worker crash.
"""
def Semaphore(self, value=1):
"""Returns a semaphore object"""
from .synchronize import Semaphore
return Semaphore(value=value)
def BoundedSemaphore(self, value):
"""Returns a bounded semaphore object"""
from .synchronize import BoundedSemaphore
return BoundedSemaphore(value)
def Lock(self):
"""Returns a lock object"""
from .synchronize import Lock
return Lock()
def RLock(self):
"""Returns a recurrent lock object"""
from .synchronize import RLock
return RLock()
def Condition(self, lock=None):
"""Returns a condition object"""
from .synchronize import Condition
return Condition(lock)
def Event(self):
"""Returns an event object"""
from .synchronize import Event
return Event()
class LokyInitMainContext(LokyContext):
"""Extra context with LokyProcess, which does load the main module
This context is used for compatibility in the case ``cloudpickle`` is not
present on the running system. This permits to load functions defined in
the ``main`` module, using proper safeguards. The declaration of the
``executor`` should be protected by ``if __name__ == "__main__":`` and the
functions and variable used from main should be out of this block.
This mimics the default behavior of multiprocessing under Windows and the
behavior of the ``spawn`` start method on a posix system.
For more details, see the end of the following section of python doc
https://docs.python.org/3/library/multiprocessing.html#multiprocessing-programming
"""
_name = "loky_init_main"
Process = LokyInitMainProcess
# Register loky context so it works with multiprocessing.get_context
ctx_loky = LokyContext()
mp.context._concrete_contexts["loky"] = ctx_loky
mp.context._concrete_contexts["loky_init_main"] = LokyInitMainContext()

View File

@@ -0,0 +1,73 @@
###############################################################################
# Launch a subprocess using forkexec and make sure only the needed fd are
# shared in the two process.
#
# author: Thomas Moreau and Olivier Grisel
#
import sys
import os
import subprocess
def fork_exec(cmd, keep_fds, env=None):
import _posixsubprocess
# Encoded command args as bytes:
cmd = [os.fsencode(arg) for arg in cmd]
# Copy the environment variables to set in the child process (also encoded
# as bytes).
env = env or {}
env = {**os.environ, **env}
encoded_env = []
for key, value in env.items():
encoded_env.append(os.fsencode(f"{key}={value}"))
# Fds with fileno larger than 3 (stdin=0, stdout=1, stderr=2) are be closed
# in the child process, except for those passed in keep_fds.
keep_fds = tuple(sorted(map(int, keep_fds)))
errpipe_read, errpipe_write = os.pipe()
if sys.version_info >= (3, 14):
# Python >= 3.14 removed allow_vfork from _posixsubprocess.fork_exec,
# see https://github.com/python/cpython/pull/121383
pgid_to_set = [-1]
allow_vfork = []
elif sys.version_info >= (3, 11):
# Python 3.11 - 3.13 has allow_vfork in _posixsubprocess.fork_exec
pgid_to_set = [-1]
allow_vfork = [subprocess._USE_VFORK]
else:
# Python < 3.11
pgid_to_set = []
allow_vfork = []
try:
return _posixsubprocess.fork_exec(
cmd, # args
cmd[0:1], # executable_list
True, # close_fds
keep_fds, # pass_fds
None, # cwd
encoded_env, # env
-1, # p2cread
-1, # p2cwrite
-1, # c2pread
-1, # c2pwrite
-1, # errread
-1, # errwrite
errpipe_read, # errpipe_read
errpipe_write, # errpipe_write
False, # restore_signal
False, # call_setsid
*pgid_to_set, # pgid_to_set
None, # gid
None, # extra_groups
None, # uid
-1, # child_umask
None, # preexec_fn
*allow_vfork, # extra flag if vfork is available
)
finally:
os.close(errpipe_read)
os.close(errpipe_write)

View File

@@ -0,0 +1,193 @@
###############################################################################
# Popen for LokyProcess.
#
# author: Thomas Moreau and Olivier Grisel
#
import os
import sys
import signal
import pickle
from io import BytesIO
from multiprocessing import util, process
from multiprocessing.connection import wait
from multiprocessing.context import set_spawning_popen
from . import reduction, resource_tracker, spawn
__all__ = ["Popen"]
#
# Wrapper for an fd used while launching a process
#
class _DupFd:
def __init__(self, fd):
self.fd = reduction._mk_inheritable(fd)
def detach(self):
return self.fd
#
# Start child process using subprocess.Popen
#
class Popen:
method = "loky"
DupFd = _DupFd
def __init__(self, process_obj):
sys.stdout.flush()
sys.stderr.flush()
self.returncode = None
self._fds = []
self._launch(process_obj)
def duplicate_for_child(self, fd):
self._fds.append(fd)
return reduction._mk_inheritable(fd)
def poll(self, flag=os.WNOHANG):
if self.returncode is None:
while True:
try:
pid, sts = os.waitpid(self.pid, flag)
except OSError:
# Child process not yet created. See #1731717
# e.errno == errno.ECHILD == 10
return None
else:
break
if pid == self.pid:
if os.WIFSIGNALED(sts):
self.returncode = -os.WTERMSIG(sts)
else:
assert os.WIFEXITED(sts)
self.returncode = os.WEXITSTATUS(sts)
return self.returncode
def wait(self, timeout=None):
if self.returncode is None:
if timeout is not None:
if not wait([self.sentinel], timeout):
return None
# This shouldn't block if wait() returned successfully.
return self.poll(os.WNOHANG if timeout == 0.0 else 0)
return self.returncode
def terminate(self):
if self.returncode is None:
try:
os.kill(self.pid, signal.SIGTERM)
except ProcessLookupError:
pass
except OSError:
if self.wait(timeout=0.1) is None:
raise
def _launch(self, process_obj):
tracker_fd = resource_tracker._resource_tracker.getfd()
fp = BytesIO()
set_spawning_popen(self)
try:
prep_data = spawn.get_preparation_data(
process_obj._name,
getattr(process_obj, "init_main_module", True),
)
reduction.dump(prep_data, fp)
reduction.dump(process_obj, fp)
finally:
set_spawning_popen(None)
try:
parent_r, child_w = os.pipe()
child_r, parent_w = os.pipe()
# for fd in self._fds:
# _mk_inheritable(fd)
cmd_python = [sys.executable]
cmd_python += ["-m", self.__module__]
cmd_python += ["--process-name", str(process_obj.name)]
cmd_python += ["--pipe", str(reduction._mk_inheritable(child_r))]
reduction._mk_inheritable(child_w)
reduction._mk_inheritable(tracker_fd)
self._fds += [child_r, child_w, tracker_fd]
if os.name == "posix":
mp_tracker_fd = prep_data["mp_tracker_fd"]
self.duplicate_for_child(mp_tracker_fd)
from .fork_exec import fork_exec
pid = fork_exec(cmd_python, self._fds, env=process_obj.env)
util.debug(
f"launched python with pid {pid} and cmd:\n{cmd_python}"
)
self.sentinel = parent_r
method = "getbuffer"
if not hasattr(fp, method):
method = "getvalue"
with os.fdopen(parent_w, "wb") as f:
f.write(getattr(fp, method)())
self.pid = pid
finally:
if parent_r is not None:
util.Finalize(self, os.close, (parent_r,))
for fd in (child_r, child_w):
if fd is not None:
os.close(fd)
@staticmethod
def thread_is_spawning():
return True
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser("Command line parser")
parser.add_argument(
"--pipe", type=int, required=True, help="File handle for the pipe"
)
parser.add_argument(
"--process-name",
type=str,
default=None,
help="Identifier for debugging purpose",
)
args = parser.parse_args()
info = {}
exitcode = 1
try:
with os.fdopen(args.pipe, "rb") as from_parent:
process.current_process()._inheriting = True
try:
prep_data = pickle.load(from_parent)
spawn.prepare(prep_data)
process_obj = pickle.load(from_parent)
finally:
del process.current_process()._inheriting
exitcode = process_obj._bootstrap()
except Exception:
print("\n\n" + "-" * 80)
print(f"{args.process_name} failed with traceback: ")
print("-" * 80)
import traceback
print(traceback.format_exc())
print("\n" + "-" * 80)
finally:
if from_parent is not None:
from_parent.close()
sys.exit(exitcode)

View File

@@ -0,0 +1,173 @@
import os
import sys
import msvcrt
import _winapi
from pickle import load
from multiprocessing import process, util
from multiprocessing.context import set_spawning_popen
from multiprocessing.popen_spawn_win32 import Popen as _Popen
from . import reduction, spawn
__all__ = ["Popen"]
#
#
#
def _path_eq(p1, p2):
return p1 == p2 or os.path.normcase(p1) == os.path.normcase(p2)
WINENV = hasattr(sys, "_base_executable") and not _path_eq(
sys.executable, sys._base_executable
)
def _close_handles(*handles):
for handle in handles:
_winapi.CloseHandle(handle)
#
# We define a Popen class similar to the one from subprocess, but
# whose constructor takes a process object as its argument.
#
class Popen(_Popen):
"""
Start a subprocess to run the code of a process object.
We differ from cpython implementation with the way we handle environment
variables, in order to be able to modify then in the child processes before
importing any library, in order to control the number of threads in C-level
threadpools.
We also use the loky preparation data, in particular to handle main_module
inits and the loky resource tracker.
"""
method = "loky"
def __init__(self, process_obj):
prep_data = spawn.get_preparation_data(
process_obj._name, getattr(process_obj, "init_main_module", True)
)
# read end of pipe will be duplicated by the child process
# -- see spawn_main() in spawn.py.
#
# bpo-33929: Previously, the read end of pipe was "stolen" by the child
# process, but it leaked a handle if the child process had been
# terminated before it could steal the handle from the parent process.
rhandle, whandle = _winapi.CreatePipe(None, 0)
wfd = msvcrt.open_osfhandle(whandle, 0)
cmd = get_command_line(parent_pid=os.getpid(), pipe_handle=rhandle)
python_exe = spawn.get_executable()
# copy the environment variables to set in the child process
child_env = {**os.environ, **process_obj.env}
# bpo-35797: When running in a venv, we bypass the redirect
# executor and launch our base Python.
if WINENV and _path_eq(python_exe, sys.executable):
cmd[0] = python_exe = sys._base_executable
child_env["__PYVENV_LAUNCHER__"] = sys.executable
cmd = " ".join(f'"{x}"' for x in cmd)
with open(wfd, "wb") as to_child:
# start process
try:
hp, ht, pid, _ = _winapi.CreateProcess(
python_exe,
cmd,
None,
None,
False,
0,
child_env,
None,
None,
)
_winapi.CloseHandle(ht)
except BaseException:
_winapi.CloseHandle(rhandle)
raise
# set attributes of self
self.pid = pid
self.returncode = None
self._handle = hp
self.sentinel = int(hp)
self.finalizer = util.Finalize(
self, _close_handles, (self.sentinel, int(rhandle))
)
# send information to child
set_spawning_popen(self)
try:
reduction.dump(prep_data, to_child)
reduction.dump(process_obj, to_child)
finally:
set_spawning_popen(None)
def get_command_line(pipe_handle, parent_pid, **kwds):
"""Returns prefix of command line used for spawning a child process."""
if getattr(sys, "frozen", False):
return [sys.executable, "--multiprocessing-fork", pipe_handle]
else:
prog = (
"from joblib.externals.loky.backend.popen_loky_win32 import main; "
f"main(pipe_handle={pipe_handle}, parent_pid={parent_pid})"
)
opts = util._args_from_interpreter_flags()
return [
spawn.get_executable(),
*opts,
"-c",
prog,
"--multiprocessing-fork",
]
def is_forking(argv):
"""Return whether commandline indicates we are forking."""
if len(argv) >= 2 and argv[1] == "--multiprocessing-fork":
return True
else:
return False
def main(pipe_handle, parent_pid=None):
"""Run code specified by data received over pipe."""
assert is_forking(sys.argv), "Not forking"
if parent_pid is not None:
source_process = _winapi.OpenProcess(
_winapi.SYNCHRONIZE | _winapi.PROCESS_DUP_HANDLE, False, parent_pid
)
else:
source_process = None
new_handle = reduction.duplicate(
pipe_handle, source_process=source_process
)
fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY)
parent_sentinel = source_process
with os.fdopen(fd, "rb", closefd=True) as from_parent:
process.current_process()._inheriting = True
try:
preparation_data = load(from_parent)
spawn.prepare(preparation_data, parent_sentinel)
self = load(from_parent)
finally:
del process.current_process()._inheriting
exitcode = self._bootstrap(parent_sentinel)
sys.exit(exitcode)

View File

@@ -0,0 +1,85 @@
###############################################################################
# LokyProcess implementation
#
# authors: Thomas Moreau and Olivier Grisel
#
# based on multiprocessing/process.py (17/02/2017)
#
import sys
from multiprocessing.context import assert_spawning
from multiprocessing.process import BaseProcess
class LokyProcess(BaseProcess):
_start_method = "loky"
def __init__(
self,
group=None,
target=None,
name=None,
args=(),
kwargs={},
daemon=None,
init_main_module=False,
env=None,
):
super().__init__(
group=group,
target=target,
name=name,
args=args,
kwargs=kwargs,
daemon=daemon,
)
self.env = {} if env is None else env
self.authkey = self.authkey
self.init_main_module = init_main_module
@staticmethod
def _Popen(process_obj):
if sys.platform == "win32":
from .popen_loky_win32 import Popen
else:
from .popen_loky_posix import Popen
return Popen(process_obj)
class LokyInitMainProcess(LokyProcess):
_start_method = "loky_init_main"
def __init__(
self,
group=None,
target=None,
name=None,
args=(),
kwargs={},
daemon=None,
):
super().__init__(
group=group,
target=target,
name=name,
args=args,
kwargs=kwargs,
daemon=daemon,
init_main_module=True,
)
#
# We subclass bytes to avoid accidental transmission of auth keys over network
#
class AuthenticationKey(bytes):
def __reduce__(self):
try:
assert_spawning(self)
except RuntimeError:
raise TypeError(
"Pickling an AuthenticationKey object is "
"disallowed for security reasons"
)
return AuthenticationKey, (bytes(self),)

View File

@@ -0,0 +1,236 @@
###############################################################################
# Queue and SimpleQueue implementation for loky
#
# authors: Thomas Moreau, Olivier Grisel
#
# based on multiprocessing/queues.py (16/02/2017)
# * Add some custom reducers for the Queues/SimpleQueue to tweak the
# pickling process. (overload Queue._feed/SimpleQueue.put)
#
import os
import sys
import errno
import weakref
import threading
from multiprocessing import util
from multiprocessing.queues import (
Full,
Queue as mp_Queue,
SimpleQueue as mp_SimpleQueue,
_sentinel,
)
from multiprocessing.context import assert_spawning
from .reduction import dumps
__all__ = ["Queue", "SimpleQueue", "Full"]
class Queue(mp_Queue):
def __init__(self, maxsize=0, reducers=None, ctx=None):
super().__init__(maxsize=maxsize, ctx=ctx)
self._reducers = reducers
# Use custom queue set/get state to be able to reduce the custom reducers
def __getstate__(self):
assert_spawning(self)
return (
self._ignore_epipe,
self._maxsize,
self._reader,
self._writer,
self._reducers,
self._rlock,
self._wlock,
self._sem,
self._opid,
)
def __setstate__(self, state):
(
self._ignore_epipe,
self._maxsize,
self._reader,
self._writer,
self._reducers,
self._rlock,
self._wlock,
self._sem,
self._opid,
) = state
if sys.version_info >= (3, 9):
self._reset()
else:
self._after_fork()
# Overload _start_thread to correctly call our custom _feed
def _start_thread(self):
util.debug("Queue._start_thread()")
# Start thread which transfers data from buffer to pipe
self._buffer.clear()
self._thread = threading.Thread(
target=Queue._feed,
args=(
self._buffer,
self._notempty,
self._send_bytes,
self._wlock,
self._writer.close,
self._reducers,
self._ignore_epipe,
self._on_queue_feeder_error,
self._sem,
),
name="QueueFeederThread",
)
self._thread.daemon = True
util.debug("doing self._thread.start()")
self._thread.start()
util.debug("... done self._thread.start()")
# On process exit we will wait for data to be flushed to pipe.
#
# However, if this process created the queue then all
# processes which use the queue will be descendants of this
# process. Therefore waiting for the queue to be flushed
# is pointless once all the child processes have been joined.
created_by_this_process = self._opid == os.getpid()
if not self._joincancelled and not created_by_this_process:
self._jointhread = util.Finalize(
self._thread,
Queue._finalize_join,
[weakref.ref(self._thread)],
exitpriority=-5,
)
# Send sentinel to the thread queue object when garbage collected
self._close = util.Finalize(
self,
Queue._finalize_close,
[self._buffer, self._notempty],
exitpriority=10,
)
# Overload the _feed methods to use our custom pickling strategy.
@staticmethod
def _feed(
buffer,
notempty,
send_bytes,
writelock,
close,
reducers,
ignore_epipe,
onerror,
queue_sem,
):
util.debug("starting thread to feed data to pipe")
nacquire = notempty.acquire
nrelease = notempty.release
nwait = notempty.wait
bpopleft = buffer.popleft
sentinel = _sentinel
if sys.platform != "win32":
wacquire = writelock.acquire
wrelease = writelock.release
else:
wacquire = None
while True:
try:
nacquire()
try:
if not buffer:
nwait()
finally:
nrelease()
try:
while True:
obj = bpopleft()
if obj is sentinel:
util.debug("feeder thread got sentinel -- exiting")
close()
return
# serialize the data before acquiring the lock
obj_ = dumps(obj, reducers=reducers)
if wacquire is None:
send_bytes(obj_)
else:
wacquire()
try:
send_bytes(obj_)
finally:
wrelease()
# Remove references early to avoid leaking memory
del obj, obj_
except IndexError:
pass
except BaseException as e:
if ignore_epipe and getattr(e, "errno", 0) == errno.EPIPE:
return
# Since this runs in a daemon thread the resources it uses
# may be become unusable while the process is cleaning up.
# We ignore errors which happen after the process has
# started to cleanup.
if util.is_exiting():
util.info(f"error in queue thread: {e}")
return
else:
queue_sem.release()
onerror(e, obj)
def _on_queue_feeder_error(self, e, obj):
"""
Private API hook called when feeding data in the background thread
raises an exception. For overriding by concurrent.futures.
"""
import traceback
traceback.print_exc()
class SimpleQueue(mp_SimpleQueue):
def __init__(self, reducers=None, ctx=None):
super().__init__(ctx=ctx)
# Add possiblity to use custom reducers
self._reducers = reducers
def close(self):
self._reader.close()
self._writer.close()
# Use custom queue set/get state to be able to reduce the custom reducers
def __getstate__(self):
assert_spawning(self)
return (
self._reader,
self._writer,
self._reducers,
self._rlock,
self._wlock,
)
def __setstate__(self, state):
(
self._reader,
self._writer,
self._reducers,
self._rlock,
self._wlock,
) = state
# Overload put to use our customizable reducer
def put(self, obj):
# serialize the data before acquiring the lock
obj = dumps(obj, reducers=self._reducers)
if self._wlock is None:
# writes to a message oriented win32 pipe are atomic
self._writer.send_bytes(obj)
else:
with self._wlock:
self._writer.send_bytes(obj)

View File

@@ -0,0 +1,223 @@
###############################################################################
# Customizable Pickler with some basic reducers
#
# author: Thomas Moreau
#
# adapted from multiprocessing/reduction.py (17/02/2017)
# * Replace the ForkingPickler with a similar _LokyPickler,
# * Add CustomizableLokyPickler to allow customizing pickling process
# on the fly.
#
import copyreg
import io
import functools
import types
import sys
import os
from multiprocessing import util
from pickle import loads, HIGHEST_PROTOCOL
###############################################################################
# Enable custom pickling in Loky.
_dispatch_table = {}
def register(type_, reduce_function):
_dispatch_table[type_] = reduce_function
###############################################################################
# Registers extra pickling routines to improve picklization for loky
# make methods picklable
def _reduce_method(m):
if m.__self__ is None:
return getattr, (m.__class__, m.__func__.__name__)
else:
return getattr, (m.__self__, m.__func__.__name__)
class _C:
def f(self):
pass
@classmethod
def h(cls):
pass
register(type(_C().f), _reduce_method)
register(type(_C.h), _reduce_method)
def _reduce_method_descriptor(m):
return getattr, (m.__objclass__, m.__name__)
register(type(list.append), _reduce_method_descriptor)
register(type(int.__add__), _reduce_method_descriptor)
# Make partial func pickable
def _reduce_partial(p):
return _rebuild_partial, (p.func, p.args, p.keywords or {})
def _rebuild_partial(func, args, keywords):
return functools.partial(func, *args, **keywords)
register(functools.partial, _reduce_partial)
if sys.platform != "win32":
from ._posix_reduction import _mk_inheritable # noqa: F401
else:
from . import _win_reduction # noqa: F401
# global variable to change the pickler behavior
try:
from joblib.externals import cloudpickle # noqa: F401
DEFAULT_ENV = "cloudpickle"
except ImportError:
# If cloudpickle is not present, fallback to pickle
DEFAULT_ENV = "pickle"
ENV_LOKY_PICKLER = os.environ.get("LOKY_PICKLER", DEFAULT_ENV)
_LokyPickler = None
_loky_pickler_name = None
def set_loky_pickler(loky_pickler=None):
global _LokyPickler, _loky_pickler_name
if loky_pickler is None:
loky_pickler = ENV_LOKY_PICKLER
loky_pickler_cls = None
# The default loky_pickler is cloudpickle
if loky_pickler in ["", None]:
loky_pickler = "cloudpickle"
if loky_pickler == _loky_pickler_name:
return
if loky_pickler == "cloudpickle":
from joblib.externals.cloudpickle import CloudPickler as loky_pickler_cls
else:
try:
from importlib import import_module
module_pickle = import_module(loky_pickler)
loky_pickler_cls = module_pickle.Pickler
except (ImportError, AttributeError) as e:
extra_info = (
"\nThis error occurred while setting loky_pickler to"
f" '{loky_pickler}', as required by the env variable "
"LOKY_PICKLER or the function set_loky_pickler."
)
e.args = (e.args[0] + extra_info,) + e.args[1:]
e.msg = e.args[0]
raise e
util.debug(
f"Using '{loky_pickler if loky_pickler else 'cloudpickle'}' for "
"serialization."
)
class CustomizablePickler(loky_pickler_cls):
_loky_pickler_cls = loky_pickler_cls
def _set_dispatch_table(self, dispatch_table):
for ancestor_class in self._loky_pickler_cls.mro():
dt_attribute = getattr(ancestor_class, "dispatch_table", None)
if isinstance(dt_attribute, types.MemberDescriptorType):
# Ancestor class (typically _pickle.Pickler) has a
# member_descriptor for its "dispatch_table" attribute. Use
# it to set the dispatch_table as a member instead of a
# dynamic attribute in the __dict__ of the instance,
# otherwise it will not be taken into account by the C
# implementation of the dump method if a subclass defines a
# class-level dispatch_table attribute as was done in
# cloudpickle 1.6.0:
# https://github.com/joblib/loky/pull/260
dt_attribute.__set__(self, dispatch_table)
break
# On top of member descriptor set, also use setattr such that code
# that directly access self.dispatch_table gets a consistent view
# of the same table.
self.dispatch_table = dispatch_table
def __init__(self, writer, reducers=None, protocol=HIGHEST_PROTOCOL):
loky_pickler_cls.__init__(self, writer, protocol=protocol)
if reducers is None:
reducers = {}
if hasattr(self, "dispatch_table"):
# Force a copy that we will update without mutating the
# any class level defined dispatch_table.
loky_dt = dict(self.dispatch_table)
else:
# Use standard reducers as bases
loky_dt = copyreg.dispatch_table.copy()
# Register loky specific reducers
loky_dt.update(_dispatch_table)
# Set the new dispatch table, taking care of the fact that we
# need to use the member_descriptor when we inherit from a
# subclass of the C implementation of the Pickler base class
# with an class level dispatch_table attribute.
self._set_dispatch_table(loky_dt)
# Register the reducers
for type, reduce_func in reducers.items():
self.register(type, reduce_func)
def register(self, type, reduce_func):
"""Attach a reducer function to a given type in the dispatch table."""
self.dispatch_table[type] = reduce_func
_LokyPickler = CustomizablePickler
_loky_pickler_name = loky_pickler
def get_loky_pickler_name():
global _loky_pickler_name
return _loky_pickler_name
def get_loky_pickler():
global _LokyPickler
return _LokyPickler
# Set it to its default value
set_loky_pickler()
def dump(obj, file, reducers=None, protocol=None):
"""Replacement for pickle.dump() using _LokyPickler."""
global _LokyPickler
_LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
def dumps(obj, reducers=None, protocol=None):
global _LokyPickler
buf = io.BytesIO()
dump(obj, buf, reducers=reducers, protocol=protocol)
return buf.getbuffer()
__all__ = ["dump", "dumps", "loads", "register", "set_loky_pickler"]
if sys.platform == "win32":
from multiprocessing.reduction import duplicate
__all__ += ["duplicate"]

View File

@@ -0,0 +1,411 @@
###############################################################################
# Server process to keep track of unlinked resources, like folders and
# semaphores and clean them.
#
# author: Thomas Moreau
#
# Adapted from multiprocessing/resource_tracker.py
# * add some VERBOSE logging,
# * add support to track folders,
# * add Windows support,
# * refcounting scheme to avoid unlinking resources still in use.
#
# On Unix we run a server process which keeps track of unlinked
# resources. The server ignores SIGINT and SIGTERM and reads from a
# pipe. The resource_tracker implements a reference counting scheme: each time
# a Python process anticipates the shared usage of a resource by another
# process, it signals the resource_tracker of this shared usage, and in return,
# the resource_tracker increments the resource's reference count by 1.
# Similarly, when access to a resource is closed by a Python process, the
# process notifies the resource_tracker by asking it to decrement the
# resource's reference count by 1. When the reference count drops to 0, the
# resource_tracker attempts to clean up the underlying resource.
# Finally, every other process connected to the resource tracker has a copy of
# the writable end of the pipe used to communicate with it, so the resource
# tracker gets EOF when all other processes have exited. Then the
# resource_tracker process unlinks any remaining leaked resources (with
# reference count above 0)
# For semaphores, this is important because the system only supports a limited
# number of named semaphores, and they will not be automatically removed till
# the next reboot. Without this resource tracker process, "killall python"
# would probably leave unlinked semaphores.
# Note that this behavior differs from CPython's resource_tracker, which only
# implements list of shared resources, and not a proper refcounting scheme.
# Also, CPython's resource tracker will only attempt to cleanup those shared
# resources once all processes connected to the resource tracker have exited.
import os
import shutil
import sys
import signal
import warnings
from multiprocessing import util
from multiprocessing.resource_tracker import (
ResourceTracker as _ResourceTracker,
)
from . import spawn
if sys.platform == "win32":
import _winapi
import msvcrt
from multiprocessing.reduction import duplicate
__all__ = ["ensure_running", "register", "unregister"]
_HAVE_SIGMASK = hasattr(signal, "pthread_sigmask")
_IGNORED_SIGNALS = (signal.SIGINT, signal.SIGTERM)
def cleanup_noop(name):
raise RuntimeError("noop should never be registered or cleaned up")
_CLEANUP_FUNCS = {
"noop": cleanup_noop,
"folder": shutil.rmtree,
"file": os.unlink,
}
if os.name == "posix":
import _multiprocessing
# Use sem_unlink() to clean up named semaphores.
#
# sem_unlink() may be missing if the Python build process detected the
# absence of POSIX named semaphores. In that case, no named semaphores were
# ever opened, so no cleanup would be necessary.
if hasattr(_multiprocessing, "sem_unlink"):
_CLEANUP_FUNCS.update(
{
"semlock": _multiprocessing.sem_unlink,
}
)
VERBOSE = False
class ResourceTracker(_ResourceTracker):
"""Resource tracker with refcounting scheme.
This class is an extension of the multiprocessing ResourceTracker class
which implements a reference counting scheme to avoid unlinking shared
resources still in use in other processes.
This feature is notably used by `joblib.Parallel` to share temporary
folders and memory mapped files between the main process and the worker
processes.
The actual implementation of the refcounting scheme is in the main
function, which is run in a dedicated process.
"""
def maybe_unlink(self, name, rtype):
"""Decrement the refcount of a resource, and delete it if it hits 0"""
self._send("MAYBE_UNLINK", name, rtype)
def ensure_running(self):
"""Make sure that resource tracker process is running.
This can be run from any process. Usually a child process will use
the resource created by its parent.
This function is necessary for backward compatibility with python
versions before 3.13.7.
"""
return self._ensure_running_and_write()
def _teardown_dead_process(self):
# Override this function for compatibility with windows and
# for python version before 3.13.7
# At this point, the resource_tracker process has been killed
# or crashed.
os.close(self._fd)
# Let's remove the process entry from the process table on POSIX system
# to avoid zombie processes.
if os.name == "posix":
try:
# _pid can be None if this process is a child from another
# python process, which has started the resource_tracker.
if self._pid is not None:
os.waitpid(self._pid, 0)
except OSError:
# The resource_tracker has already been terminated.
pass
self._fd = None
self._pid = None
warnings.warn(
"resource_tracker: process died unexpectedly, relaunching. "
"Some folders/semaphores might leak."
)
def _launch(self):
# This is the overridden part of the resource tracker, which launches
# loky's version, which is compatible with windows and allow to track
# folders with external ref counting.
fds_to_pass = []
try:
fds_to_pass.append(sys.stderr.fileno())
except Exception:
pass
# Create a pipe for posix and windows
r, w = os.pipe()
if sys.platform == "win32":
_r = duplicate(msvcrt.get_osfhandle(r), inheritable=True)
os.close(r)
r = _r
cmd = f"from {main.__module__} import main; main({r}, {VERBOSE})"
try:
fds_to_pass.append(r)
# process will out live us, so no need to wait on pid
exe = spawn.get_executable()
args = [exe, *util._args_from_interpreter_flags(), "-c", cmd]
util.debug(f"launching resource tracker: {args}")
# bpo-33613: Register a signal mask that will block the
# signals. This signal mask will be inherited by the child
# that is going to be spawned and will protect the child from a
# race condition that can make the child die before it
# registers signal handlers for SIGINT and SIGTERM. The mask is
# unregistered after spawning the child.
try:
if _HAVE_SIGMASK:
signal.pthread_sigmask(signal.SIG_BLOCK, _IGNORED_SIGNALS)
pid = spawnv_passfds(exe, args, fds_to_pass)
finally:
if _HAVE_SIGMASK:
signal.pthread_sigmask(
signal.SIG_UNBLOCK, _IGNORED_SIGNALS
)
except BaseException:
os.close(w)
raise
else:
self._fd = w
self._pid = pid
finally:
if sys.platform == "win32":
_winapi.CloseHandle(r)
else:
os.close(r)
def _ensure_running_and_write(self, msg=None):
"""Make sure that resource tracker process is running.
This can be run from any process. Usually a child process will use
the resource created by its parent.
This function is added for compatibility with python version before 3.13.7.
"""
with self._lock:
if (
self._fd is not None
): # resource tracker was launched before, is it still running?
if msg is None:
to_send = b"PROBE:0:noop\n"
else:
to_send = msg
try:
self._write(to_send)
except OSError:
self._teardown_dead_process()
self._launch()
msg = None # message was sent in probe
else:
self._launch()
if msg is not None:
self._write(msg)
def _write(self, msg):
nbytes = os.write(self._fd, msg)
assert nbytes == len(msg), f"{nbytes=} != {len(msg)=}"
def __del__(self):
# ignore error due to trying to clean up child process which has already been
# shutdown on windows. See https://github.com/joblib/loky/pull/450
# This is only required if __del__ is defined
if not hasattr(_ResourceTracker, "__del__"):
return
try:
super().__del__()
except ChildProcessError:
pass
_resource_tracker = ResourceTracker()
ensure_running = _resource_tracker.ensure_running
register = _resource_tracker.register
maybe_unlink = _resource_tracker.maybe_unlink
unregister = _resource_tracker.unregister
getfd = _resource_tracker.getfd
def main(fd, verbose=0):
"""Run resource tracker."""
if verbose:
util.log_to_stderr(level=util.DEBUG)
# protect the process from ^C and "killall python" etc
signal.signal(signal.SIGINT, signal.SIG_IGN)
signal.signal(signal.SIGTERM, signal.SIG_IGN)
if _HAVE_SIGMASK:
signal.pthread_sigmask(signal.SIG_UNBLOCK, _IGNORED_SIGNALS)
for f in (sys.stdin, sys.stdout):
try:
f.close()
except Exception:
pass
if verbose:
util.debug("Main resource tracker is running")
registry = {rtype: {} for rtype in _CLEANUP_FUNCS.keys()}
try:
if sys.platform == "win32":
fd = msvcrt.open_osfhandle(fd, os.O_RDONLY)
# keep track of registered/unregistered resources
with open(fd, "rb") as f:
for line in f:
try:
splitted = line.strip().decode("ascii").split(":")
# name can potentially contain separator symbols (for
# instance folders on Windows)
cmd, name, rtype = (
splitted[0],
":".join(splitted[1:-1]),
splitted[-1],
)
if rtype not in _CLEANUP_FUNCS:
raise ValueError(
f"Cannot register {name} for automatic cleanup: "
f"unknown resource type ({rtype}). Resource type "
"should be one of the following: "
f"{list(_CLEANUP_FUNCS.keys())}"
)
if cmd == "PROBE":
pass
elif cmd == "REGISTER":
if name not in registry[rtype]:
registry[rtype][name] = 1
else:
registry[rtype][name] += 1
if verbose:
util.debug(
"[ResourceTracker] incremented refcount of "
f"{rtype} {name} "
f"(current {registry[rtype][name]})"
)
elif cmd == "UNREGISTER":
del registry[rtype][name]
if verbose:
util.debug(
f"[ResourceTracker] unregister {name} {rtype}: "
f"registry({len(registry)})"
)
elif cmd == "MAYBE_UNLINK":
registry[rtype][name] -= 1
if verbose:
util.debug(
"[ResourceTracker] decremented refcount of "
f"{rtype} {name} "
f"(current {registry[rtype][name]})"
)
if registry[rtype][name] == 0:
del registry[rtype][name]
try:
if verbose:
util.debug(
f"[ResourceTracker] unlink {name}"
)
_CLEANUP_FUNCS[rtype](name)
except Exception as e:
warnings.warn(
f"resource_tracker: {name}: {e!r}"
)
else:
raise RuntimeError(f"unrecognized command {cmd!r}")
except BaseException:
try:
sys.excepthook(*sys.exc_info())
except BaseException:
pass
finally:
# all processes have terminated; cleanup any remaining resources
def _unlink_resources(rtype_registry, rtype):
if rtype_registry:
try:
warnings.warn(
"resource_tracker: There appear to be "
f"{len(rtype_registry)} leaked {rtype} objects to "
"clean up at shutdown"
)
except Exception:
pass
for name in rtype_registry:
# For some reason the process which created and registered this
# resource has failed to unregister it. Presumably it has
# died. We therefore clean it up.
try:
_CLEANUP_FUNCS[rtype](name)
if verbose:
util.debug(f"[ResourceTracker] unlink {name}")
except Exception as e:
warnings.warn(f"resource_tracker: {name}: {e!r}")
for rtype, rtype_registry in registry.items():
if rtype == "folder":
continue
else:
_unlink_resources(rtype_registry, rtype)
# The default cleanup routine for folders deletes everything inside
# those folders recursively, which can include other resources tracked
# by the resource tracker). To limit the risk of the resource tracker
# attempting to delete twice a resource (once as part of a tracked
# folder, and once as a resource), we delete the folders after all
# other resource types.
if "folder" in registry:
_unlink_resources(registry["folder"], "folder")
if verbose:
util.debug("resource tracker shut down")
def spawnv_passfds(path, args, passfds):
if sys.platform != "win32":
args = [arg.encode("utf-8") for arg in args]
path = path.encode("utf-8")
return util.spawnv_passfds(path, args, passfds)
else:
passfds = sorted(passfds)
cmd = " ".join(f'"{x}"' for x in args)
try:
_, ht, pid, _ = _winapi.CreateProcess(
path, cmd, None, None, True, 0, None, None, None
)
_winapi.CloseHandle(ht)
except BaseException:
pass
return pid

View File

@@ -0,0 +1,244 @@
###############################################################################
# Prepares and processes the data to setup the new process environment
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/spawn.py (17/02/2017)
# * Improve logging data
#
import os
import sys
import runpy
import textwrap
import types
from multiprocessing import process, util
if sys.platform != "win32":
WINEXE = False
WINSERVICE = False
else:
import msvcrt
from multiprocessing.reduction import duplicate
WINEXE = sys.platform == "win32" and getattr(sys, "frozen", False)
WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
if WINSERVICE:
_python_exe = os.path.join(sys.exec_prefix, "python.exe")
else:
_python_exe = sys.executable
def get_executable():
return _python_exe
def _check_not_importing_main():
if getattr(process.current_process(), "_inheriting", False):
raise RuntimeError(
textwrap.dedent(
"""\
An attempt has been made to start a new process before the
current process has finished its bootstrapping phase.
This probably means that you are not using fork to start your
child processes and you have forgotten to use the proper idiom
in the main module:
if __name__ == '__main__':
freeze_support()
...
The "freeze_support()" line can be omitted if the program
is not going to be frozen to produce an executable."""
)
)
def get_preparation_data(name, init_main_module=True):
"""Return info about parent needed by child to unpickle process object."""
_check_not_importing_main()
d = dict(
log_to_stderr=util._log_to_stderr,
authkey=bytes(process.current_process().authkey),
name=name,
sys_argv=sys.argv,
orig_dir=process.ORIGINAL_DIR,
dir=os.getcwd(),
)
# Send sys_path and make sure the current directory will not be changed
d["sys_path"] = [p if p != "" else process.ORIGINAL_DIR for p in sys.path]
# Make sure to pass the information if the multiprocessing logger is active
if util._logger is not None:
d["log_level"] = util._logger.getEffectiveLevel()
if util._logger.handlers:
h = util._logger.handlers[0]
d["log_fmt"] = h.formatter._fmt
# Tell the child how to communicate with the resource_tracker
from .resource_tracker import _resource_tracker
_resource_tracker.ensure_running()
if sys.platform == "win32":
d["tracker_fd"] = msvcrt.get_osfhandle(_resource_tracker._fd)
else:
d["tracker_fd"] = _resource_tracker._fd
if os.name == "posix":
# joblib/loky#242: allow loky processes to retrieve the resource
# tracker of their parent in case the child processes depickles
# shared_memory objects, that are still tracked by multiprocessing's
# resource_tracker by default.
# XXX: this is a workaround that may be error prone: in the future, it
# would be better to have loky subclass multiprocessing's shared_memory
# to force registration of shared_memory segments via loky's
# resource_tracker.
from multiprocessing.resource_tracker import (
_resource_tracker as mp_resource_tracker,
)
# multiprocessing's resource_tracker must be running before loky
# process is created (othewise the child won't be able to use it if it
# is created later on)
mp_resource_tracker.ensure_running()
d["mp_tracker_fd"] = mp_resource_tracker._fd
# Figure out whether to initialise main in the subprocess as a module
# or through direct execution (or to leave it alone entirely)
if init_main_module:
main_module = sys.modules["__main__"]
try:
main_mod_name = getattr(main_module.__spec__, "name", None)
except BaseException:
main_mod_name = None
if main_mod_name is not None:
d["init_main_from_name"] = main_mod_name
elif sys.platform != "win32" or (not WINEXE and not WINSERVICE):
main_path = getattr(main_module, "__file__", None)
if main_path is not None:
if (
not os.path.isabs(main_path)
and process.ORIGINAL_DIR is not None
):
main_path = os.path.join(process.ORIGINAL_DIR, main_path)
d["init_main_from_path"] = os.path.normpath(main_path)
return d
#
# Prepare current process
#
old_main_modules = []
def prepare(data, parent_sentinel=None):
"""Try to get current process ready to unpickle process object."""
if "name" in data:
process.current_process().name = data["name"]
if "authkey" in data:
process.current_process().authkey = data["authkey"]
if "log_to_stderr" in data and data["log_to_stderr"]:
util.log_to_stderr()
if "log_level" in data:
util.get_logger().setLevel(data["log_level"])
if "log_fmt" in data:
import logging
util.get_logger().handlers[0].setFormatter(
logging.Formatter(data["log_fmt"])
)
if "sys_path" in data:
sys.path = data["sys_path"]
if "sys_argv" in data:
sys.argv = data["sys_argv"]
if "dir" in data:
os.chdir(data["dir"])
if "orig_dir" in data:
process.ORIGINAL_DIR = data["orig_dir"]
if "mp_tracker_fd" in data:
from multiprocessing.resource_tracker import (
_resource_tracker as mp_resource_tracker,
)
mp_resource_tracker._fd = data["mp_tracker_fd"]
if "tracker_fd" in data:
from .resource_tracker import _resource_tracker
if sys.platform == "win32":
handle = data["tracker_fd"]
handle = duplicate(handle, source_process=parent_sentinel)
_resource_tracker._fd = msvcrt.open_osfhandle(handle, os.O_RDONLY)
else:
_resource_tracker._fd = data["tracker_fd"]
if "init_main_from_name" in data:
_fixup_main_from_name(data["init_main_from_name"])
elif "init_main_from_path" in data:
_fixup_main_from_path(data["init_main_from_path"])
# Multiprocessing module helpers to fix up the main module in
# spawned subprocesses
def _fixup_main_from_name(mod_name):
# __main__.py files for packages, directories, zip archives, etc, run
# their "main only" code unconditionally, so we don't even try to
# populate anything in __main__, nor do we make any changes to
# __main__ attributes
current_main = sys.modules["__main__"]
if mod_name == "__main__" or mod_name.endswith(".__main__"):
return
# If this process was forked, __main__ may already be populated
if getattr(current_main.__spec__, "name", None) == mod_name:
return
# Otherwise, __main__ may contain some non-main code where we need to
# support unpickling it properly. We rerun it as __mp_main__ and make
# the normal __main__ an alias to that
old_main_modules.append(current_main)
main_module = types.ModuleType("__mp_main__")
main_content = runpy.run_module(
mod_name, run_name="__mp_main__", alter_sys=True
)
main_module.__dict__.update(main_content)
sys.modules["__main__"] = sys.modules["__mp_main__"] = main_module
def _fixup_main_from_path(main_path):
# If this process was forked, __main__ may already be populated
current_main = sys.modules["__main__"]
# Unfortunately, the main ipython launch script historically had no
# "if __name__ == '__main__'" guard, so we work around that
# by treating it like a __main__.py file
# See https://github.com/ipython/ipython/issues/4698
main_name = os.path.splitext(os.path.basename(main_path))[0]
if main_name == "ipython":
return
# Otherwise, if __file__ already has the setting we expect,
# there's nothing more to do
if getattr(current_main, "__file__", None) == main_path:
return
# If the parent process has sent a path through rather than a module
# name we assume it is an executable script that may contain
# non-main code that needs to be executed
old_main_modules.append(current_main)
main_module = types.ModuleType("__mp_main__")
main_content = runpy.run_path(main_path, run_name="__mp_main__")
main_module.__dict__.update(main_content)
sys.modules["__main__"] = sys.modules["__mp_main__"] = main_module

View File

@@ -0,0 +1,409 @@
###############################################################################
# Synchronization primitives based on our SemLock implementation
#
# author: Thomas Moreau and Olivier Grisel
#
# adapted from multiprocessing/synchronize.py (17/02/2017)
# * Remove ctx argument for compatibility reason
# * Registers a cleanup function with the loky resource_tracker to remove the
# semaphore when the process dies instead.
#
# TODO: investigate which Python version is required to be able to use
# multiprocessing.resource_tracker and therefore multiprocessing.synchronize
# instead of a loky-specific fork.
import os
import sys
import tempfile
import threading
import _multiprocessing
from time import time as _time
from multiprocessing import process, util
from multiprocessing.context import assert_spawning
from . import resource_tracker
__all__ = [
"Lock",
"RLock",
"Semaphore",
"BoundedSemaphore",
"Condition",
"Event",
]
# Try to import the mp.synchronize module cleanly, if it fails
# raise ImportError for platforms lacking a working sem_open implementation.
# See issue 3770
try:
from _multiprocessing import SemLock as _SemLock
from _multiprocessing import sem_unlink
except ImportError:
raise ImportError(
"This platform lacks a functioning sem_open"
" implementation, therefore, the required"
" synchronization primitives needed will not"
" function, see issue 3770."
)
#
# Constants
#
RECURSIVE_MUTEX, SEMAPHORE = range(2)
SEM_VALUE_MAX = _multiprocessing.SemLock.SEM_VALUE_MAX
#
# Base class for semaphores and mutexes; wraps `_multiprocessing.SemLock`
#
class SemLock:
_rand = tempfile._RandomNameSequence()
def __init__(self, kind, value, maxvalue, name=None):
# unlink_now is only used on win32 or when we are using fork.
unlink_now = False
if name is None:
# Try to find an unused name for the SemLock instance.
for _ in range(100):
try:
self._semlock = _SemLock(
kind, value, maxvalue, SemLock._make_name(), unlink_now
)
except FileExistsError: # pragma: no cover
pass
else:
break
else: # pragma: no cover
raise FileExistsError("cannot find name for semaphore")
else:
self._semlock = _SemLock(kind, value, maxvalue, name, unlink_now)
self.name = name
util.debug(
f"created semlock with handle {self._semlock.handle} and name "
f'"{self.name}"'
)
self._make_methods()
def _after_fork(obj):
obj._semlock._after_fork()
util.register_after_fork(self, _after_fork)
# When the object is garbage collected or the
# process shuts down we unlink the semaphore name
resource_tracker.register(self._semlock.name, "semlock")
util.Finalize(
self, SemLock._cleanup, (self._semlock.name,), exitpriority=0
)
@staticmethod
def _cleanup(name):
try:
sem_unlink(name)
except FileNotFoundError:
# Already unlinked, possibly by user code: ignore and make sure to
# unregister the semaphore from the resource tracker.
pass
finally:
resource_tracker.unregister(name, "semlock")
def _make_methods(self):
self.acquire = self._semlock.acquire
self.release = self._semlock.release
def __enter__(self):
return self._semlock.acquire()
def __exit__(self, *args):
return self._semlock.release()
def __getstate__(self):
assert_spawning(self)
sl = self._semlock
h = sl.handle
return (h, sl.kind, sl.maxvalue, sl.name)
def __setstate__(self, state):
self._semlock = _SemLock._rebuild(*state)
util.debug(
f'recreated blocker with handle {state[0]!r} and name "{state[3]}"'
)
self._make_methods()
@staticmethod
def _make_name():
# OSX does not support long names for semaphores
return f"/loky-{os.getpid()}-{next(SemLock._rand)}"
#
# Semaphore
#
class Semaphore(SemLock):
def __init__(self, value=1):
SemLock.__init__(self, SEMAPHORE, value, SEM_VALUE_MAX)
def get_value(self):
if sys.platform == "darwin":
raise NotImplementedError("OSX does not implement sem_getvalue")
return self._semlock._get_value()
def __repr__(self):
try:
value = self._semlock._get_value()
except Exception:
value = "unknown"
return f"<{self.__class__.__name__}(value={value})>"
#
# Bounded semaphore
#
class BoundedSemaphore(Semaphore):
def __init__(self, value=1):
SemLock.__init__(self, SEMAPHORE, value, value)
def __repr__(self):
try:
value = self._semlock._get_value()
except Exception:
value = "unknown"
return (
f"<{self.__class__.__name__}(value={value}, "
f"maxvalue={self._semlock.maxvalue})>"
)
#
# Non-recursive lock
#
class Lock(SemLock):
def __init__(self):
super().__init__(SEMAPHORE, 1, 1)
def __repr__(self):
try:
if self._semlock._is_mine():
name = process.current_process().name
if threading.current_thread().name != "MainThread":
name = f"{name}|{threading.current_thread().name}"
elif self._semlock._get_value() == 1:
name = "None"
elif self._semlock._count() > 0:
name = "SomeOtherThread"
else:
name = "SomeOtherProcess"
except Exception:
name = "unknown"
return f"<{self.__class__.__name__}(owner={name})>"
#
# Recursive lock
#
class RLock(SemLock):
def __init__(self):
super().__init__(RECURSIVE_MUTEX, 1, 1)
def __repr__(self):
try:
if self._semlock._is_mine():
name = process.current_process().name
if threading.current_thread().name != "MainThread":
name = f"{name}|{threading.current_thread().name}"
count = self._semlock._count()
elif self._semlock._get_value() == 1:
name, count = "None", 0
elif self._semlock._count() > 0:
name, count = "SomeOtherThread", "nonzero"
else:
name, count = "SomeOtherProcess", "nonzero"
except Exception:
name, count = "unknown", "unknown"
return f"<{self.__class__.__name__}({name}, {count})>"
#
# Condition variable
#
class Condition:
def __init__(self, lock=None):
self._lock = lock or RLock()
self._sleeping_count = Semaphore(0)
self._woken_count = Semaphore(0)
self._wait_semaphore = Semaphore(0)
self._make_methods()
def __getstate__(self):
assert_spawning(self)
return (
self._lock,
self._sleeping_count,
self._woken_count,
self._wait_semaphore,
)
def __setstate__(self, state):
(
self._lock,
self._sleeping_count,
self._woken_count,
self._wait_semaphore,
) = state
self._make_methods()
def __enter__(self):
return self._lock.__enter__()
def __exit__(self, *args):
return self._lock.__exit__(*args)
def _make_methods(self):
self.acquire = self._lock.acquire
self.release = self._lock.release
def __repr__(self):
try:
num_waiters = (
self._sleeping_count._semlock._get_value()
- self._woken_count._semlock._get_value()
)
except Exception:
num_waiters = "unknown"
return f"<{self.__class__.__name__}({self._lock}, {num_waiters})>"
def wait(self, timeout=None):
assert (
self._lock._semlock._is_mine()
), "must acquire() condition before using wait()"
# indicate that this thread is going to sleep
self._sleeping_count.release()
# release lock
count = self._lock._semlock._count()
for _ in range(count):
self._lock.release()
try:
# wait for notification or timeout
return self._wait_semaphore.acquire(True, timeout)
finally:
# indicate that this thread has woken
self._woken_count.release()
# reacquire lock
for _ in range(count):
self._lock.acquire()
def notify(self):
assert self._lock._semlock._is_mine(), "lock is not owned"
assert not self._wait_semaphore.acquire(False)
# to take account of timeouts since last notify() we subtract
# woken_count from sleeping_count and rezero woken_count
while self._woken_count.acquire(False):
res = self._sleeping_count.acquire(False)
assert res
if self._sleeping_count.acquire(False): # try grabbing a sleeper
self._wait_semaphore.release() # wake up one sleeper
self._woken_count.acquire() # wait for the sleeper to wake
# rezero _wait_semaphore in case a timeout just happened
self._wait_semaphore.acquire(False)
def notify_all(self):
assert self._lock._semlock._is_mine(), "lock is not owned"
assert not self._wait_semaphore.acquire(False)
# to take account of timeouts since last notify*() we subtract
# woken_count from sleeping_count and rezero woken_count
while self._woken_count.acquire(False):
res = self._sleeping_count.acquire(False)
assert res
sleepers = 0
while self._sleeping_count.acquire(False):
self._wait_semaphore.release() # wake up one sleeper
sleepers += 1
if sleepers:
for _ in range(sleepers):
self._woken_count.acquire() # wait for a sleeper to wake
# rezero wait_semaphore in case some timeouts just happened
while self._wait_semaphore.acquire(False):
pass
def wait_for(self, predicate, timeout=None):
result = predicate()
if result:
return result
if timeout is not None:
endtime = _time() + timeout
else:
endtime = None
waittime = None
while not result:
if endtime is not None:
waittime = endtime - _time()
if waittime <= 0:
break
self.wait(waittime)
result = predicate()
return result
#
# Event
#
class Event:
def __init__(self):
self._cond = Condition(Lock())
self._flag = Semaphore(0)
def is_set(self):
with self._cond:
if self._flag.acquire(False):
self._flag.release()
return True
return False
def set(self):
with self._cond:
self._flag.acquire(False)
self._flag.release()
self._cond.notify_all()
def clear(self):
with self._cond:
self._flag.acquire(False)
def wait(self, timeout=None):
with self._cond:
if self._flag.acquire(False):
self._flag.release()
else:
self._cond.wait(timeout)
if self._flag.acquire(False):
self._flag.release()
return True
return False

View File

@@ -0,0 +1,181 @@
import os
import sys
import time
import errno
import signal
import warnings
import subprocess
import traceback
try:
import psutil
except ImportError:
psutil = None
def kill_process_tree(process, use_psutil=True):
"""Terminate process and its descendants with SIGKILL"""
if use_psutil and psutil is not None:
_kill_process_tree_with_psutil(process)
else:
_kill_process_tree_without_psutil(process)
def recursive_terminate(process, use_psutil=True):
warnings.warn(
"recursive_terminate is deprecated in loky 3.2, use kill_process_tree"
"instead",
DeprecationWarning,
)
kill_process_tree(process, use_psutil=use_psutil)
def _kill_process_tree_with_psutil(process):
try:
descendants = psutil.Process(process.pid).children(recursive=True)
except psutil.NoSuchProcess:
return
# Kill the descendants in reverse order to avoid killing the parents before
# the descendant in cases where there are more processes nested.
for descendant in descendants[::-1]:
try:
descendant.kill()
except psutil.NoSuchProcess:
pass
try:
psutil.Process(process.pid).kill()
except psutil.NoSuchProcess:
pass
process.join()
def _kill_process_tree_without_psutil(process):
"""Terminate a process and its descendants."""
try:
if sys.platform == "win32":
_windows_taskkill_process_tree(process.pid)
else:
_posix_recursive_kill(process.pid)
except Exception: # pragma: no cover
details = traceback.format_exc()
warnings.warn(
"Failed to kill subprocesses on this platform. Please install"
"psutil: https://github.com/giampaolo/psutil\n"
f"Details:\n{details}"
)
# In case we cannot introspect or kill the descendants, we fall back to
# only killing the main process.
#
# Note: on Windows, process.kill() is an alias for process.terminate()
# which in turns calls the Win32 API function TerminateProcess().
process.kill()
process.join()
def _windows_taskkill_process_tree(pid):
# On windows, the taskkill function with option `/T` terminate a given
# process pid and its children.
try:
subprocess.check_output(
["taskkill", "/F", "/T", "/PID", str(pid)], stderr=None
)
except subprocess.CalledProcessError as e:
# In Windows, taskkill returns 128, 255 for no process found.
if e.returncode not in [128, 255]:
# Let's raise to let the caller log the error details in a
# warning and only kill the root process.
raise # pragma: no cover
def _kill(pid):
# Not all systems (e.g. Windows) have a SIGKILL, but the C specification
# mandates a SIGTERM signal. While Windows is handled specifically above,
# let's try to be safe for other hypothetic platforms that only have
# SIGTERM without SIGKILL.
kill_signal = getattr(signal, "SIGKILL", signal.SIGTERM)
try:
os.kill(pid, kill_signal)
except OSError as e:
# if OSError is raised with [Errno 3] no such process, the process
# is already terminated, else, raise the error and let the top
# level function raise a warning and retry to kill the process.
if e.errno != errno.ESRCH:
raise # pragma: no cover
def _posix_recursive_kill(pid):
"""Recursively kill the descendants of a process before killing it."""
try:
children_pids = subprocess.check_output(
["pgrep", "-P", str(pid)], stderr=None, text=True
)
except subprocess.CalledProcessError as e:
# `ps` returns 1 when no child process has been found
if e.returncode == 1:
children_pids = ""
else:
raise # pragma: no cover
# Decode the result, split the cpid and remove the trailing line
for cpid in children_pids.splitlines():
cpid = int(cpid)
_posix_recursive_kill(cpid)
_kill(pid)
def get_exitcodes_terminated_worker(processes):
"""Return a formatted string with the exitcodes of terminated workers.
If necessary, wait (up to .25s) for the system to correctly set the
exitcode of one terminated worker.
"""
patience = 5
# Catch the exitcode of the terminated workers. There should at least be
# one. If not, wait a bit for the system to correctly set the exitcode of
# the terminated worker.
exitcodes = [
p.exitcode for p in list(processes.values()) if p.exitcode is not None
]
while not exitcodes and patience > 0:
patience -= 1
exitcodes = [
p.exitcode
for p in list(processes.values())
if p.exitcode is not None
]
time.sleep(0.05)
return _format_exitcodes(exitcodes)
def _format_exitcodes(exitcodes):
"""Format a list of exit code with names of the signals if possible"""
str_exitcodes = [
f"{_get_exitcode_name(e)}({e})" for e in exitcodes if e is not None
]
return "{" + ", ".join(str_exitcodes) + "}"
def _get_exitcode_name(exitcode):
if sys.platform == "win32":
# The exitcode are unreliable on windows (see bpo-31863).
# For this case, return UNKNOWN
return "UNKNOWN"
if exitcode < 0:
try:
import signal
return signal.Signals(-exitcode).name
except ValueError:
return "UNKNOWN"
elif exitcode != 255:
# The exitcode are unreliable on forkserver were 255 is always returned
# (see bpo-30589). For this case, return UNKNOWN
return "EXIT"
return "UNKNOWN"

View File

@@ -0,0 +1,102 @@
import inspect
from functools import partial
from joblib.externals.cloudpickle import dumps, loads
WRAP_CACHE = {}
class CloudpickledObjectWrapper:
def __init__(self, obj, keep_wrapper=False):
self._obj = obj
self._keep_wrapper = keep_wrapper
def __reduce__(self):
_pickled_object = dumps(self._obj)
if not self._keep_wrapper:
return loads, (_pickled_object,)
return _reconstruct_wrapper, (_pickled_object, self._keep_wrapper)
def __getattr__(self, attr):
# Ensure that the wrapped object can be used seemlessly as the
# previous object.
if attr not in ["_obj", "_keep_wrapper"]:
return getattr(self._obj, attr)
return getattr(self, attr)
# Make sure the wrapped object conserves the callable property
class CallableObjectWrapper(CloudpickledObjectWrapper):
def __call__(self, *args, **kwargs):
return self._obj(*args, **kwargs)
def _wrap_non_picklable_objects(obj, keep_wrapper):
if callable(obj):
return CallableObjectWrapper(obj, keep_wrapper=keep_wrapper)
return CloudpickledObjectWrapper(obj, keep_wrapper=keep_wrapper)
def _reconstruct_wrapper(_pickled_object, keep_wrapper):
obj = loads(_pickled_object)
return _wrap_non_picklable_objects(obj, keep_wrapper)
def _wrap_objects_when_needed(obj):
# Function to introspect an object and decide if it should be wrapped or
# not.
need_wrap = "__main__" in getattr(obj, "__module__", "")
if isinstance(obj, partial):
return partial(
_wrap_objects_when_needed(obj.func),
*[_wrap_objects_when_needed(a) for a in obj.args],
**{
k: _wrap_objects_when_needed(v)
for k, v in obj.keywords.items()
},
)
if callable(obj):
# Need wrap if the object is a function defined in a local scope of
# another function.
func_code = getattr(obj, "__code__", "")
need_wrap |= getattr(func_code, "co_flags", 0) & inspect.CO_NESTED
# Need wrap if the obj is a lambda expression
func_name = getattr(obj, "__name__", "")
need_wrap |= "<lambda>" in func_name
if not need_wrap:
return obj
wrapped_obj = WRAP_CACHE.get(obj)
if wrapped_obj is None:
wrapped_obj = _wrap_non_picklable_objects(obj, keep_wrapper=False)
WRAP_CACHE[obj] = wrapped_obj
return wrapped_obj
def wrap_non_picklable_objects(obj, keep_wrapper=True):
"""Wrapper for non-picklable object to use cloudpickle to serialize them.
Note that this wrapper tends to slow down the serialization process as it
is done with cloudpickle which is typically slower compared to pickle. The
proper way to solve serialization issues is to avoid defining functions and
objects in the main scripts and to implement __reduce__ functions for
complex classes.
"""
# If obj is a class, create a CloudpickledClassWrapper which instantiates
# the object internally and wrap it directly in a CloudpickledObjectWrapper
if inspect.isclass(obj):
class CloudpickledClassWrapper(CloudpickledObjectWrapper):
def __init__(self, *args, **kwargs):
self._obj = obj(*args, **kwargs)
self._keep_wrapper = keep_wrapper
CloudpickledClassWrapper.__name__ = obj.__name__
return CloudpickledClassWrapper
# If obj is an instance of a class, just wrap it in a regular
# CloudpickledObjectWrapper
return _wrap_non_picklable_objects(obj, keep_wrapper=keep_wrapper)

View File

@@ -0,0 +1,80 @@
import warnings
def _viztracer_init(init_kwargs):
"""Initialize viztracer's profiler in worker processes"""
from viztracer import VizTracer
tracer = VizTracer(**init_kwargs)
tracer.register_exit()
tracer.start()
def _make_viztracer_initializer_and_initargs():
try:
import viztracer
tracer = viztracer.get_tracer()
if tracer is not None and getattr(tracer, "enable", False):
# Profiler is active: introspect its configuration to
# initialize the workers with the same configuration.
return _viztracer_init, (tracer.init_kwargs,)
except ImportError:
# viztracer is not installed: nothing to do
pass
except Exception as e:
# In case viztracer's API evolve, we do not want to crash loky but
# we want to know about it to be able to update loky.
warnings.warn(f"Unable to introspect viztracer state: {e}")
return None, ()
class _ChainedInitializer:
"""Compound worker initializer
This is meant to be used in conjunction with _chain_initializers to
produce the necessary chained_args list to be passed to __call__.
"""
def __init__(self, initializers):
self._initializers = initializers
def __call__(self, *chained_args):
for initializer, args in zip(self._initializers, chained_args):
initializer(*args)
def _chain_initializers(initializer_and_args):
"""Convenience helper to combine a sequence of initializers.
If some initializers are None, they are filtered out.
"""
filtered_initializers = []
filtered_initargs = []
for initializer, initargs in initializer_and_args:
if initializer is not None:
filtered_initializers.append(initializer)
filtered_initargs.append(initargs)
if not filtered_initializers:
return None, ()
elif len(filtered_initializers) == 1:
return filtered_initializers[0], filtered_initargs[0]
else:
return _ChainedInitializer(filtered_initializers), filtered_initargs
def _prepare_initializer(initializer, initargs):
if initializer is not None and not callable(initializer):
raise TypeError(
f"initializer must be a callable, got: {initializer!r}"
)
# Introspect runtime to determine if we need to propagate the viztracer
# profiler information to the workers:
return _chain_initializers(
[
(initializer, initargs),
_make_viztracer_initializer_and_initargs(),
]
)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,294 @@
###############################################################################
# Reusable ProcessPoolExecutor
#
# author: Thomas Moreau and Olivier Grisel
#
import time
import warnings
import threading
import multiprocessing as mp
from .process_executor import ProcessPoolExecutor, EXTRA_QUEUED_CALLS
from .backend.context import cpu_count
from .backend import get_context
__all__ = ["get_reusable_executor"]
# Singleton executor and id management
_executor_lock = threading.RLock()
_next_executor_id = 0
_executor = None
_executor_kwargs = None
def _get_next_executor_id():
"""Ensure that each successive executor instance has a unique, monotonic id.
The purpose of this monotonic id is to help debug and test automated
instance creation.
"""
global _next_executor_id
with _executor_lock:
executor_id = _next_executor_id
_next_executor_id += 1
return executor_id
def get_reusable_executor(
max_workers=None,
context=None,
timeout=10,
kill_workers=False,
reuse="auto",
job_reducers=None,
result_reducers=None,
initializer=None,
initargs=(),
env=None,
):
"""Return the current ReusableExectutor instance.
Start a new instance if it has not been started already or if the previous
instance was left in a broken state.
If the previous instance does not have the requested number of workers, the
executor is dynamically resized to adjust the number of workers prior to
returning.
Reusing a singleton instance spares the overhead of starting new worker
processes and importing common python packages each time.
``max_workers`` controls the maximum number of tasks that can be running in
parallel in worker processes. By default this is set to the number of
CPUs on the host.
Setting ``timeout`` (in seconds) makes idle workers automatically shutdown
so as to release system resources. New workers are respawn upon submission
of new tasks so that ``max_workers`` are available to accept the newly
submitted tasks. Setting ``timeout`` to around 100 times the time required
to spawn new processes and import packages in them (on the order of 100ms)
ensures that the overhead of spawning workers is negligible.
Setting ``kill_workers=True`` makes it possible to forcibly interrupt
previously spawned jobs to get a new instance of the reusable executor
with new constructor argument values.
The ``job_reducers`` and ``result_reducers`` are used to customize the
pickling of tasks and results send to the executor.
When provided, the ``initializer`` is run first in newly spawned
processes with argument ``initargs``.
The environment variable in the child process are a copy of the values in
the main process. One can provide a dict ``{ENV: VAL}`` where ``ENV`` and
``VAL`` are string literals to overwrite the environment variable ``ENV``
in the child processes to value ``VAL``. The environment variables are set
in the children before any module is loaded. This only works with the
``loky`` context.
"""
_executor, _ = _ReusablePoolExecutor.get_reusable_executor(
max_workers=max_workers,
context=context,
timeout=timeout,
kill_workers=kill_workers,
reuse=reuse,
job_reducers=job_reducers,
result_reducers=result_reducers,
initializer=initializer,
initargs=initargs,
env=env,
)
return _executor
class _ReusablePoolExecutor(ProcessPoolExecutor):
def __init__(
self,
submit_resize_lock,
max_workers=None,
context=None,
timeout=None,
executor_id=0,
job_reducers=None,
result_reducers=None,
initializer=None,
initargs=(),
env=None,
):
super().__init__(
max_workers=max_workers,
context=context,
timeout=timeout,
job_reducers=job_reducers,
result_reducers=result_reducers,
initializer=initializer,
initargs=initargs,
env=env,
)
self.executor_id = executor_id
self._submit_resize_lock = submit_resize_lock
@classmethod
def get_reusable_executor(
cls,
max_workers=None,
context=None,
timeout=10,
kill_workers=False,
reuse="auto",
job_reducers=None,
result_reducers=None,
initializer=None,
initargs=(),
env=None,
):
with _executor_lock:
global _executor, _executor_kwargs
executor = _executor
if max_workers is None:
if reuse is True and executor is not None:
max_workers = executor._max_workers
else:
max_workers = cpu_count()
elif max_workers <= 0:
raise ValueError(
f"max_workers must be greater than 0, got {max_workers}."
)
if isinstance(context, str):
context = get_context(context)
if context is not None and context.get_start_method() == "fork":
raise ValueError(
"Cannot use reusable executor with the 'fork' context"
)
kwargs = dict(
context=context,
timeout=timeout,
job_reducers=job_reducers,
result_reducers=result_reducers,
initializer=initializer,
initargs=initargs,
env=env,
)
if executor is None:
is_reused = False
mp.util.debug(
f"Create a executor with max_workers={max_workers}."
)
executor_id = _get_next_executor_id()
_executor_kwargs = kwargs
_executor = executor = cls(
_executor_lock,
max_workers=max_workers,
executor_id=executor_id,
**kwargs,
)
else:
if reuse == "auto":
reuse = kwargs == _executor_kwargs
if (
executor._flags.broken
or executor._flags.shutdown
or not reuse
or executor.queue_size < max_workers
):
if executor._flags.broken:
reason = "broken"
elif executor._flags.shutdown:
reason = "shutdown"
elif executor.queue_size < max_workers:
# Do not reuse the executor if the queue size is too
# small as this would lead to limited parallelism.
reason = "queue size is too small"
else:
reason = "arguments have changed"
mp.util.debug(
"Creating a new executor with max_workers="
f"{max_workers} as the previous instance cannot be "
f"reused ({reason})."
)
executor.shutdown(wait=True, kill_workers=kill_workers)
_executor = executor = _executor_kwargs = None
# Recursive call to build a new instance
return cls.get_reusable_executor(
max_workers=max_workers, **kwargs
)
else:
mp.util.debug(
"Reusing existing executor with "
f"max_workers={executor._max_workers}."
)
is_reused = True
executor._resize(max_workers)
return executor, is_reused
def submit(self, fn, *args, **kwargs):
with self._submit_resize_lock:
return super().submit(fn, *args, **kwargs)
def _resize(self, max_workers):
with self._submit_resize_lock:
if max_workers is None:
raise ValueError("Trying to resize with max_workers=None")
elif max_workers == self._max_workers:
return
if self._executor_manager_thread is None:
# If the executor_manager_thread has not been started
# then no processes have been spawned and we can just
# update _max_workers and return
self._max_workers = max_workers
return
self._wait_job_completion()
# Some process might have returned due to timeout so check how many
# children are still alive. Use the _process_management_lock to
# ensure that no process are spawned or timeout during the resize.
with self._processes_management_lock:
processes = list(self._processes.values())
nb_children_alive = sum(p.is_alive() for p in processes)
self._max_workers = max_workers
for _ in range(max_workers, nb_children_alive):
self._call_queue.put(None)
while (
len(self._processes) > max_workers and not self._flags.broken
):
time.sleep(1e-3)
self._adjust_process_count()
processes = list(self._processes.values())
while not all(p.is_alive() for p in processes):
time.sleep(1e-3)
def _wait_job_completion(self):
"""Wait for the cache to be empty before resizing the pool."""
# Issue a warning to the user about the bad effect of this usage.
if self._pending_work_items:
warnings.warn(
"Trying to resize an executor with running jobs: "
"waiting for jobs completion before resizing.",
UserWarning,
)
mp.util.debug(
f"Executor {self.executor_id} waiting for jobs completion "
"before resizing"
)
# Wait for the completion of the jobs
while self._pending_work_items:
time.sleep(1e-3)
def _setup_queues(self, job_reducers, result_reducers):
# As this executor can be resized, use a large queue size to avoid
# underestimating capacity and introducing overhead
# Also handle the case where the user set max_workers to a value larger
# than cpu_count(), to avoid limiting the number of parallel jobs.
min_queue_size = max(cpu_count(), self._max_workers)
self.queue_size = 2 * min_queue_size + EXTRA_QUEUED_CALLS
super()._setup_queues(
job_reducers, result_reducers, queue_size=self.queue_size
)