ENH: update joblib

This commit is contained in:
GaelVaroquaux 2012-01-07 16:37:02 +01:00
parent 1543ba214c
commit 54fcfc2da3
5 changed files with 460 additions and 259 deletions

View File

@ -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

View File

@ -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

View File

@ -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()

View File

@ -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)

View File

@ -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.
"""