ENH: update joblib
This commit is contained in:
parent
1543ba214c
commit
54fcfc2da3
|
|
@ -38,14 +38,15 @@ solution.
|
|||
issue is error-prone and often leads to unreproducible results
|
||||
|
||||
* **Persist to disk transparently**: persisting in an efficient way
|
||||
arbitrary objects containing large data is hard. In addition,
|
||||
hand-written persistence does not link easily the file on disk to the
|
||||
execution context of the original Python object. As a result, it is
|
||||
challenging to resume a application status or computational job, eg
|
||||
arbitrary objects containing large data is hard. Using
|
||||
joblib's caching mechanism avoids hand-written persistence and
|
||||
implicitely 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.
|
||||
|
||||
It strives to address these problems while **leaving your code and your
|
||||
flow control as unmodified as possible** (no framework, no new
|
||||
Joblib strives to address these problems while **leaving your code and
|
||||
your flow control as unmodified as possible** (no framework, no new
|
||||
paradigms).
|
||||
|
||||
Main features
|
||||
|
|
@ -76,7 +77,7 @@ Main features
|
|||
>>> # The above call did not trigger an evaluation
|
||||
|
||||
2) **Embarrassingly parallel helper:** to make is easy to write readable
|
||||
parallel code and debug it quickly:
|
||||
parallel code and debug it quickly::
|
||||
|
||||
>>> from sklearn.externals.joblib import Parallel, delayed
|
||||
>>> from math import sqrt
|
||||
|
|
@ -91,12 +92,16 @@ Main features
|
|||
display streams, and provide a way of compiling a report.
|
||||
We want to be able to quickly inspect what has been run.
|
||||
|
||||
4) **Fast compressed Persistence**: a replacement for pickle to work
|
||||
efficiently on Python objects containing large data (
|
||||
*joblib.dump* & *joblib.load* ).
|
||||
|
||||
..
|
||||
>>> import shutil ; shutil.rmtree('/tmp/joblib/')
|
||||
|
||||
"""
|
||||
|
||||
__version__ = '0.5.7'
|
||||
__version__ = '0.6.0b'
|
||||
|
||||
|
||||
from .memory import Memory
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
"""
|
||||
A pickler to save numpy arrays in separate .npy files.
|
||||
Utilities for fast persistence of big data, with optional compression.
|
||||
"""
|
||||
|
||||
# Author: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
|
||||
|
|
@ -10,46 +10,175 @@ import pickle
|
|||
import traceback
|
||||
import sys
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import zipfile
|
||||
import zlib
|
||||
import warnings
|
||||
|
||||
if sys.version_info[0] == 3:
|
||||
from pickle import _Unpickler as Unpickler
|
||||
from cStringIO import StringIO as BytesIO
|
||||
else:
|
||||
if sys.version_info[0] >= 3:
|
||||
from io import BytesIO
|
||||
from pickle import _Unpickler as Unpickler
|
||||
def asbytes(s):
|
||||
if isinstance(s, bytes):
|
||||
return s
|
||||
return s.encode('latin1')
|
||||
else:
|
||||
try:
|
||||
from io import BytesIO
|
||||
except ImportError:
|
||||
# BytesIO has been added in Python 2.5
|
||||
from cStringIO import StringIO as BytesIO
|
||||
from pickle import Unpickler
|
||||
asbytes = str
|
||||
|
||||
_MEGA = 2**20
|
||||
_MAX_LEN = len(hex(2**64))
|
||||
|
||||
# To detect file types
|
||||
_ZFILE_PREFIX = asbytes('ZF')
|
||||
|
||||
###############################################################################
|
||||
# Compressed file with Zlib
|
||||
|
||||
def _read_magic(file_handle):
|
||||
""" Utility to check the magic signature of a file identifying it as a
|
||||
Zfile
|
||||
"""
|
||||
magic = file_handle.read(len(_ZFILE_PREFIX))
|
||||
# Pickling needs file-handles at the beginning of the file
|
||||
file_handle.seek(0)
|
||||
return magic
|
||||
|
||||
|
||||
def read_zfile(file_handle):
|
||||
"""Read the z-file and return the content as a string
|
||||
|
||||
Z-files are raw data compressed with zlib used internally by joblib
|
||||
for persistence. Backward compatibility is not garantied. Do not
|
||||
use for external purposes.
|
||||
"""
|
||||
file_handle.seek(0)
|
||||
assert _read_magic(file_handle) == _ZFILE_PREFIX, \
|
||||
"File does not have the right magic"
|
||||
length = file_handle.read(len(_ZFILE_PREFIX) + _MAX_LEN)
|
||||
length = length[len(_ZFILE_PREFIX):]
|
||||
length = int(length, 16)
|
||||
# We use the known length of the data to tell Zlib the size of the
|
||||
# buffer to allocate.
|
||||
data = zlib.decompress(file_handle.read(), 15, length)
|
||||
assert len(data) == length, (
|
||||
"Incorrect data length while decompressing %s."
|
||||
"The file could be corrupted." % file_handle)
|
||||
return data
|
||||
|
||||
|
||||
def write_zfile(file_handle, data, compress=1):
|
||||
"""Write the data in the given file as a Z-file.
|
||||
|
||||
Z-files are raw data compressed with zlib used internally by joblib
|
||||
for persistence. Backward compatibility is not garantied. Do not
|
||||
use for external purposes.
|
||||
"""
|
||||
file_handle.write(_ZFILE_PREFIX)
|
||||
length = hex(len(data))
|
||||
if sys.version_info[0] < 3 and type(length) is long:
|
||||
# We need to remove the trailing 'L' in the hex representation
|
||||
length = length[:-1]
|
||||
# Store the length of the data
|
||||
file_handle.write(length.ljust(_MAX_LEN))
|
||||
file_handle.write(zlib.compress(data, compress))
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Utility objects for persistence.
|
||||
|
||||
|
||||
class NDArrayWrapper(object):
|
||||
""" An object to be persisted instead of numpy arrays.
|
||||
|
||||
The only thing this object does, is store the filename in wich
|
||||
the array has been persisted.
|
||||
The only thing this object does, is to carrus the filename in wich
|
||||
the array has been persisted, and the array subclass.
|
||||
"""
|
||||
def __init__(self, filename, subclass=None):
|
||||
def __init__(self, filename, subclass):
|
||||
"Store the useful information for later"
|
||||
self.filename = filename
|
||||
self.subclass = subclass
|
||||
|
||||
def read(self, unpickler):
|
||||
"Reconstruct the array"
|
||||
filename = os.path.join(unpickler._dirname, self.filename)
|
||||
# Load the array from the disk
|
||||
if unpickler.np.__version__ >= '1.3':
|
||||
array = unpickler.np.load(filename,
|
||||
mmap_mode=unpickler.mmap_mode)
|
||||
else:
|
||||
# Numpy does not have mmap_mode before 1.3
|
||||
array = unpickler.np.load(filename)
|
||||
# Reconstruct subclasses
|
||||
if not self.subclass in (unpickler.np.ndarray,
|
||||
unpickler.np.memmap):
|
||||
# We need to reconstruct another subclass
|
||||
new_array = unpickler.np.core.multiarray._reconstruct(
|
||||
self.subclass, (0,), 'b')
|
||||
new_array.__array_prepare__(array)
|
||||
array = new_array
|
||||
return array
|
||||
|
||||
|
||||
class ZNDArrayWrapper(NDArrayWrapper):
|
||||
"""An object to be persisted instead of numpy arrays.
|
||||
|
||||
This object store the Zfile filename in wich
|
||||
the data array has been persisted, and the meta information to
|
||||
retrieve it.
|
||||
|
||||
The reason that we store the raw buffer data of the array and
|
||||
the meta information, rather than array representation routine
|
||||
(tostring) is that it enables us to use completely the strided
|
||||
model to avoid memory copies (a and a.T store as fast). In
|
||||
addition saving the heavy information separately can avoid
|
||||
creating large temporary buffers when unpickling data with
|
||||
large arrays.
|
||||
"""
|
||||
def __init__(self, filename, init_args, state):
|
||||
"Store the useful information for later"
|
||||
self.filename = filename
|
||||
self.state = state
|
||||
self.init_args = init_args
|
||||
|
||||
def read(self, unpickler):
|
||||
"Reconstruct the array from the meta-information and the z-file"
|
||||
# Here we a simply reproducing the unpickling mechanism for numpy
|
||||
# arrays
|
||||
filename = os.path.join(unpickler._dirname, self.filename)
|
||||
array = unpickler.np.core.multiarray._reconstruct(*self.init_args)
|
||||
data = read_zfile(open(filename, 'rb'))
|
||||
state = self.state + (data,)
|
||||
array.__setstate__(state)
|
||||
return array
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Pickler classes
|
||||
|
||||
class NumpyPickler(pickle.Pickler):
|
||||
""" A pickler subclass that extracts ndarrays and saves them in .npy
|
||||
files outside of the pickle.
|
||||
"""A pickler to persist of big data efficiently.
|
||||
|
||||
The main features of this object are:
|
||||
|
||||
* persistence of numpy arrays in separate .npy files, for which
|
||||
I/O is fast.
|
||||
|
||||
* optional compression using Zlib, with a special care on avoid
|
||||
temporaries.
|
||||
"""
|
||||
|
||||
def __init__(self, filename):
|
||||
def __init__(self, filename, compress=0, cache_size=100):
|
||||
self._filename = filename
|
||||
self._filenames = [filename, ]
|
||||
self.file = open(filename, 'wb')
|
||||
self.cache_size = cache_size
|
||||
self.compress = compress
|
||||
if not self.compress:
|
||||
self.file = open(filename, 'wb')
|
||||
else:
|
||||
self.file = BytesIO()
|
||||
# Count the number of npy files that we have created:
|
||||
self._npy_counter = 0
|
||||
pickle.Pickler.__init__(self, self.file,
|
||||
|
|
@ -61,45 +190,72 @@ class NumpyPickler(pickle.Pickler):
|
|||
np = None
|
||||
self.np = np
|
||||
|
||||
def _write_array(self, array, filename):
|
||||
if not self.compress:
|
||||
self.np.save(filename, array)
|
||||
container = NDArrayWrapper(os.path.basename(filename),
|
||||
type(array))
|
||||
else:
|
||||
filename += '.z'
|
||||
# Efficient compressed storage:
|
||||
# The meta data is stored in the container, and the core
|
||||
# numerics in a z-file
|
||||
_, init_args, state = array.__reduce__()
|
||||
# the last entry of 'state' is the data itself
|
||||
write_zfile(open(filename, 'w'), state[-1],
|
||||
compress=self.compress)
|
||||
state = state[:-1]
|
||||
container = ZNDArrayWrapper(os.path.basename(filename),
|
||||
init_args, state)
|
||||
return container, filename
|
||||
|
||||
def save(self, obj):
|
||||
""" Subclass the save method, to save ndarray subclasses in npy
|
||||
files, rather than pickling them. Of course, this is a
|
||||
total abuse of the Pickler class.
|
||||
"""
|
||||
if self.np is not None and type(obj) in (self.np.ndarray,
|
||||
self.np.matrix, self.np.memmap):
|
||||
self.np.matrix, self.np.memmap):
|
||||
size = obj.size * obj.itemsize
|
||||
if self.compress and size < self.cache_size * _MEGA:
|
||||
# When compressing, as we are not writing directly to the
|
||||
# disk, it is more efficient to use standard pickling
|
||||
if type(obj) is self.np.memmap:
|
||||
# Pickling doesn't work with memmaped arrays
|
||||
obj = self.np.asarray(obj)
|
||||
return pickle.Pickler.save(self, obj)
|
||||
self._npy_counter += 1
|
||||
try:
|
||||
filename = '%s_%02i.npy' % (self._filename,
|
||||
self._npy_counter)
|
||||
self.np.save(filename, obj)
|
||||
# This converts the array in a container
|
||||
obj, filename = self._write_array(obj, filename)
|
||||
self._filenames.append(filename)
|
||||
obj = NDArrayWrapper(os.path.basename(filename),
|
||||
type(obj))
|
||||
except:
|
||||
self._npy_counter -= 1
|
||||
# XXX: We should have a logging mechanism
|
||||
print 'Failed to save %s to .npy file:\n%s' % (
|
||||
type(obj),
|
||||
traceback.format_exc())
|
||||
pickle.Pickler.save(self, obj)
|
||||
return pickle.Pickler.save(self, obj)
|
||||
|
||||
def close(self):
|
||||
if self.compress:
|
||||
write_zfile(open(self._filename, 'wb'),
|
||||
self.file.getvalue(), self.compress)
|
||||
# The file handes are closed in the dump function
|
||||
|
||||
|
||||
class NumpyUnpickler(Unpickler):
|
||||
""" A subclass of the Unpickler to unpickle our numpy pickles.
|
||||
"""A subclass of the Unpickler to unpickle our numpy pickles.
|
||||
"""
|
||||
dispatch = Unpickler.dispatch.copy()
|
||||
|
||||
def __init__(self, filename, file_handle=None, mmap_mode=None):
|
||||
def __init__(self, filename, file_handle, mmap_mode=None):
|
||||
self._filename = os.path.basename(filename)
|
||||
self.mmap_mode = mmap_mode
|
||||
self._dirname = os.path.dirname(filename)
|
||||
if file_handle is None:
|
||||
file_handle = self._open_file(self._filename)
|
||||
if isinstance(file_handle, basestring):
|
||||
# To handle memmap, we need to have file names
|
||||
file_handle = open(file_handle, 'rb')
|
||||
self.file_handle = file_handle
|
||||
self.mmap_mode = mmap_mode
|
||||
self.file_handle = self._open_pickle(file_handle)
|
||||
Unpickler.__init__(self, self.file_handle)
|
||||
try:
|
||||
import numpy as np
|
||||
|
|
@ -107,9 +263,8 @@ class NumpyUnpickler(Unpickler):
|
|||
np = None
|
||||
self.np = np
|
||||
|
||||
def _open_file(self, name):
|
||||
"Return the path of the given file in our store"
|
||||
return os.path.join(self._dirname, name)
|
||||
def _open_pickle(self, file_handle):
|
||||
return file_handle
|
||||
|
||||
def load_build(self):
|
||||
""" This method is called to set the state of a newly created
|
||||
|
|
@ -117,7 +272,7 @@ class NumpyUnpickler(Unpickler):
|
|||
|
||||
We capture it to replace our place-holder objects,
|
||||
NDArrayWrapper, by the array we are interested in. We
|
||||
replace directly in the stack of pickler.
|
||||
replace them directly in the stack of pickler.
|
||||
"""
|
||||
Unpickler.load_build(self)
|
||||
if isinstance(self.stack[-1], NDArrayWrapper):
|
||||
|
|
@ -125,21 +280,7 @@ class NumpyUnpickler(Unpickler):
|
|||
raise ImportError('Trying to unpickle an ndarray, '
|
||||
"but numpy didn't import correctly")
|
||||
nd_array_wrapper = self.stack.pop()
|
||||
if self.np.__version__ >= '1.3':
|
||||
array = self.np.load(
|
||||
self._open_file(nd_array_wrapper.filename),
|
||||
mmap_mode=self.mmap_mode)
|
||||
else:
|
||||
# Numpy does not have mmap_mode before 1.3
|
||||
array = self.np.load(
|
||||
self._open_file(nd_array_wrapper.filename),
|
||||
mmap_mode=self.mmap_mode)
|
||||
if not nd_array_wrapper.subclass is self.np.ndarray:
|
||||
# We need to reconstruct another subclass
|
||||
new_array = self.np.core.multiarray._reconstruct(
|
||||
nd_array_wrapper.subclass, (0,), 'b')
|
||||
new_array.__array_prepare__(array)
|
||||
array = new_array
|
||||
array = nd_array_wrapper.read(self)
|
||||
self.stack.append(array)
|
||||
|
||||
# Be careful to register our new method.
|
||||
|
|
@ -147,65 +288,64 @@ class NumpyUnpickler(Unpickler):
|
|||
|
||||
|
||||
class ZipNumpyUnpickler(NumpyUnpickler):
|
||||
""" A subclass of our Unpickler to unpickle on the fly from zips.
|
||||
"""
|
||||
"""A subclass of our Unpickler to unpickle on the fly from
|
||||
compressed storage."""
|
||||
|
||||
def __init__(self, file_handle):
|
||||
kwargs = dict(compression=zipfile.ZIP_DEFLATED)
|
||||
if sys.version_info >= (2, 5):
|
||||
kwargs['allowZip64'] = True
|
||||
self._zip_file = zipfile.ZipFile(file_handle, **kwargs)
|
||||
NumpyUnpickler.__init__(self, 'joblib_dump.pkl',
|
||||
def __init__(self, filename, file_handle):
|
||||
NumpyUnpickler.__init__(self, filename,
|
||||
file_handle,
|
||||
mmap_mode=None)
|
||||
|
||||
def _open_file(self, name):
|
||||
"Return the path of the given file in our store"
|
||||
decompression_buffer = BytesIO(
|
||||
self._zip_file.read(os.path.join('dump_file', name)))
|
||||
return decompression_buffer
|
||||
def _open_pickle(self, file_handle):
|
||||
return BytesIO(read_zfile(file_handle))
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Utility functions
|
||||
|
||||
def dump(value, filename, compress=False):
|
||||
""" Persist an arbitrary Python object into a filename, with numpy arrays
|
||||
saved as separate .npy files.
|
||||
def dump(value, filename, compress=0, cache_size=100):
|
||||
"""Fast persistence of an arbitrary Python object into a files, with
|
||||
dedicated storage for numpy arrays.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
value: any Python object
|
||||
The object to store to disk
|
||||
filename: string
|
||||
The name of the file in which it is to be stored
|
||||
compress: boolean, optional
|
||||
Whether to compress the data on the disk or not
|
||||
Parameters
|
||||
-----------
|
||||
value: any Python object
|
||||
The object to store to disk
|
||||
filename: string
|
||||
The name of the file in which it is to be stored
|
||||
compress: integer for 0 to 9, optional
|
||||
Optional compression level for the data. 0 is no compression.
|
||||
Higher means more compression, but also slower read and
|
||||
write times. Using a value of 3 is often a good compromise.
|
||||
See the notes for more details.
|
||||
cache_size: positive number, optional
|
||||
Fixes the order of magnitude (in megabytes) of the cache used
|
||||
for in-memory compression. Note that this is just an order of
|
||||
magnitude estimate and that for big arrays, the code will go
|
||||
over this value at dump and at load time.
|
||||
|
||||
Returns
|
||||
-------
|
||||
filenames: list of strings
|
||||
The list of file names in which the data is stored. If
|
||||
compress is false, each array is stored in a different file.
|
||||
Returns
|
||||
-------
|
||||
filenames: list of strings
|
||||
The list of file names in which the data is stored. If
|
||||
compress is false, each array is stored in a different file.
|
||||
|
||||
See Also
|
||||
--------
|
||||
joblib.load : corresponding loader
|
||||
See Also
|
||||
--------
|
||||
joblib.load : corresponding loader
|
||||
|
||||
Notes
|
||||
-----
|
||||
compressed files take extra disk space during the dump, and extra
|
||||
memory during the loading.
|
||||
Notes
|
||||
-----
|
||||
Memmapping on load cannot be used for compressed files. Thus
|
||||
using compression can significantly slow down loading. In
|
||||
addition, compressed files take extra extra memory during
|
||||
dump and load.
|
||||
"""
|
||||
if compress:
|
||||
return _dump_zipped(value, filename)
|
||||
else:
|
||||
return _dump(value, filename)
|
||||
|
||||
|
||||
def _dump(value, filename):
|
||||
try:
|
||||
pickler = NumpyPickler(filename)
|
||||
pickler = NumpyPickler(filename, compress=compress,
|
||||
cache_size=cache_size)
|
||||
pickler.dump(value)
|
||||
pickler.close()
|
||||
finally:
|
||||
if 'pickler' in locals() and hasattr(pickler, 'file'):
|
||||
pickler.file.flush()
|
||||
|
|
@ -213,84 +353,49 @@ def _dump(value, filename):
|
|||
return pickler._filenames
|
||||
|
||||
|
||||
def _dump_zipped(value, filename):
|
||||
""" Persist an arbitrary Python object into a compressed zip
|
||||
filename.
|
||||
"""
|
||||
kwargs = dict(compression=zipfile.ZIP_DEFLATED, mode='w')
|
||||
if sys.version_info >= (2, 5):
|
||||
kwargs['allowZip64'] = True
|
||||
dump_file = zipfile.ZipFile(filename, **kwargs)
|
||||
|
||||
# Stage file in a temporary dir on disk, before writing to zip.
|
||||
tmp_dir = tempfile.mkdtemp(prefix='joblib-',
|
||||
dir=os.path.dirname(filename))
|
||||
try:
|
||||
_dump(value, os.path.join(tmp_dir, 'joblib_dump.pkl'))
|
||||
for sub_file in os.listdir(tmp_dir):
|
||||
# We use a different arcname (archive name) to avoid having
|
||||
# the name of our tmp_dir in the archive
|
||||
dump_file.write(os.path.join(tmp_dir, sub_file),
|
||||
arcname=os.path.join('dump_file', sub_file))
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
dump_file.close()
|
||||
return [filename]
|
||||
|
||||
|
||||
def load(filename, mmap_mode=None):
|
||||
""" Reconstruct a Python object and the numpy arrays it contains from
|
||||
a persisted file.
|
||||
"""Reconstruct a Python object from a file persisted with joblib.load.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
filename: string
|
||||
The name of the file from which to load the object
|
||||
mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, optional
|
||||
If not None, the arrays are memory-mapped from the disk. This
|
||||
mode has not effect for compressed files. Note that in this
|
||||
case the reconstructed object might not longer match exactly
|
||||
the originally pickled object.
|
||||
Parameters
|
||||
-----------
|
||||
filename: string
|
||||
The name of the file from which to load the object
|
||||
mmap_mode: {None, 'r+', 'r', 'w+', 'c'}, optional
|
||||
If not None, the arrays are memory-mapped from the disk. This
|
||||
mode has not effect for compressed files. Note that in this
|
||||
case the reconstructed object might not longer match exactly
|
||||
the originally pickled object.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result: any Python object
|
||||
The object stored in the file.
|
||||
Returns
|
||||
-------
|
||||
result: any Python object
|
||||
The object stored in the file.
|
||||
|
||||
See Also
|
||||
--------
|
||||
joblib.dump : function to save an object
|
||||
See Also
|
||||
--------
|
||||
joblib.dump : function to save an object
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
This function loads the numpy array files saved separately. If
|
||||
the mmap_mode argument is given, it is passed to np.save and
|
||||
arrays are loaded as memmaps. As a consequence, the reconstructed
|
||||
object might not match the original pickled object.
|
||||
Notes
|
||||
-----
|
||||
|
||||
This function can load numpy array files saved separately during the
|
||||
dump. If the mmap_mode argument is given, it is passed to np.load and
|
||||
arrays are loaded as memmaps. As a consequence, the reconstructed
|
||||
object might not match the original pickled object. Note that if the
|
||||
file was saved with compression, the arrays cannot be memmaped.
|
||||
"""
|
||||
# Code to detect zip files
|
||||
_ZIP_PREFIX = 'PK\x03\x04'
|
||||
try:
|
||||
# Py3k compatibility
|
||||
from numpy.compat import asbytes
|
||||
_ZIP_PREFIX = asbytes(_ZIP_PREFIX)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
file_handle = open(filename, 'rb')
|
||||
if file_handle.read(len(_ZIP_PREFIX)) == _ZIP_PREFIX:
|
||||
# We are careful to open the file hanlde early and keep it open to
|
||||
# avoid race-conditions on renames. That said, if data are stored in
|
||||
# companion files, moving the directory will create a race when
|
||||
# joblib tries to access the companion files.
|
||||
if _read_magic(file_handle) == _ZFILE_PREFIX:
|
||||
if mmap_mode is not None:
|
||||
warnings.warn('file "%(filename)s" appears to be a zip, '
|
||||
'ignoring mmap_mode "%(mmap_mode)s" flag passed'
|
||||
% locals(),
|
||||
Warning, stacklevel=2)
|
||||
unpickler = ZipNumpyUnpickler(file_handle=file_handle)
|
||||
% locals(), Warning, stacklevel=2)
|
||||
unpickler = ZipNumpyUnpickler(filename, file_handle=file_handle)
|
||||
else:
|
||||
# Pickling needs file-handles at the beginning of the file
|
||||
file_handle.seek(0)
|
||||
unpickler = NumpyUnpickler(filename,
|
||||
file_handle=file_handle,
|
||||
mmap_mode=mmap_mode)
|
||||
|
|
@ -298,11 +403,8 @@ def load(filename, mmap_mode=None):
|
|||
try:
|
||||
obj = unpickler.load()
|
||||
finally:
|
||||
if 'unpickler' in locals():
|
||||
if hasattr(unpickler, 'file'):
|
||||
unpickler.file.close()
|
||||
if hasattr(unpickler, '_zip_file'):
|
||||
unpickler._zip_file.close()
|
||||
if hasattr(unpickler, 'file_handle'):
|
||||
unpickler.file_handle.close()
|
||||
return obj
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ Helpers for embarrassingly parallel code.
|
|||
|
||||
import os
|
||||
import sys
|
||||
import warnings
|
||||
from math import sqrt
|
||||
import functools
|
||||
import time
|
||||
import threading
|
||||
|
|
@ -41,6 +43,27 @@ def cpu_count():
|
|||
return multiprocessing.cpu_count()
|
||||
|
||||
|
||||
###############################################################################
|
||||
# For verbosity
|
||||
|
||||
def _verbosity_filter(index, verbose):
|
||||
""" Returns False for indices increasingly appart, the distance
|
||||
depending on the value of verbose.
|
||||
|
||||
We use a lag increasing as the square of index
|
||||
"""
|
||||
if not verbose:
|
||||
return True
|
||||
elif verbose > 10:
|
||||
return False
|
||||
if index == 0:
|
||||
return False
|
||||
verbose = .5*(11 - verbose)**2
|
||||
scale = sqrt(index/verbose)
|
||||
next_scale = sqrt((index + 1)/verbose)
|
||||
return (int(next_scale) == int(scale))
|
||||
|
||||
|
||||
###############################################################################
|
||||
class WorkerInterrupt(Exception):
|
||||
""" An exception that is not KeyboardInterrupt to allow subprocesses
|
||||
|
|
@ -114,40 +137,10 @@ class CallBack(object):
|
|||
self.index = index
|
||||
|
||||
def __call__(self, out):
|
||||
if self.parallel.verbose:
|
||||
self.print_progress()
|
||||
self.parallel.print_progress(self.index)
|
||||
if self.parallel._iterable:
|
||||
self.parallel.dispatch_next()
|
||||
|
||||
def print_progress(self):
|
||||
# XXX: Not using the logger framework: need to
|
||||
# learn to use logger better.
|
||||
n_jobs = len(self.parallel._pool._pool)
|
||||
if self.parallel.n_dispatched > 2 * n_jobs:
|
||||
# Report less often
|
||||
if not self.index % n_jobs == 0:
|
||||
return
|
||||
elapsed_time = time.time() - self.parallel._start_time
|
||||
remaining_time = (elapsed_time / (self.index + 1) *
|
||||
(self.parallel.n_dispatched - self.index - 1.))
|
||||
if self.parallel._iterable:
|
||||
# The object is still building its job list
|
||||
total = "%3i+" % self.parallel.n_dispatched
|
||||
else:
|
||||
total = "%3i " % self.parallel.n_dispatched
|
||||
|
||||
if self.parallel.verbose < 50:
|
||||
writer = sys.stderr.write
|
||||
else:
|
||||
writer = sys.stdout.write
|
||||
writer('[%s]: Done %3i out of %s |elapsed: %s remaining: %s\n'
|
||||
% (self.parallel,
|
||||
self.index + 1,
|
||||
total,
|
||||
short_format_time(elapsed_time),
|
||||
short_format_time(remaining_time),
|
||||
))
|
||||
|
||||
|
||||
###############################################################################
|
||||
class Parallel(Logger):
|
||||
|
|
@ -158,11 +151,14 @@ class Parallel(Logger):
|
|||
n_jobs: int
|
||||
The number of jobs to use for the computation. If -1 all CPUs
|
||||
are used. If 1 is given, no parallel computing code is used
|
||||
at all, which is useful for debuging.
|
||||
at all, which is useful for debuging. For n_jobs below -1,
|
||||
(n_cpus + 1 - n_jobs) are used. Thus for n_jobs = -2, all
|
||||
CPUs but one are used.
|
||||
verbose: int, optional
|
||||
The verbosity level. If 1 is given, the elapsed time as well
|
||||
as the estimated remaining time are displayed. Above 100, the
|
||||
output is sent to stdout.
|
||||
The verbosity level: if non zero, progress messages are
|
||||
printed. Above 50, the output is sent to stdout.
|
||||
The frequency of the messages increases with the verbosity level.
|
||||
If it more than 10, all iterations are reported.
|
||||
pre_dispatch: {'all', integer, or expression, as in '3*n_jobs'}
|
||||
The amount of jobs to be pre-dispatched. Default is 'all',
|
||||
but it may be memory consuming, for instance if each job
|
||||
|
|
@ -212,16 +208,17 @@ class Parallel(Logger):
|
|||
>>> i
|
||||
(0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0)
|
||||
|
||||
The progress meter::
|
||||
The progress meter: the higher the value of `verbose`, the more
|
||||
messages::
|
||||
|
||||
>>> from time import sleep
|
||||
>>> from sklearn.externals.joblib import Parallel, delayed
|
||||
>>> r = Parallel(n_jobs=2, verbose=1)(delayed(sleep)(.1) for _ in range(10)) #doctest: +SKIP
|
||||
[Parallel(n_jobs=2)]: Done 1 out of 10 |elapsed: 0.1s remaining: 0.9s
|
||||
[Parallel(n_jobs=2)]: Done 3 out of 10 |elapsed: 0.2s remaining: 0.5s
|
||||
[Parallel(n_jobs=2)]: Done 5 out of 10 |elapsed: 0.3s remaining: 0.3s
|
||||
[Parallel(n_jobs=2)]: Done 7 out of 10 |elapsed: 0.4s remaining: 0.2s
|
||||
[Parallel(n_jobs=2)]: Done 9 out of 10 |elapsed: 0.5s remaining: 0.1s
|
||||
>>> r = Parallel(n_jobs=2, verbose=5)(delayed(sleep)(.1) for _ in range(10)) #doctest: +SKIP
|
||||
[Parallel(n_jobs=2)]: Done 1 out of 10 | elapsed: 0.1s remaining: 0.9s
|
||||
[Parallel(n_jobs=2)]: Done 3 out of 10 | elapsed: 0.2s remaining: 0.5s
|
||||
[Parallel(n_jobs=2)]: Done 6 out of 10 | elapsed: 0.3s remaining: 0.2s
|
||||
[Parallel(n_jobs=2)]: Done 9 out of 10 | elapsed: 0.5s remaining: 0.1s
|
||||
[Parallel(n_jobs=2)]: Done 10 out of 10 | elapsed: 0.5s finished
|
||||
|
||||
Traceback example, note how the line of the error is indicated
|
||||
as well as the values of the parameter passed to the function that
|
||||
|
|
@ -257,7 +254,7 @@ class Parallel(Logger):
|
|||
data is generated on the fly. Note how the producer is first
|
||||
called a 3 times before the parallel loop is initiated, and then
|
||||
called to generate new data on the fly. In this case the total
|
||||
number of iterations reported is underestimated::
|
||||
number of iterations cannot be reported in the progress messages::
|
||||
|
||||
>>> from math import sqrt
|
||||
>>> from sklearn.externals.joblib import Parallel, delayed
|
||||
|
|
@ -272,11 +269,15 @@ class Parallel(Logger):
|
|||
Produced 0
|
||||
Produced 1
|
||||
Produced 2
|
||||
[Parallel(n_jobs=2)]: Done 1 out of 3+ |elapsed: ...s remaining: ...s
|
||||
[Parallel(n_jobs=2)]: Done 1 jobs | elapsed: 0.0s
|
||||
Produced 3
|
||||
[Parallel(n_jobs=2)]: Done ... out of 4+ |elapsed: ...s remaining: ...s
|
||||
...
|
||||
|
||||
[Parallel(n_jobs=2)]: Done 2 jobs | elapsed: 0.0s
|
||||
Produced 4
|
||||
[Parallel(n_jobs=2)]: Done 3 jobs | elapsed: 0.0s
|
||||
Produced 5
|
||||
[Parallel(n_jobs=2)]: Done 4 jobs | elapsed: 0.0s
|
||||
[Parallel(n_jobs=2)]: Done 5 out of 6 | elapsed: 0.0s remaining: 0.0s
|
||||
[Parallel(n_jobs=2)]: Done 6 out of 6 | elapsed: 0.0s finished
|
||||
'''
|
||||
def __init__(self, n_jobs=None, verbose=0, pre_dispatch='all'):
|
||||
self.verbose = verbose
|
||||
|
|
@ -293,11 +294,12 @@ class Parallel(Logger):
|
|||
"""
|
||||
if self._pool is None:
|
||||
job = ImmediateApply(func, args, kwargs)
|
||||
if self.verbose:
|
||||
print '[%s]: Done job %3i | elapsed: %s' % (
|
||||
self, len(self._jobs) + 1,
|
||||
short_format_time(time.time() - self._start_time)
|
||||
)
|
||||
index = len(self._jobs)
|
||||
if not _verbosity_filter(index, self.verbose):
|
||||
self._print('Done %3i jobs | elapsed: %s',
|
||||
(index + 1,
|
||||
short_format_time(time.time() - self._start_time)
|
||||
))
|
||||
self._jobs.append(job)
|
||||
self.n_dispatched += 1
|
||||
else:
|
||||
|
|
@ -332,6 +334,59 @@ class Parallel(Logger):
|
|||
self._iterable = None
|
||||
return
|
||||
|
||||
def _print(self, msg, msg_args):
|
||||
""" Display the message on stout or stderr depending on verbosity
|
||||
"""
|
||||
# XXX: Not using the logger framework: need to
|
||||
# learn to use logger better.
|
||||
if not self.verbose:
|
||||
return
|
||||
if self.verbose < 50:
|
||||
writer = sys.stderr.write
|
||||
else:
|
||||
writer = sys.stdout.write
|
||||
msg = msg % msg_args
|
||||
writer('[%s]: %s\n' % (self, msg))
|
||||
|
||||
def print_progress(self, index):
|
||||
"""Display the process of the parallel execution only a fraction
|
||||
of time, controled by self.verbose.
|
||||
"""
|
||||
if not self.verbose:
|
||||
return
|
||||
elapsed_time = time.time() - self._start_time
|
||||
|
||||
# This is heuristic code to print only 'verbose' times a messages
|
||||
# The challenge is that we may not know the queue length
|
||||
if self._iterable:
|
||||
if _verbosity_filter(index, self.verbose):
|
||||
return
|
||||
self._print('Done %3i jobs | elapsed: %s',
|
||||
(index + 1,
|
||||
short_format_time(elapsed_time),
|
||||
))
|
||||
else:
|
||||
# We are finished dispatching
|
||||
queue_length = self.n_dispatched
|
||||
# We always display the first loop
|
||||
if not index == 0:
|
||||
# Display depending on the number of remaining items
|
||||
# A message as soon as we finish dispatching, cursor is 0
|
||||
cursor = (queue_length - index + 1
|
||||
- self._pre_dispatch_amount)
|
||||
frequency = (queue_length // self.verbose) + 1
|
||||
is_last_item = (index + 1 == queue_length)
|
||||
if (is_last_item or cursor % frequency):
|
||||
return
|
||||
remaining_time = (elapsed_time / (index + 1) *
|
||||
(self.n_dispatched - index - 1.))
|
||||
self._print('Done %3i out of %3i | elapsed: %s remaining: %s',
|
||||
(index + 1,
|
||||
queue_length,
|
||||
short_format_time(elapsed_time),
|
||||
short_format_time(remaining_time),
|
||||
))
|
||||
|
||||
def retrieve(self):
|
||||
self._output = list()
|
||||
while self._jobs:
|
||||
|
|
@ -376,8 +431,8 @@ Sub-process traceback:
|
|||
if self._jobs:
|
||||
raise ValueError('This Parallel instance is already running')
|
||||
n_jobs = self.n_jobs
|
||||
if n_jobs == -1 and multiprocessing is not None:
|
||||
n_jobs = multiprocessing.cpu_count()
|
||||
if n_jobs < 0 and multiprocessing is not None:
|
||||
n_jobs = min(multiprocessing.cpu_count() + 1 + n_jobs, 1)
|
||||
|
||||
# The list of exceptions that we will capture
|
||||
self.exceptions = [TransportableException]
|
||||
|
|
@ -385,21 +440,30 @@ Sub-process traceback:
|
|||
n_jobs = 1
|
||||
self._pool = None
|
||||
else:
|
||||
self._pool = multiprocessing.Pool(n_jobs)
|
||||
self._lock = threading.Lock()
|
||||
# We are using multiprocessing, we also want to capture
|
||||
# KeyboardInterrupts
|
||||
self.exceptions.extend([KeyboardInterrupt, WorkerInterrupt])
|
||||
if multiprocessing.current_process()._daemonic:
|
||||
# Daemonic processes cannot have children
|
||||
n_jobs = 1
|
||||
self._pool = None
|
||||
warnings.warn(
|
||||
'Parallel loops cannot be nested, setting n_jobs=1',
|
||||
stacklevel=2)
|
||||
else:
|
||||
self._pool = multiprocessing.Pool(n_jobs)
|
||||
self._lock = threading.Lock()
|
||||
# We are using multiprocessing, we also want to capture
|
||||
# KeyboardInterrupts
|
||||
self.exceptions.extend([KeyboardInterrupt, WorkerInterrupt])
|
||||
|
||||
if self.pre_dispatch == 'all' or n_jobs == 1:
|
||||
self._iterable = None
|
||||
self._pre_dispatch_amount = 0
|
||||
else:
|
||||
self._iterable = iterable
|
||||
self._dispatch_amount = 0
|
||||
pre_dispatch = self.pre_dispatch
|
||||
if hasattr(pre_dispatch, 'endswith'):
|
||||
pre_dispatch = eval(pre_dispatch)
|
||||
pre_dispatch = int(pre_dispatch)
|
||||
self._pre_dispatch_amount = pre_dispatch = int(pre_dispatch)
|
||||
iterable = itertools.islice(iterable, pre_dispatch)
|
||||
|
||||
self._start_time = time.time()
|
||||
|
|
@ -409,6 +473,14 @@ Sub-process traceback:
|
|||
self.dispatch(function, args, kwargs)
|
||||
|
||||
self.retrieve()
|
||||
# Make sure that we get a last message telling us we are done
|
||||
elapsed_time = time.time() - self._start_time
|
||||
self._print('Done %3i out of %3i | elapsed: %s finished',
|
||||
(len(self._output),
|
||||
len(self._output),
|
||||
short_format_time(elapsed_time)
|
||||
))
|
||||
|
||||
finally:
|
||||
if n_jobs > 1:
|
||||
self._pool.close()
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from tempfile import mkdtemp
|
|||
import copy
|
||||
import shutil
|
||||
import os
|
||||
import random
|
||||
|
||||
import nose
|
||||
|
||||
|
|
@ -110,13 +111,14 @@ def teardown_module():
|
|||
# Tests
|
||||
|
||||
def test_standard_types():
|
||||
#""" Test pickling and saving with standard types.
|
||||
#"""
|
||||
# Test pickling and saving with standard types.
|
||||
filename = env['filename']
|
||||
for compress in [True, False]:
|
||||
for compress in [0, 1]:
|
||||
for member in typelist:
|
||||
numpy_pickle.dump(member, filename, compress=compress)
|
||||
_member = numpy_pickle.load(filename)
|
||||
# Change the file name to avoid side effects between tests
|
||||
this_filename = filename + str(random.randint(0, 1000))
|
||||
numpy_pickle.dump(member, this_filename, compress=compress)
|
||||
_member = numpy_pickle.load(this_filename)
|
||||
# We compare the pickled instance to the reloaded one only if it
|
||||
# can be compared to a copied one
|
||||
if member == copy.deepcopy(member):
|
||||
|
|
@ -126,40 +128,45 @@ def test_standard_types():
|
|||
@with_numpy
|
||||
def test_numpy_persistence():
|
||||
filename = env['filename']
|
||||
a = np.random.random(10)
|
||||
for compress in [True, False]:
|
||||
for obj in (a,), (a, a), [a, a, a]:
|
||||
filenames = numpy_pickle.dump(obj, filename, compress=compress)
|
||||
a = np.random.random((10, 2))
|
||||
for compress, cache_size in ((0, 0), (1, 0), (1, 10)):
|
||||
# We use 'a.T' to have a non C-contiguous array.
|
||||
for index, obj in enumerate(((a,), (a.T,), (a, a), [a, a, a])):
|
||||
# Change the file name to avoid side effects between tests
|
||||
this_filename = filename + str(random.randint(0, 1000))
|
||||
filenames = numpy_pickle.dump(obj, this_filename,
|
||||
compress=compress,
|
||||
cache_size=cache_size)
|
||||
# Check that one file was created per array
|
||||
if not compress:
|
||||
# Check that one file was created per array
|
||||
yield nose.tools.assert_equal, len(filenames), len(obj) + 1
|
||||
# Check that these files do exist
|
||||
for file in filenames:
|
||||
yield nose.tools.assert_true, \
|
||||
os.path.exists(os.path.join(env['dir'], file))
|
||||
else:
|
||||
yield nose.tools.assert_equal, len(filenames), 1
|
||||
nose.tools.assert_equal(len(filenames), len(obj) + 1)
|
||||
# Check that these files do exist
|
||||
for file in filenames:
|
||||
nose.tools.assert_true(
|
||||
os.path.exists(os.path.join(env['dir'], file)))
|
||||
|
||||
# Unpickle the object
|
||||
obj_ = numpy_pickle.load(filename)
|
||||
obj_ = numpy_pickle.load(this_filename)
|
||||
# Check that the items are indeed arrays
|
||||
for item in obj_:
|
||||
yield nose.tools.assert_true, isinstance(item, np.ndarray)
|
||||
nose.tools.assert_true(isinstance(item, np.ndarray))
|
||||
# And finally, check that all the values are equal.
|
||||
yield nose.tools.assert_true, np.all(np.array(obj) ==
|
||||
np.array(obj_))
|
||||
nose.tools.assert_true(np.all(np.array(obj) ==
|
||||
np.array(obj_)))
|
||||
|
||||
# Now test with array subclasses
|
||||
obj = np.matrix(np.zeros(10))
|
||||
filenames = numpy_pickle.dump(obj, filename, compress=compress)
|
||||
obj_ = numpy_pickle.load(filename)
|
||||
yield nose.tools.assert_true, isinstance(obj_, np.matrix)
|
||||
this_filename = filename + str(random.randint(0, 1000))
|
||||
filenames = numpy_pickle.dump(obj, this_filename, compress=compress,
|
||||
cache_size=cache_size)
|
||||
obj_ = numpy_pickle.load(this_filename)
|
||||
nose.tools.assert_true(isinstance(obj_, np.matrix))
|
||||
|
||||
|
||||
@with_numpy
|
||||
def test_memmap_persistence():
|
||||
a = np.random.random(10)
|
||||
filename = env['filename']
|
||||
filename = env['filename'] + str(random.randint(0, 1000))
|
||||
numpy_pickle.dump(a, filename)
|
||||
b = numpy_pickle.load(filename, mmap_mode='r')
|
||||
if np.__version__ >= '1.3':
|
||||
|
|
@ -172,9 +179,16 @@ def test_masked_array_persistence():
|
|||
# not implemented, but it just delegates to the standard pickler.
|
||||
a = np.random.random(10)
|
||||
a = np.ma.masked_greater(a, 0.5)
|
||||
filename = env['filename']
|
||||
filename = env['filename'] + str(random.randint(0, 1000))
|
||||
numpy_pickle.dump(a, filename)
|
||||
b = numpy_pickle.load(filename, mmap_mode='r')
|
||||
nose.tools.assert_true(isinstance(b, np.ma.masked_array))
|
||||
|
||||
|
||||
def test_z_file():
|
||||
# Test saving and loading data with Zfiles
|
||||
filename = env['filename'] + str(random.randint(0, 1000))
|
||||
data = 'Foo, \n Bar, baz, \n\nfoobar'
|
||||
numpy_pickle.write_zfile(file(filename, 'wb'), data)
|
||||
data_read = numpy_pickle.read_zfile(file(filename, 'rb'))
|
||||
nose.tools.assert_equal(data, data_read)
|
||||
|
|
|
|||
|
|
@ -61,11 +61,19 @@ def test_cpu_count():
|
|||
# Test parallel
|
||||
def test_simple_parallel():
|
||||
X = range(10)
|
||||
for n_jobs in (1, 2, -1):
|
||||
for n_jobs in (1, 2, -1, -2):
|
||||
yield (nose.tools.assert_equal, [square(x) for x in X],
|
||||
Parallel(n_jobs=-1)(delayed(square)(x) for x in X))
|
||||
|
||||
|
||||
def nested_loop():
|
||||
Parallel(n_jobs=2)(delayed(square)(.01) for _ in range(2))
|
||||
|
||||
|
||||
def test_nested_loop():
|
||||
Parallel(n_jobs=2)(delayed(nested_loop)() for _ in range(2))
|
||||
|
||||
|
||||
def test_parallel_kwargs():
|
||||
""" Check the keyword argument processing of pmap.
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue