Make sure arrays are bytes aligned in joblib pickles

joblib/numpy_pickle.py

@@ -7,6 +7,7 @@
 import pickle
 import os
 import warnings
+import io
 from pathlib import Path
 
 from .compressor import lz4, LZ4_NOT_INSTALLED_ERROR
@@ -39,6 +40,8 @@
 ###############################################################################
 # Utility objects for persistence.
 
+NUMPY_ARRAY_ALIGNMENT_BYTES = 16
+
 
 class NumpyArrayWrapper(object):
     """An object to be persisted instead of numpy arrays.
@@ -92,6 +95,17 @@ def write_array(self, array, pickler):
             # pickle protocol.
             pickle.dump(array, pickler.file_handle, protocol=2)
         else:
+            try:
+                current_pos = pickler.file_handle.tell()
+                alignment = current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES
+
+                if alignment != 0:
+                    padding = b' ' * (NUMPY_ARRAY_ALIGNMENT_BYTES - alignment)
+                    pickler.file_handle.write(padding)
+            except io.UnsupportedOperation:
+                # TODO log something somewhere?
+                pass
+
             for chunk in pickler.np.nditer(array,
                                            flags=['external_loop',
                                                   'buffered',
@@ -118,6 +132,21 @@ def read_array(self, unpickler):
             # The array contained Python objects. We need to unpickle the data.
             array = pickle.load(unpickler.file_handle)
         else:
+            try:
+                current_pos = unpickler.file_handle.tell()
+                alignment = current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES
+
+                # peek not supported in io.BytesIO ...
+                current_byte = unpickler.file_handle.read(1)
+                unpickler.file_handle.seek(current_pos)
+
+                if alignment != 0 and current_byte == b' ':
+                    padding_length = NUMPY_ARRAY_ALIGNMENT_BYTES - alignment
+                    unpickler.file_handle.seek(current_pos + padding_length)
+            except io.UnsupportedOperation:
+                # TODO log something somewhere?
+                pass
+
             # This is not a real file. We have to read it the
             # memory-intensive way.
             # crc32 module fails on reads greater than 2 ** 32 bytes,
@@ -150,7 +179,18 @@ def read_array(self, unpickler):
 
     def read_mmap(self, unpickler):
         """Read an array using numpy memmap."""
-        offset = unpickler.file_handle.tell()
+        current_pos = unpickler.file_handle.tell()
+        offset = current_pos
+        alignment = current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES
+
+        # peek not supported in io.BytesIO ...
+        current_byte = unpickler.file_handle.read(1)
+        unpickler.file_handle.seek(current_pos)
+
+        if alignment != 0 and current_byte == b' ':
+            padding_length = NUMPY_ARRAY_ALIGNMENT_BYTES - alignment
+            offset += padding_length
+
         if unpickler.mmap_mode == 'w+':
             unpickler.mmap_mode = 'r+'
 

joblib/test/test_numpy_pickle.py

@@ -465,6 +465,27 @@ def test_joblib_pickle_across_python_versions():
         _check_pickle(fname, expected_list)
 
 
+@with_numpy
+def test_joblib_pickle_across_python_versions_with_mmap():
+    expected_list = [np.arange(5, dtype=np.dtype('<i8')),
+                     np.arange(5, dtype=np.dtype('<f8')),
+                     np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'),
+                     np.arange(256, dtype=np.uint8).tobytes(),
+                     # np.matrix is a subclass of np.ndarray, here we want
+                     # to verify this type of object is correctly unpickled
+                     # among versions.
+                     np.matrix([0, 1, 2], dtype=np.dtype('<i8')),
+                     u"C'est l'\xe9t\xe9 !"]
+
+    test_data_dir = os.path.dirname(os.path.abspath(data.__file__))
+
+    pickle_filenames = [
+        os.path.join(test_data_dir, fn)
+        for fn in os.listdir(test_data_dir) if fn.endswith('.pkl')]
+    for fname in pickle_filenames:
+        _check_pickle(fname, expected_list)
+
+
 @with_numpy
 def test_numpy_array_byte_order_mismatch_detection():
     # List of numpy arrays with big endian byteorder.
@@ -1054,3 +1075,55 @@ def test_lz4_compression_without_lz4(tmpdir):
     with raises(ValueError) as excinfo:
         numpy_pickle.dump(data, fname + '.lz4')
     excinfo.match(msg)
+
+
+@with_numpy
+@parametrize('protocol', range(0, pickle.HIGHEST_PROTOCOL + 1))
+def test_memmap_alignment_padding(tmpdir, protocol):
+    # Test that memmaped arrays returned by numpy.load are correctly aligned
+    fname = tmpdir.join('test.mmap').strpath
+
+    a = np.random.randn(2)
+    numpy_pickle.dump(a, fname, protocol=protocol)
+    memmap = numpy_pickle.load(fname, mmap_mode='r')
+    assert isinstance(memmap, np.memmap)
+    np.testing.assert_array_equal(a, memmap)
+    assert (
+        memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
+    assert memmap.flags.aligned
+
+    array_list = [
+        np.random.randn(2), np.random.randn(2),
+        np.random.randn(2), np.random.randn(2)
+    ]
+
+    # On Windows OSError 22 if reusing the same path for memmap ...
+    fname = tmpdir.join('test1.mmap').strpath
+    numpy_pickle.dump(array_list, fname, protocol=protocol)
+    l_reloaded = numpy_pickle.load(fname, mmap_mode='r')
+
+    for idx, memmap in enumerate(l_reloaded):
+        assert isinstance(memmap, np.memmap)
+        np.testing.assert_array_equal(array_list[idx], memmap)
+        assert (
+            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
+        assert memmap.flags.aligned
+
+    array_dict = {
+        'a1': np.random.randn(100),
+        'a2': np.random.randn(200),
+        'a3': np.random.randn(300),
+        'a4': np.random.randn(400)
+    }
+
+    # On Windows OSError 22 if reusing the same path for memmap ...
+    fname = tmpdir.join('test2.mmap').strpath
+    numpy_pickle.dump(array_dict, fname, protocol=protocol)
+    d_reloaded = numpy_pickle.load(fname, mmap_mode='r')
+
+    for key, memmap in d_reloaded.items():
+        assert isinstance(memmap, np.memmap)
+        np.testing.assert_array_equal(array_dict[key], memmap)
+        assert (
+            memmap.ctypes.data % numpy_pickle.NUMPY_ARRAY_ALIGNMENT_BYTES == 0)
+        assert memmap.flags.aligned