Repository: joblib/joblib
Pull request #479 (2017-01-18 13:12:14)

[MRG] Provide backport for numpy memmap fix when offset is greater than mmap.ALLOCATIONGRANULARITY
Fixes https://github.com/joblib/joblib/issues/451.

ping @ogrisel @aabadie.

lesteve: OK I'll leave the dtype as strings and merge this one.

Patch:
diff --git a/joblib/backports.py b/joblib/backports.py
new file mode 100644
index 0000000..5fcc45c
--- /dev/null
+++ b/joblib/backports.py
@@ -0,0 +1,26 @@
+"""
+Backports of fixes for joblib dependencies
+"""
+
+from distutils.version import LooseVersion
+
+try:
+    import numpy as np
+
+    def make_memmap(filename, dtype='uint8', mode='r+', offset=0,
+                    shape=None, order='C'):
+        """Backport of numpy memmap offset fix
+
+        See https://github.com/numpy/numpy/pull/8443 for more details.
+        """
+        mm = np.memmap(filename, dtype=dtype, mode=mode, offset=offset,
+                       shape=shape, order=order)
+        if LooseVersion(np.__version__) < '1.13':
+            mm.offset = offset
+        return mm
+except ImportError:
+    def make_memmap(filename, dtype='uint8', mode='r+', offset=0,
+                    shape=None, order='C'):
+        raise NotImplementedError(
+            "'joblib.backports.make_memmap' should not be used "
+            'if numpy is not installed.')
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index bb897a9..87a1a61 100644
--- a/joblib/numpy_pickle.py
+++ b/joblib/numpy_pickle.py
@@ -26,6 +26,7 @@ from .numpy_pickle_compat import NDArrayWrapper
 # which we don't care.
 from .numpy_pickle_compat import ZNDArrayWrapper  # noqa
 from ._compat import _basestring, PY3_OR_LATER
+from .backports import make_memmap
 
 ###############################################################################
 # Utility objects for persistence.
@@ -151,12 +152,12 @@ class NumpyArrayWrapper(object):
         if unpickler.mmap_mode == 'w+':
             unpickler.mmap_mode = 'r+'
 
-        marray = unpickler.np.memmap(unpickler.filename,
-                                     dtype=self.dtype,
-                                     shape=self.shape,
-                                     order=self.order,
-                                     mode=unpickler.mmap_mode,
-                                     offset=offset)
+        marray = make_memmap(unpickler.filename,
+                             dtype=self.dtype,
+                             shape=self.shape,
+                             order=self.order,
+                             mode=unpickler.mmap_mode,
+                             offset=offset)
         # update the offset so that it corresponds to the end of the read array
         unpickler.file_handle.seek(offset + marray.nbytes)
 
diff --git a/joblib/pool.py b/joblib/pool.py
index e5d0464..0867782 100644
--- a/joblib/pool.py
+++ b/joblib/pool.py
@@ -61,7 +61,7 @@ except ImportError:
 from .numpy_pickle import load
 from .numpy_pickle import dump
 from .hashing import hash
-
+from .backports import make_memmap
 # 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'
@@ -107,13 +107,13 @@ def _strided_from_memmap(filename, dtype, mode, offset, order, shape, strides,
 
     if strides is None:
         # Simple, contiguous memmap
-        return np.memmap(filename, dtype=dtype, shape=shape, mode=mode,
-                         offset=offset, order=order)
+        return make_memmap(filename, dtype=dtype, shape=shape, mode=mode,
+                           offset=offset, order=order)
     else:
         # For non-contiguous data, memmap the total enclosing buffer and then
         # extract the non-contiguous view with the stride-tricks API
-        base = np.memmap(filename, dtype=dtype, shape=total_buffer_len,
-                         mode=mode, offset=offset, order=order)
+        base = make_memmap(filename, dtype=dtype, shape=total_buffer_len,
+                           mode=mode, offset=offset, order=order)
         return as_strided(base, shape=shape, strides=strides)
 
 


Test patch:
diff --git a/joblib/test/test_backports.py b/joblib/test/test_backports.py
new file mode 100644
index 0000000..0199d02
--- /dev/null
+++ b/joblib/test/test_backports.py
@@ -0,0 +1,13 @@
+import mmap
+
+from joblib.backports import make_memmap
+from joblib.test.common import with_numpy
+
+
+@with_numpy
+def test_memmap(tmpdir):
+    fname = tmpdir.join('test.mmap').strpath
+    size = 5 * mmap.ALLOCATIONGRANULARITY
+    offset = mmap.ALLOCATIONGRANULARITY + 1
+    memmap_obj = make_memmap(fname, shape=size, mode='w+', offset=offset)
+    assert memmap_obj.offset == offset
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index 1f5196d..b0b29d7 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -13,6 +13,7 @@ import bz2
 import pickle
 import socket
 from contextlib import closing
+import mmap
 
 from joblib.test.common import np, with_numpy
 from joblib.test.common import with_memory_profiler, memory_used
@@ -898,3 +899,19 @@ def test_pickle_in_socket():
         array_reloaded = numpy_pickle.load(cf)
 
     np.testing.assert_array_equal(array_reloaded, test_array)
+
+
+@with_numpy
+def test_load_memmap_with_big_offset(tmpdir):
+    # Test that numpy memmap offset is set correctly if greater than
+    # mmap.ALLOCATIONGRANULARITY, see
+    # https://github.com/joblib/joblib/issues/451 and
+    # https://github.com/numpy/numpy/pull/8443 for more details.
+    fname = tmpdir.join('test.mmap').strpath
+    size = mmap.ALLOCATIONGRANULARITY
+    obj = [np.zeros(size, dtype='uint8'), np.ones(size, dtype='uint8')]
+    numpy_pickle.dump(obj, fname)
+    memmaps = numpy_pickle.load(fname, mmap_mode='r')
+    assert isinstance(memmaps[1], np.memmap)
+    assert memmaps[1].offset > size
+    np.testing.assert_array_equal(obj, memmaps)
diff --git a/joblib/test/test_parallel.py b/joblib/test/test_parallel.py
index 46e86fd..1ee29c9 100644
--- a/joblib/test/test_parallel.py
+++ b/joblib/test/test_parallel.py
@@ -13,7 +13,9 @@ from math import sqrt
 import threading
 from multiprocessing import TimeoutError
 from time import sleep
+import mmap
 
+from joblib import dump, load
 from joblib import parallel
 
 from joblib.test.common import np, with_numpy
@@ -729,3 +731,21 @@ def test_nested_parallel_warnings(capfd):
                        for _ in range(5))
     out, err = capfd.readouterr()
     assert 'Multiprocessing-backed parallel loops cannot be nested' in err
