PK aZZZ)�O joblib/__init__.py"""Joblib is a set of tools to provide **lightweight pipelining in
Python**. In particular:
1. transparent disk-caching of functions and lazy re-evaluation
(memoize pattern)
2. easy simple parallel computing
Joblib is optimized to be **fast** and **robust** on large
data in particular and has specific optimizations for `numpy` arrays. It is
**BSD-licensed**.
==================== ===============================================
**Documentation:** https://joblib.readthedocs.io
**Download:** https://pypi.python.org/pypi/joblib#downloads
**Source code:** https://github.com/joblib/joblib
**Report issues:** https://github.com/joblib/joblib/issues
==================== ===============================================
Vision
--------
The vision is to provide tools to easily achieve better performance and
reproducibility when working with long running jobs.
* **Avoid computing the same thing twice**: code is often rerun again and
again, for instance when prototyping computational-heavy jobs (as in
scientific development), but hand-crafted solutions to alleviate this
issue are error-prone and often lead to unreproducible results.
* **Persist to disk transparently**: efficiently persisting
arbitrary objects containing large data is hard. Using
joblib's caching mechanism avoids hand-written persistence and
implicitly links the file on disk to the execution context of
the original Python object. As a result, joblib's persistence is
good for resuming an application status or computational job, eg
after a crash.
Joblib addresses these problems while **leaving your code and your flow
control as unmodified as possible** (no framework, no new paradigms).
Main features
------------------
1) **Transparent and fast disk-caching of output value:** a memoize or
make-like functionality for Python functions that works well for
arbitrary Python objects, including very large numpy arrays. Separate
persistence and flow-execution logic from domain logic or algorithmic
code by writing the operations as a set of steps with well-defined
inputs and outputs: Python functions. Joblib can save their
computation to disk and rerun it only if necessary::
>>> from joblib import Memory
>>> cachedir = 'your_cache_dir_goes_here'
>>> mem = Memory(cachedir)
>>> import numpy as np
>>> a = np.vander(np.arange(3)).astype(float)
>>> square = mem.cache(np.square)
>>> b = square(a) # doctest: +ELLIPSIS
______________________________________________________________________...
[Memory] Calling square...
square(array([[0., 0., 1.],
[1., 1., 1.],
[4., 2., 1.]]))
_________________________________________________...square - ...s, 0.0min
>>> c = square(a)
>>> # The above call did not trigger an evaluation
2) **Embarrassingly parallel helper:** to make it easy to write readable
parallel code and debug it quickly::
>>> from joblib import Parallel, delayed
>>> from math import sqrt
>>> Parallel(n_jobs=1)(delayed(sqrt)(i**2) for i in range(10))
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
3) **Fast compressed Persistence**: a replacement for pickle to work
efficiently on Python objects containing large data (
*joblib.dump* & *joblib.load* ).
..
>>> import shutil ; shutil.rmtree(cachedir)
"""
# PEP0440 compatible formatted version, see:
# https://www.python.org/dev/peps/pep-0440/
#
# Generic release markers:
# X.Y
# X.Y.Z # For bugfix releases
#
# Admissible pre-release markers:
# X.YaN # Alpha release
# X.YbN # Beta release
# X.YrcN # Release Candidate
# X.Y # Final release
#
# Dev branch marker is: 'X.Y.dev' or 'X.Y.devN' where N is an integer.
# 'X.Y.dev0' is the canonical version of 'X.Y.dev'
#
__version__ = '1.4.0'
import os
from .memory import Memory
from .memory import MemorizedResult
from .memory import register_store_backend
from .memory import expires_after
from .logger import PrintTime
from .logger import Logger
from .hashing import hash
from .numpy_pickle import dump
from .numpy_pickle import load
from .compressor import register_compressor
from .parallel import Parallel
from .parallel import delayed
from .parallel import cpu_count
from .parallel import register_parallel_backend
from .parallel import parallel_backend
from .parallel import parallel_config
from .parallel import effective_n_jobs
from ._cloudpickle_wrapper import wrap_non_picklable_objects
__all__ = ['Memory', 'MemorizedResult', 'PrintTime', 'Logger', 'hash', 'dump',
'load', 'Parallel', 'delayed', 'cpu_count', 'effective_n_jobs',
'register_parallel_backend', 'parallel_backend', 'expires_after',
'register_store_backend', 'register_compressor',
'wrap_non_picklable_objects', 'parallel_config']
# Workaround issue discovered in intel-openmp 2019.5:
# https://github.com/ContinuumIO/anaconda-issues/issues/11294
os.environ.setdefault("KMP_INIT_AT_FORK", "FALSE")
PK aZZZGy� � joblib/_cloudpickle_wrapper.py"""
Small shim of loky's cloudpickle_wrapper to avoid failure when
multiprocessing is not available.
"""
from ._multiprocessing_helpers import mp
def _my_wrap_non_picklable_objects(obj, keep_wrapper=True):
return obj
if mp is not None:
from .externals.loky import wrap_non_picklable_objects
else:
wrap_non_picklable_objects = _my_wrap_non_picklable_objects
__all__ = ["wrap_non_picklable_objects"]
PK aZZZ��V4 4 joblib/_dask.pyfrom __future__ import print_function, division, absolute_import
import asyncio
import concurrent.futures
import contextlib
import time
from uuid import uuid4
import weakref
from .parallel import parallel_config
from .parallel import AutoBatchingMixin, ParallelBackendBase
from ._utils import (
_TracebackCapturingWrapper,
_retrieve_traceback_capturing_wrapped_call
)
try:
import dask
import distributed
except ImportError:
dask = None
distributed = None
if dask is not None and distributed is not None:
from dask.utils import funcname
from dask.sizeof import sizeof
from dask.distributed import (
Client,
as_completed,
get_client,
secede,
rejoin,
)
from distributed.utils import thread_state
try:
# asyncio.TimeoutError, Python3-only error thrown by recent versions of
# distributed
from distributed.utils import TimeoutError as _TimeoutError
except ImportError:
from tornado.gen import TimeoutError as _TimeoutError
def is_weakrefable(obj):
try:
weakref.ref(obj)
return True
except TypeError:
return False
class _WeakKeyDictionary:
"""A variant of weakref.WeakKeyDictionary for unhashable objects.
This datastructure is used to store futures for broadcasted data objects
such as large numpy arrays or pandas dataframes that are not hashable and
therefore cannot be used as keys of traditional python dicts.
Furthermore using a dict with id(array) as key is not safe because the
Python is likely to reuse id of recently collected arrays.
"""
def __init__(self):
self._data = {}
def __getitem__(self, obj):
ref, val = self._data[id(obj)]
if ref() is not obj:
# In case of a race condition with on_destroy.
raise KeyError(obj)
return val
def __setitem__(self, obj, value):
key = id(obj)
try:
ref, _ = self._data[key]
if ref() is not obj:
# In case of race condition with on_destroy.
raise KeyError(obj)
except KeyError:
# Insert the new entry in the mapping along with a weakref
# callback to automatically delete the entry from the mapping
# as soon as the object used as key is garbage collected.
def on_destroy(_):
del self._data[key]
ref = weakref.ref(obj, on_destroy)
self._data[key] = ref, value
def __len__(self):
return len(self._data)
def clear(self):
self._data.clear()
def _funcname(x):
try:
if isinstance(x, list):
x = x[0][0]
except Exception:
pass
return funcname(x)
def _make_tasks_summary(tasks):
"""Summarize of list of (func, args, kwargs) function calls"""
unique_funcs = {func for func, args, kwargs in tasks}
if len(unique_funcs) == 1:
mixed = False
else:
mixed = True
return len(tasks), mixed, _funcname(tasks)
class Batch:
"""dask-compatible wrapper that executes a batch of tasks"""
def __init__(self, tasks):
# collect some metadata from the tasks to ease Batch calls
# introspection when debugging
self._num_tasks, self._mixed, self._funcname = _make_tasks_summary(
tasks
)
def __call__(self, tasks=None):
results = []
with parallel_config(backend='dask'):
for func, args, kwargs in tasks:
results.append(func(*args, **kwargs))
return results
def __repr__(self):
descr = f"batch_of_{self._funcname}_{self._num_tasks}_calls"
if self._mixed:
descr = "mixed_" + descr
return descr
def _joblib_probe_task():
# Noop used by the joblib connector to probe when workers are ready.
pass
class DaskDistributedBackend(AutoBatchingMixin, ParallelBackendBase):
MIN_IDEAL_BATCH_DURATION = 0.2
MAX_IDEAL_BATCH_DURATION = 1.0
supports_retrieve_callback = True
default_n_jobs = -1
def __init__(self, scheduler_host=None, scatter=None,
client=None, loop=None, wait_for_workers_timeout=10,
**submit_kwargs):
super().__init__()
if distributed is None:
msg = ("You are trying to use 'dask' as a joblib parallel backend "
"but dask is not installed. Please install dask "
"to fix this error.")
raise ValueError(msg)
if client is None:
if scheduler_host:
client = Client(scheduler_host, loop=loop,
set_as_default=False)
else:
try:
client = get_client()
except ValueError as e:
msg = ("To use Joblib with Dask first create a Dask Client"
"\n\n"
" from dask.distributed import Client\n"
" client = Client()\n"
"or\n"
" client = Client('scheduler-address:8786')")
raise ValueError(msg) from e
self.client = client
if scatter is not None and not isinstance(scatter, (list, tuple)):
raise TypeError("scatter must be a list/tuple, got "
"`%s`" % type(scatter).__name__)
if scatter is not None and len(scatter) > 0:
# Keep a reference to the scattered data to keep the ids the same
self._scatter = list(scatter)
scattered = self.client.scatter(scatter, broadcast=True)
self.data_futures = {id(x): f for x, f in zip(scatter, scattered)}
else:
self._scatter = []
self.data_futures = {}
self.wait_for_workers_timeout = wait_for_workers_timeout
self.submit_kwargs = submit_kwargs
self.waiting_futures = as_completed(
[],
loop=client.loop,
with_results=True,
raise_errors=False
)
self._results = {}
self._callbacks = {}
async def _collect(self):
while self._continue:
async for future, result in self.waiting_futures:
cf_future = self._results.pop(future)
callback = self._callbacks.pop(future)
if future.status == "error":
typ, exc, tb = result
cf_future.set_exception(exc)
else:
cf_future.set_result(result)
callback(result)
await asyncio.sleep(0.01)
def __reduce__(self):
return (DaskDistributedBackend, ())
def get_nested_backend(self):
return DaskDistributedBackend(client=self.client), -1
def configure(self, n_jobs=1, parallel=None, **backend_args):
self.parallel = parallel
return self.effective_n_jobs(n_jobs)
def start_call(self):
self._continue = True
self.client.loop.add_callback(self._collect)
self.call_data_futures = _WeakKeyDictionary()
def stop_call(self):
# The explicit call to clear is required to break a cycling reference
# to the futures.
self._continue = False
# wait for the future collection routine (self._backend._collect) to
# finish in order to limit asyncio warnings due to aborting _collect
# during a following backend termination call
time.sleep(0.01)
self.call_data_futures.clear()
def effective_n_jobs(self, n_jobs):
effective_n_jobs = sum(self.client.ncores().values())
if effective_n_jobs != 0 or not self.wait_for_workers_timeout:
return effective_n_jobs
# If there is no worker, schedule a probe task to wait for the workers
# to come up and be available. If the dask cluster is in adaptive mode
# task might cause the cluster to provision some workers.
try:
self.client.submit(_joblib_probe_task).result(
timeout=self.wait_for_workers_timeout
)
except _TimeoutError as e:
error_msg = (
"DaskDistributedBackend has no worker after {} seconds. "
"Make sure that workers are started and can properly connect "
"to the scheduler and increase the joblib/dask connection "
"timeout with:\n\n"
"parallel_config(backend='dask', wait_for_workers_timeout={})"
).format(self.wait_for_workers_timeout,
max(10, 2 * self.wait_for_workers_timeout))
raise TimeoutError(error_msg) from e
return sum(self.client.ncores().values())
async def _to_func_args(self, func):
itemgetters = dict()
# Futures that are dynamically generated during a single call to
# Parallel.__call__.
call_data_futures = getattr(self, 'call_data_futures', None)
async def maybe_to_futures(args):
out = []
for arg in args:
arg_id = id(arg)
if arg_id in itemgetters:
out.append(itemgetters[arg_id])
continue
f = self.data_futures.get(arg_id, None)
if f is None and call_data_futures is not None:
try:
f = await call_data_futures[arg]
except KeyError:
pass
if f is None:
if is_weakrefable(arg) and sizeof(arg) > 1e3:
# Automatically scatter large objects to some of
# the workers to avoid duplicated data transfers.
# Rely on automated inter-worker data stealing if
# more workers need to reuse this data
# concurrently.
# set hash=False - nested scatter calls (i.e
# calling client.scatter inside a dask worker)
# using hash=True often raise CancelledError,
# see dask/distributed#3703
_coro = self.client.scatter(
arg,
asynchronous=True,
hash=False
)
# Centralize the scattering of identical arguments
# between concurrent apply_async callbacks by
# exposing the running coroutine in
# call_data_futures before it completes.
t = asyncio.Task(_coro)
call_data_futures[arg] = t
f = await t
if f is not None:
out.append(f)
else:
out.append(arg)
return out
tasks = []
for f, args, kwargs in func.items:
args = list(await maybe_to_futures(args))
kwargs = dict(zip(kwargs.keys(),
await maybe_to_futures(kwargs.values())))
tasks.append((f, args, kwargs))
return (Batch(tasks), tasks)
def apply_async(self, func, callback=None):
cf_future = concurrent.futures.Future()
cf_future.get = cf_future.result # achieve AsyncResult API
async def f(func, callback):
batch, tasks = await self._to_func_args(func)
key = f'{repr(batch)}-{uuid4().hex}'
dask_future = self.client.submit(
_TracebackCapturingWrapper(batch),
tasks=tasks,
key=key,
**self.submit_kwargs
)
self.waiting_futures.add(dask_future)
self._callbacks[dask_future] = callback
self._results[dask_future] = cf_future
self.client.loop.add_callback(f, func, callback)
return cf_future
def retrieve_result_callback(self, out):
return _retrieve_traceback_capturing_wrapped_call(out)
def abort_everything(self, ensure_ready=True):
""" Tell the client to cancel any task submitted via this instance
joblib.Parallel will never access those results
"""
with self.waiting_futures.lock:
self.waiting_futures.futures.clear()
while not self.waiting_futures.queue.empty():
self.waiting_futures.queue.get()
@contextlib.contextmanager
def retrieval_context(self):
"""Override ParallelBackendBase.retrieval_context to avoid deadlocks.
This removes thread from the worker's thread pool (using 'secede').
Seceding avoids deadlock in nested parallelism settings.
"""
# See 'joblib.Parallel.__call__' and 'joblib.Parallel.retrieve' for how
# this is used.
if hasattr(thread_state, 'execution_state'):
# we are in a worker. Secede to avoid deadlock.
secede()
yield
if hasattr(thread_state, 'execution_state'):
rejoin()
PK aZZZ~�`�m �m joblib/_memmapping_reducer.py"""
Reducer using memory mapping for numpy arrays
"""
# Author: Thomas Moreau <thomas.moreau.2010@gmail.com>
# Copyright: 2017, Thomas Moreau
# License: BSD 3 clause
from mmap import mmap
import errno
import os
import stat
import threading
import atexit
import tempfile
import time
import warnings
import weakref
from uuid import uuid4
from multiprocessing import util
from pickle import whichmodule, loads, dumps, HIGHEST_PROTOCOL, PicklingError
try:
WindowsError
except NameError:
WindowsError = type(None)
try:
import numpy as np
from numpy.lib.stride_tricks import as_strided
except ImportError:
np = None
from .numpy_pickle import dump, load, load_temporary_memmap
from .backports import make_memmap
from .disk import delete_folder
from .externals.loky.backend import resource_tracker
# Some system have a ramdisk mounted by default, we can use it instead of /tmp
# as the default folder to dump big arrays to share with subprocesses.
SYSTEM_SHARED_MEM_FS = '/dev/shm'
# Minimal number of bytes available on SYSTEM_SHARED_MEM_FS to consider using
# it as the default folder to dump big arrays to share with subprocesses.
SYSTEM_SHARED_MEM_FS_MIN_SIZE = int(2e9)
# Folder and file permissions to chmod temporary files generated by the
# memmapping pool. Only the owner of the Python process can access the
# temporary files and folder.
FOLDER_PERMISSIONS = stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR
FILE_PERMISSIONS = stat.S_IRUSR | stat.S_IWUSR
# Set used in joblib workers, referencing the filenames of temporary memmaps
# created by joblib to speed up data communication. In child processes, we add
# a finalizer to these memmaps that sends a maybe_unlink call to the
# resource_tracker, in order to free main memory as fast as possible.
JOBLIB_MMAPS = set()
def _log_and_unlink(filename):
from .externals.loky.backend.resource_tracker import _resource_tracker
util.debug(
"[FINALIZER CALL] object mapping to {} about to be deleted,"
" decrementing the refcount of the file (pid: {})".format(
os.path.basename(filename), os.getpid()))
_resource_tracker.maybe_unlink(filename, "file")
def add_maybe_unlink_finalizer(memmap):
util.debug(
"[FINALIZER ADD] adding finalizer to {} (id {}, filename {}, pid {})"
"".format(type(memmap), id(memmap), os.path.basename(memmap.filename),
os.getpid()))
weakref.finalize(memmap, _log_and_unlink, memmap.filename)
def unlink_file(filename):
"""Wrapper around os.unlink with a retry mechanism.
The retry mechanism has been implemented primarily to overcome a race
condition happening during the finalizer of a np.memmap: when a process
holding the last reference to a mmap-backed np.memmap/np.array is about to
delete this array (and close the reference), it sends a maybe_unlink
request to the resource_tracker. This request can be processed faster than
it takes for the last reference of the memmap to be closed, yielding (on
Windows) a PermissionError in the resource_tracker loop.
"""
NUM_RETRIES = 10
for retry_no in range(1, NUM_RETRIES + 1):
try:
os.unlink(filename)
break
except PermissionError:
util.debug(
'[ResourceTracker] tried to unlink {}, got '
'PermissionError'.format(filename)
)
if retry_no == NUM_RETRIES:
raise
else:
time.sleep(.2)
except FileNotFoundError:
# In case of a race condition when deleting the temporary folder,
# avoid noisy FileNotFoundError exception in the resource tracker.
pass
resource_tracker._CLEANUP_FUNCS['file'] = unlink_file
class _WeakArrayKeyMap:
"""A variant of weakref.WeakKeyDictionary for unhashable numpy arrays.
This datastructure will be used with numpy arrays as obj keys, therefore we
do not use the __get__ / __set__ methods to avoid any conflict with the
numpy fancy indexing syntax.
"""
def __init__(self):
self._data = {}
def get(self, obj):
ref, val = self._data[id(obj)]
if ref() is not obj:
# In case of race condition with on_destroy: could never be
# triggered by the joblib tests with CPython.
raise KeyError(obj)
return val
def set(self, obj, value):
key = id(obj)
try:
ref, _ = self._data[key]
if ref() is not obj:
# In case of race condition with on_destroy: could never be
# triggered by the joblib tests with CPython.
raise KeyError(obj)
except KeyError:
# Insert the new entry in the mapping along with a weakref
# callback to automatically delete the entry from the mapping
# as soon as the object used as key is garbage collected.
def on_destroy(_):
del self._data[key]
ref = weakref.ref(obj, on_destroy)
self._data[key] = ref, value
def __getstate__(self):
raise PicklingError("_WeakArrayKeyMap is not pickleable")
###############################################################################
# Support for efficient transient pickling of numpy data structures
def _get_backing_memmap(a):
"""Recursively look up the original np.memmap instance base if any."""
b = getattr(a, 'base', None)
if b is None:
# TODO: check scipy sparse datastructure if scipy is installed
# a nor its descendants do not have a memmap base
return None
elif isinstance(b, mmap):
# a is already a real memmap instance.
return a
else:
# Recursive exploration of the base ancestry
return _get_backing_memmap(b)
def _get_temp_dir(pool_folder_name, temp_folder=None):
"""Get the full path to a subfolder inside the temporary folder.
Parameters
----------
pool_folder_name : str
Sub-folder name used for the serialization of a pool instance.
temp_folder: str, optional
Folder to be used by the pool for memmapping large arrays
for sharing memory with worker processes. If None, this will try in
order:
- a folder pointed by the JOBLIB_TEMP_FOLDER environment
variable,
- /dev/shm if the folder exists and is writable: this is a
RAMdisk filesystem available by default on modern Linux
distributions,
- the default system temporary folder that can be
overridden with TMP, TMPDIR or TEMP environment
variables, typically /tmp under Unix operating systems.
Returns
-------
pool_folder : str
full path to the temporary folder
use_shared_mem : bool
whether the temporary folder is written to the system shared memory
folder or some other temporary folder.
"""
use_shared_mem = False
if temp_folder is None:
temp_folder = os.environ.get('JOBLIB_TEMP_FOLDER', None)
if temp_folder is None:
if os.path.exists(SYSTEM_SHARED_MEM_FS) and hasattr(os, 'statvfs'):
try:
shm_stats = os.statvfs(SYSTEM_SHARED_MEM_FS)
available_nbytes = shm_stats.f_bsize * shm_stats.f_bavail
if available_nbytes > SYSTEM_SHARED_MEM_FS_MIN_SIZE:
# Try to see if we have write access to the shared mem
# folder only if it is reasonably large (that is 2GB or
# more).
temp_folder = SYSTEM_SHARED_MEM_FS
pool_folder = os.path.join(temp_folder, pool_folder_name)
if not os.path.exists(pool_folder):
os.makedirs(pool_folder)
use_shared_mem = True
except (IOError, OSError):
# Missing rights in the /dev/shm partition, fallback to regular
# temp folder.
temp_folder = None
if temp_folder is None:
# Fallback to the default tmp folder, typically /tmp
temp_folder = tempfile.gettempdir()
temp_folder = os.path.abspath(os.path.expanduser(temp_folder))
pool_folder = os.path.join(temp_folder, pool_folder_name)
return pool_folder, use_shared_mem
def has_shareable_memory(a):
"""Return True if a is backed by some mmap buffer directly or not."""
return _get_backing_memmap(a) is not None
def _strided_from_memmap(filename, dtype, mode, offset, order, shape, strides,
total_buffer_len, unlink_on_gc_collect):
"""Reconstruct an array view on a memory mapped file."""
if mode == 'w+':
# Do not zero the original data when unpickling
mode = 'r+'
if strides is None:
# Simple, contiguous memmap
return make_memmap(
filename, dtype=dtype, shape=shape, mode=mode, offset=offset,
order=order, unlink_on_gc_collect=unlink_on_gc_collect
)
else:
# For non-contiguous data, memmap the total enclosing buffer and then
# extract the non-contiguous view with the stride-tricks API
base = make_memmap(
filename, dtype=dtype, shape=total_buffer_len, offset=offset,
mode=mode, order=order, unlink_on_gc_collect=unlink_on_gc_collect
)
return as_strided(base, shape=shape, strides=strides)
def _reduce_memmap_backed(a, m):
"""Pickling reduction for memmap backed arrays.
a is expected to be an instance of np.ndarray (or np.memmap)
m is expected to be an instance of np.memmap on the top of the ``base``
attribute ancestry of a. ``m.base`` should be the real python mmap object.
"""
# offset that comes from the striding differences between a and m
util.debug('[MEMMAP REDUCE] reducing a memmap-backed array '
'(shape, {}, pid: {})'.format(a.shape, os.getpid()))
try:
from numpy.lib.array_utils import byte_bounds
except (ModuleNotFoundError, ImportError):
# Backward-compat for numpy < 2.0
from numpy import byte_bounds
a_start, a_end = byte_bounds(a)
m_start = byte_bounds(m)[0]
offset = a_start - m_start
# offset from the backing memmap
offset += m.offset
if m.flags['F_CONTIGUOUS']:
order = 'F'
else:
# The backing memmap buffer is necessarily contiguous hence C if not
# Fortran
order = 'C'
if a.flags['F_CONTIGUOUS'] or a.flags['C_CONTIGUOUS']:
# If the array is a contiguous view, no need to pass the strides
strides = None
total_buffer_len = None
else:
# Compute the total number of items to map from which the strided
# view will be extracted.
strides = a.strides
total_buffer_len = (a_end - a_start) // a.itemsize
return (_strided_from_memmap,
(m.filename, a.dtype, m.mode, offset, order, a.shape, strides,
total_buffer_len, False))
def reduce_array_memmap_backward(a):
"""reduce a np.array or a np.memmap from a child process"""
m = _get_backing_memmap(a)
if isinstance(m, np.memmap) and m.filename not in JOBLIB_MMAPS:
# if a is backed by a memmaped file, reconstruct a using the
# memmaped file.
return _reduce_memmap_backed(a, m)
else:
# a is either a regular (not memmap-backed) numpy array, or an array
# backed by a shared temporary file created by joblib. In the latter
# case, in order to limit the lifespan of these temporary files, we
# serialize the memmap as a regular numpy array, and decref the
# file backing the memmap (done implicitly in a previously registered
# finalizer, see ``unlink_on_gc_collect`` for more details)
return (
loads, (dumps(np.asarray(a), protocol=HIGHEST_PROTOCOL), )
)
class ArrayMemmapForwardReducer(object):
"""Reducer callable to dump large arrays to memmap files.
Parameters
----------
max_nbytes: int
Threshold to trigger memmapping of large arrays to files created
a folder.
temp_folder_resolver: callable
An callable in charge of resolving a temporary folder name where files
for backing memmapped arrays are created.
mmap_mode: 'r', 'r+' or 'c'
Mode for the created memmap datastructure. See the documentation of
numpy.memmap for more details. Note: 'w+' is coerced to 'r+'
automatically to avoid zeroing the data on unpickling.
verbose: int, optional, 0 by default
If verbose > 0, memmap creations are logged.
If verbose > 1, both memmap creations, reuse and array pickling are
logged.
prewarm: bool, optional, False by default.
Force a read on newly memmapped array to make sure that OS pre-cache it
memory. This can be useful to avoid concurrent disk access when the
same data array is passed to different worker processes.
"""
def __init__(self, max_nbytes, temp_folder_resolver, mmap_mode,
unlink_on_gc_collect, verbose=0, prewarm=True):
self._max_nbytes = max_nbytes
self._temp_folder_resolver = temp_folder_resolver
self._mmap_mode = mmap_mode
self.verbose = int(verbose)
if prewarm == "auto":
self._prewarm = not self._temp_folder.startswith(
SYSTEM_SHARED_MEM_FS
)
else:
self._prewarm = prewarm
self._prewarm = prewarm
self._memmaped_arrays = _WeakArrayKeyMap()
self._temporary_memmaped_filenames = set()
self._unlink_on_gc_collect = unlink_on_gc_collect
@property
def _temp_folder(self):
return self._temp_folder_resolver()
def __reduce__(self):
# The ArrayMemmapForwardReducer is passed to the children processes: it
# needs to be pickled but the _WeakArrayKeyMap need to be skipped as
# it's only guaranteed to be consistent with the parent process memory
# garbage collection.
# Although this reducer is pickled, it is not needed in its destination
# process (child processes), as we only use this reducer to send
# memmaps from the parent process to the children processes. For this
# reason, we can afford skipping the resolver, (which would otherwise
# be unpicklable), and pass it as None instead.
args = (self._max_nbytes, None, self._mmap_mode,
self._unlink_on_gc_collect)
kwargs = {
'verbose': self.verbose,
'prewarm': self._prewarm,
}
return ArrayMemmapForwardReducer, args, kwargs
def __call__(self, a):
m = _get_backing_memmap(a)
if m is not None and isinstance(m, np.memmap):
# a is already backed by a memmap file, let's reuse it directly
return _reduce_memmap_backed(a, m)
if (not a.dtype.hasobject and self._max_nbytes is not None and
a.nbytes > self._max_nbytes):
# check that the folder exists (lazily create the pool temp folder
# if required)
try:
os.makedirs(self._temp_folder)
os.chmod(self._temp_folder, FOLDER_PERMISSIONS)
except OSError as e:
if e.errno != errno.EEXIST:
raise e
try:
basename = self._memmaped_arrays.get(a)
except KeyError:
# Generate a new unique random filename. The process and thread
# ids are only useful for debugging purpose and to make it
# easier to cleanup orphaned files in case of hard process
# kill (e.g. by "kill -9" or segfault).
basename = "{}-{}-{}.pkl".format(
os.getpid(), id(threading.current_thread()), uuid4().hex)
self._memmaped_arrays.set(a, basename)
filename = os.path.join(self._temp_folder, basename)
# In case the same array with the same content is passed several
# times to the pool subprocess children, serialize it only once
is_new_memmap = filename not in self._temporary_memmaped_filenames
# add the memmap to the list of temporary memmaps created by joblib
self._temporary_memmaped_filenames.add(filename)
if self._unlink_on_gc_collect:
# Bump reference count of the memmap by 1 to account for
# shared usage of the memmap by a child process. The
# corresponding decref call will be executed upon calling
# resource_tracker.maybe_unlink, registered as a finalizer in
# the child.
# the incref/decref calls here are only possible when the child
# and the parent share the same resource_tracker. It is not the
# case for the multiprocessing backend, but it does not matter
# because unlinking a memmap from a child process is only
# useful to control the memory usage of long-lasting child
# processes, while the multiprocessing-based pools terminate
# their workers at the end of a map() call.
resource_tracker.register(filename, "file")
if is_new_memmap:
# Incref each temporary memmap created by joblib one extra
# time. This means that these memmaps will only be deleted
# once an extra maybe_unlink() is called, which is done once
# all the jobs have completed (or been canceled) in the
# Parallel._terminate_backend() method.
resource_tracker.register(filename, "file")
if not os.path.exists(filename):
util.debug(
"[ARRAY DUMP] Pickling new array (shape={}, dtype={}) "
"creating a new memmap at {}".format(
a.shape, a.dtype, filename))
for dumped_filename in dump(a, filename):
os.chmod(dumped_filename, FILE_PERMISSIONS)
if self._prewarm:
# Warm up the data by accessing it. This operation ensures
# that the disk access required to create the memmapping
# file are performed in the reducing process and avoids
# concurrent memmap creation in multiple children
# processes.
load(filename, mmap_mode=self._mmap_mode).max()
else:
util.debug(
"[ARRAY DUMP] Pickling known array (shape={}, dtype={}) "
"reusing memmap file: {}".format(
a.shape, a.dtype, os.path.basename(filename)))
# The worker process will use joblib.load to memmap the data
return (
(load_temporary_memmap, (filename, self._mmap_mode,
self._unlink_on_gc_collect))
)
else:
# do not convert a into memmap, let pickler do its usual copy with
# the default system pickler
util.debug(
'[ARRAY DUMP] Pickling array (NO MEMMAPPING) (shape={}, '
' dtype={}).'.format(a.shape, a.dtype))
return (loads, (dumps(a, protocol=HIGHEST_PROTOCOL),))
def get_memmapping_reducers(
forward_reducers=None, backward_reducers=None,
temp_folder_resolver=None, max_nbytes=1e6, mmap_mode='r', verbose=0,
prewarm=False, unlink_on_gc_collect=True, **kwargs):
"""Construct a pair of memmapping reducer linked to a tmpdir.
This function manage the creation and the clean up of the temporary folders
underlying the memory maps and should be use to get the reducers necessary
to construct joblib pool or executor.
"""
if forward_reducers is None:
forward_reducers = dict()
if backward_reducers is None:
backward_reducers = dict()
if np is not None:
# Register smart numpy.ndarray reducers that detects memmap backed
# arrays and that is also able to dump to memmap large in-memory
# arrays over the max_nbytes threshold
forward_reduce_ndarray = ArrayMemmapForwardReducer(
max_nbytes, temp_folder_resolver, mmap_mode, unlink_on_gc_collect,
verbose, prewarm=prewarm)
forward_reducers[np.ndarray] = forward_reduce_ndarray
forward_reducers[np.memmap] = forward_reduce_ndarray
# Communication from child process to the parent process always
# pickles in-memory numpy.ndarray without dumping them as memmap
# to avoid confusing the caller and make it tricky to collect the
# temporary folder
backward_reducers[np.ndarray] = reduce_array_memmap_backward
backward_reducers[np.memmap] = reduce_array_memmap_backward
return forward_reducers, backward_reducers
class TemporaryResourcesManager(object):
"""Stateful object able to manage temporary folder and pickles
It exposes:
- a per-context folder name resolving API that memmap-based reducers will
rely on to know where to pickle the temporary memmaps
- a temporary file/folder management API that internally uses the
resource_tracker.
"""
def __init__(self, temp_folder_root=None, context_id=None):
self._current_temp_folder = None
self._temp_folder_root = temp_folder_root
self._use_shared_mem = None
self._cached_temp_folders = dict()
self._id = uuid4().hex
self._finalizers = {}
if context_id is None:
# It would be safer to not assign a default context id (less silent
# bugs), but doing this while maintaining backward compatibility
# with the previous, context-unaware version get_memmaping_executor
# exposes too many low-level details.
context_id = uuid4().hex
self.set_current_context(context_id)
def set_current_context(self, context_id):
self._current_context_id = context_id
self.register_new_context(context_id)
def register_new_context(self, context_id):
# Prepare a sub-folder name specific to a context (usually a unique id
# generated by each instance of the Parallel class). Do not create in
# advance to spare FS write access if no array is to be dumped).
if context_id in self._cached_temp_folders:
return
else:
# During its lifecycle, one Parallel object can have several
# executors associated to it (for instance, if a loky worker raises
# an exception, joblib shutdowns the executor and instantly
# recreates a new one before raising the error - see
# ``ensure_ready``. Because we don't want two executors tied to
# the same Parallel object (and thus the same context id) to
# register/use/delete the same folder, we also add an id specific
# to the current Manager (and thus specific to its associated
# executor) to the folder name.
new_folder_name = (
"joblib_memmapping_folder_{}_{}_{}".format(
os.getpid(), self._id, context_id)
)
new_folder_path, _ = _get_temp_dir(
new_folder_name, self._temp_folder_root
)
self.register_folder_finalizer(new_folder_path, context_id)
self._cached_temp_folders[context_id] = new_folder_path
def resolve_temp_folder_name(self):
"""Return a folder name specific to the currently activated context"""
return self._cached_temp_folders[self._current_context_id]
# resource management API
def register_folder_finalizer(self, pool_subfolder, context_id):
# Register the garbage collector at program exit in case caller forgets
# to call terminate explicitly: note we do not pass any reference to
# ensure that this callback won't prevent garbage collection of
# parallel instance and related file handler resources such as POSIX
# semaphores and pipes
pool_module_name = whichmodule(delete_folder, 'delete_folder')
resource_tracker.register(pool_subfolder, "folder")
def _cleanup():
# In some cases the Python runtime seems to set delete_folder to
# None just before exiting when accessing the delete_folder
# function from the closure namespace. So instead we reimport
# the delete_folder function explicitly.
# https://github.com/joblib/joblib/issues/328
# We cannot just use from 'joblib.pool import delete_folder'
# because joblib should only use relative imports to allow
# easy vendoring.
delete_folder = __import__(
pool_module_name, fromlist=['delete_folder']
).delete_folder
try:
delete_folder(pool_subfolder, allow_non_empty=True)
resource_tracker.unregister(pool_subfolder, "folder")
except OSError:
warnings.warn("Failed to delete temporary folder: {}"
.format(pool_subfolder))
self._finalizers[context_id] = atexit.register(_cleanup)
def _clean_temporary_resources(self, context_id=None, force=False,
allow_non_empty=False):
"""Clean temporary resources created by a process-based pool"""
if context_id is None:
# Iterates over a copy of the cache keys to avoid Error due to
# iterating over a changing size dictionary.
for context_id in list(self._cached_temp_folders):
self._clean_temporary_resources(
context_id, force=force, allow_non_empty=allow_non_empty
)
else:
temp_folder = self._cached_temp_folders.get(context_id)
if temp_folder and os.path.exists(temp_folder):
for filename in os.listdir(temp_folder):
if force:
# Some workers have failed and the ref counted might
# be off. The workers should have shut down by this
# time so forcefully clean up the files.
resource_tracker.unregister(
os.path.join(temp_folder, filename), "file"
)
else:
resource_tracker.maybe_unlink(
os.path.join(temp_folder, filename), "file"
)
# When forcing clean-up, try to delete the folder even if some
# files are still in it. Otherwise, try to delete the folder
allow_non_empty |= force
# Clean up the folder if possible, either if it is empty or
# if none of the files in it are in used and allow_non_empty.
try:
delete_folder(
temp_folder, allow_non_empty=allow_non_empty
)
# Forget the folder once it has been deleted
self._cached_temp_folders.pop(context_id, None)
resource_tracker.unregister(temp_folder, "folder")
# Also cancel the finalizers that gets triggered at gc.
finalizer = self._finalizers.pop(context_id, None)
if finalizer is not None:
atexit.unregister(finalizer)
except OSError:
# Temporary folder cannot be deleted right now.
# This folder will be cleaned up by an atexit
# finalizer registered by the memmapping_reducer.
pass
PK aZZZ��1�� � "