+
+
+def identity(arg):
+    return arg
+
+
+@with_numpy
+@with_multiprocessing
+def test_memmap_with_big_offset(tmpdir):
+    fname = tmpdir.join('test.mmap').strpath
+    size = mmap.ALLOCATIONGRANULARITY
+    obj = [np.zeros(size, dtype='uint8'), np.ones(size, dtype='uint8')]
+    dump(obj, fname)
+    memmap = load(fname, mmap_mode='r')
+    result, = Parallel(n_jobs=2)(delayed(identity)(memmap) for _ in [0])
+    assert isinstance(memmap[1], np.memmap)
+    assert memmap[1].offset > size
+    np.testing.assert_array_equal(obj, result)
diff --git a/joblib/test/test_pool.py b/joblib/test/test_pool.py
index 04f7dca..a24648d 100644
--- a/joblib/test/test_pool.py
+++ b/joblib/test/test_pool.py
@@ -1,4 +1,5 @@
 import os
+import mmap
 
 from joblib.test.common import with_numpy, np
 from joblib.test.common import setup_autokill
@@ -6,14 +7,14 @@ from joblib.test.common import teardown_autokill
 from joblib.test.common import with_multiprocessing
 from joblib.test.common import with_dev_shm
 from joblib.testing import raises
+from joblib.backports import make_memmap
 
-
-from joblib._multiprocessing_helpers import mp
-if mp is not None:
-    from joblib.pool import MemmapingPool
-    from joblib.pool import has_shareable_memory
-    from joblib.pool import ArrayMemmapReducer
-    from joblib.pool import reduce_memmap
+from joblib.pool import MemmapingPool
+from joblib.pool import has_shareable_memory
+from joblib.pool import ArrayMemmapReducer
+from joblib.pool import reduce_memmap
+from joblib.pool import _strided_from_memmap
+from joblib.pool import _get_backing_memmap
 
 
 def setup_module():
@@ -433,3 +434,47 @@ def test_workaround_against_bad_memmap_with_copied_buffers(tmpdir):
     finally:
         p.terminate()
         del p
+
+
+@with_numpy
+def test__strided_from_memmap(tmpdir):
+    fname = tmpdir.join('test.mmap').strpath
+    size = 5 * mmap.ALLOCATIONGRANULARITY
+    offset = mmap.ALLOCATIONGRANULARITY + 1
+    # This line creates the mmap file that is reused later
+    memmap_obj = np.memmap(fname, mode='w+', shape=size + offset)
+    # filename, dtype, mode, offset, order, shape, strides, total_buffer_len
+    memmap_obj = _strided_from_memmap(fname, dtype='uint8', mode='r',
+                                      offset=offset, order='C', shape=size,
+                                      strides=None, total_buffer_len=None)
+    assert isinstance(memmap_obj, np.memmap)
+    assert memmap_obj.offset == offset
+    memmap_backed_obj = _strided_from_memmap(fname, dtype='uint8', mode='r',
+                                             offset=offset, order='C',
+                                             shape=(size // 2,), strides=(2,),
+                                             total_buffer_len=size)
+    assert _get_backing_memmap(memmap_backed_obj).offset == offset
+
+
+def identity(arg):
+    return arg
+
+
+@with_numpy
+@with_multiprocessing
+def test_pool_memmap_with_big_offset(tmpdir):
+    # Test that numpy memmap offset is set correctly if greater than
+    # mmap.ALLOCATIONGRANULARITY, see
+    # https://github.com/joblib/joblib/issues/451 and
+    # https://github.com/numpy/numpy/pull/8443 for more details.
+    fname = tmpdir.join('test.mmap').strpath
+    size = 5 * mmap.ALLOCATIONGRANULARITY
+    offset = mmap.ALLOCATIONGRANULARITY + 1
+    obj = make_memmap(fname, mode='w+', shape=size, dtype='uint8',
+                      offset=offset)
+
+    p = MemmapingPool(2, temp_folder=tmpdir.strpath)
+    result = p.apply_async(identity, args=(obj,)).get()
+    assert isinstance(result, np.memmap)
+    assert result.offset == offset
+    np.testing.assert_array_equal(obj, result)