Repository: joblib/joblib

Issue:
Alignment issue when passing arrays though mmap
Arrays that are transfered to subprocesses though mmap appear not always to be aligned:

```python
import joblib
import numpy as np

def func(x):
    print(type(x))
    print(x.flags)
    print(x.ctypes.data / 8)

if __name__ == '__main__':
    # Make the array larger than 1Mbyte so that mmap is used
    n = 1024 ** 2 // 8 + 1
    x = np.random.randn(n)
    jobs = joblib.Parallel(2)
    print(x.nbytes)
    res = jobs([joblib.delayed(func)(x), joblib.delayed(func)(x)])
```

This returns
```
<class 'numpy.core.memmap.memmap'>
1048584
  C_CONTIGUOUS : True
  F_CONTIGUOUS : True
  OWNDATA : False
  WRITEABLE : False
  ALIGNED : False
  UPDATEIFCOPY : False
570054683.375
```

Note the last line, the address of the first element of the array is indeed not aligned.

I'm using version 0.11, on OSx sierra.

This came up in https://github.com/pymc-devs/pymc3/issues/2640, and seems to happen on several OSs.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
> Merging [#1254](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) (6b72bcf) into [master](https://codecov.io/gh/joblib/joblib/commit/7742f5882273889f7aaf1d483a8a1c72a97d57e3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) (7742f58) will **decrease** coverage by `19.60%`.
> The diff coverage is `4.65%`.

[![Impacted file tree graph](https://codecov.io/gh/joblib/joblib/pull/1254/graphs/tree.svg?width=650&height=150&src=pr&token=gA6LF5DGTW&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)

```diff
@@             Coverage Diff             @@
##           master    #1254       +/-   ##
===========================================
- Coverage   88.43%   68.83%   -19.61%     
===========================================
  Files          47       47               
  Lines        7082     6917      -165     
===========================================
- Hits         6263     4761     -1502     
- Misses        819     2156     +1337     
```


| [Impacted Files](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | Coverage Δ | |
|---|---|---|
| [joblib/numpy\_pickle.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL251bXB5X3BpY2tsZS5weQ==) | `58.06% <0.00%> (-39.96%)` | :arrow_down: |
| [joblib/test/test\_numpy\_pickle.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9udW1weV9waWNrbGUucHk=) | `43.31% <7.40%> (-50.02%)` | :arrow_down: |
| [joblib/test/test\_memmapping.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9tZW1tYXBwaW5nLnB5) | `23.44% <0.00%> (-75.80%)` | :arrow_down: |
| [joblib/backports.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL2JhY2twb3J0cy5weQ==) | `28.94% <0.00%> (-55.27%)` | :arrow_down: |
| [joblib/test/test\_hashing.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9oYXNoaW5nLnB5) | `49.26% <0.00%> (-49.82%)` | :arrow_down: |
| [joblib/numpy\_pickle\_compat.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL251bXB5X3BpY2tsZV9jb21wYXQucHk=) | `45.54% <0.00%> (-45.55%)` | :arrow_down: |
| [joblib/\_memmapping\_reducer.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19tZW1tYXBwaW5nX3JlZHVjZXIucHk=) | `52.65% <0.00%> (-43.96%)` | :arrow_down: |
| [joblib/format\_stack.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL2Zvcm1hdF9zdGFjay5weQ==) | `58.33% <0.00%> (-41.67%)` | :arrow_down: |
| [joblib/my\_exceptions.py](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL215X2V4Y2VwdGlvbnMucHk=) | `64.28% <0.00%> (-35.72%)` | :arrow_down: |
| ... and [30 more](https://codecov.io/gh/joblib/joblib/pull/1254/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib). Last update [7742f58...6b72bcf](https://codecov.io/gh/joblib/joblib/pull/1254?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).

lesteve: > Instead I think we should adopt an explicit [prefix-code](https://en.wikipedia.org/wiki/Prefix_code) of the padding size at write time (without the try / except io.UnsupportedOperation boilerplate for readability's sake):

That does seem cleaner indeed. The only problem I see is that for old pickles (i.e. written without the change in this PR), there is no prefix code (the array bytes starts right away) so you are going to interpret the first array byte as the prefix code. Somehow there needs to be some kind of "sentinel" byte to indicate whether you have a new or old pickle file.
ogrisel: > That does seem cleaner indeed. The only problem I see is that for old pickles (i.e. written without the change in this PR), there is no prefix code (the array bytes starts right away) so you are going to interpret the first array byte as the prefix code.

This is a valid concern indeed. Good catch.

> Somehow there needs to be some kind of "sentinel" byte to indicate whether you have a new or old pickle file.

Yes we need to put some joblib-dump format versioning info somewhere at the beginning of the pickle file. But I don't think we can do it without breaking the pickle format. But we are already breaking the pickle format so maybe this is not a problem.
lesteve: I think this is ready enough for a more thorough review, removing the draft status.

lesteve: The failure is on MacOS `test_child_raises_parent_exits_cleanly[multiprocessing]` probably unrelated to this PR. The build [log](https://dev.azure.com/joblib/joblib/_build/results?buildId=1803&view=logs&j=ee31ff9a-0e27-5056-2892-3ca864802f2f&t=75da501c-a363-5ac9-fd11-59ca96cb91ff&l=1318)
ogrisel: It might also be worth pinging the people who reported or linked issues in various third party projects to the original bug report #563.
lesteve: > Could you please try to open a DRAFT PR in scikit-learn to run its CI against this branch of joblib with a revert of this workaround: [scikit-learn/scikit-learn#21654 (files)](https://github.com/scikit-learn/scikit-learn/pull/21654/files) ? In particular trigger a `[cd build]` to build an test scikit-learn with old GCC versions from the old manylinux docker images.

This has been done and it looks like this branch does fix the segmentation fault seen in scikit-learn: https://github.com/scikit-learn/scikit-learn/pull/22607#issuecomment-1050998451.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index c5c428a..56c8e76 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -16,6 +16,18 @@ Development version
   worker processes in case of a crash.
   https://github.com/joblib/joblib/pull/1269
 
+- Fix memory alignment bug for pickles containing numpy arrays.
+  This is especially important when loading the pickle with
+  ``mmap_mode != None`` as the resulting ``numpy.memmap`` object
+  would not be able to correct the misalignment without performing
+  a memory copy.
+  This bug would cause invalid computation and segmentation faults
+  with native code that would directly access the underlying data
+  buffer of a numpy array, for instance C/C++/Cython code compiled
+  with older GCC versions or some old OpenBLAS written in platform
+  specific assembly.
+  https://github.com/joblib/joblib/pull/1254
+
 Release 1.1.0
 --------------
 
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index 881d1f0..fa450fb 100644
--- a/joblib/numpy_pickle.py
+++ b/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,11 @@ register_compressor('lz4', LZ4CompressorWrapper())
 ###############################################################################
 # Utility objects for persistence.
 
+# For convenience, 16 bytes are used to be sure to cover all the possible
+# dtypes' alignments. For reference, see:
+# https://numpy.org/devdocs/dev/alignment.html
+NUMPY_ARRAY_ALIGNMENT_BYTES = 16
+
 
 class NumpyArrayWrapper(object):
     """An object to be persisted instead of numpy arrays.
@@ -70,13 +76,23 @@ class NumpyArrayWrapper(object):
         Default: False.
     """
 
-    def __init__(self, subclass, shape, order, dtype, allow_mmap=False):
+    def __init__(self, subclass, shape, order, dtype, allow_mmap=False,
+                 numpy_array_alignment_bytes=NUMPY_ARRAY_ALIGNMENT_BYTES):
         """Constructor. Store the useful information for later."""
         self.subclass = subclass
         self.shape = shape
         self.order = order
         self.dtype = dtype
         self.allow_mmap = allow_mmap
+        # We make numpy_array_alignment_bytes an instance attribute to allow us
+        # to change our mind about the default alignment and still load the old
+        # pickles (with the previous alignment) correctly
+        self.numpy_array_alignment_bytes = numpy_array_alignment_bytes
+
+    def safe_get_numpy_array_alignment_bytes(self):
+        # NumpyArrayWrapper instances loaded from joblib <= 1.1 pickles don't
+        # have an numpy_array_alignment_bytes attribute
+        return getattr(self, 'numpy_array_alignment_bytes', None)
 
     def write_array(self, array, pickler):
         """Write array bytes to pickler file handle.
@@ -92,6 +108,23 @@ class NumpyArrayWrapper(object):
             # pickle protocol.
             pickle.dump(array, pickler.file_handle, protocol=2)
         else:
+            numpy_array_alignment_bytes = \
+                self.safe_get_numpy_array_alignment_bytes()
+            if numpy_array_alignment_bytes is not None:
+                current_pos = pickler.file_handle.tell()
+                pos_after_padding_byte = current_pos + 1
+                padding_length = numpy_array_alignment_bytes - (
+                    pos_after_padding_byte % numpy_array_alignment_bytes)
+                # A single byte is written that contains the padding length in
+                # bytes
+                padding_length_byte = int.to_bytes(
+                    padding_length, length=1, byteorder='little')
+                pickler.file_handle.write(padding_length_byte)
+
+                if padding_length != 0:
+                    padding = b'\xff' * padding_length
+                    pickler.file_handle.write(padding)
+
             for chunk in pickler.np.nditer(array,
                                            flags=['external_loop',
                                                   'buffered',
@@ -118,6 +151,15 @@ class NumpyArrayWrapper(object):
             # The array contained Python objects. We need to unpickle the data.
             array = pickle.load(unpickler.file_handle)
         else:
+            numpy_array_alignment_bytes = \
+                self.safe_get_numpy_array_alignment_bytes()
+            if numpy_array_alignment_bytes is not None:
+                padding_byte = unpickler.file_handle.read(1)
+                padding_length = int.from_bytes(
+                    padding_byte, byteorder='little')
+                if padding_length != 0:
+                    unpickler.file_handle.read(padding_length)
+
             # 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 +192,17 @@ class NumpyArrayWrapper(object):
 
     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
+        numpy_array_alignment_bytes = \
+            self.safe_get_numpy_array_alignment_bytes()
+
+        if numpy_array_alignment_bytes is not None:
+            padding_byte = unpickler.file_handle.read(1)
+            padding_length = int.from_bytes(padding_byte, byteorder='little')
+            # + 1 is for the padding byte
+            offset += padding_length + 1
+
         if unpickler.mmap_mode == 'w+':
             unpickler.mmap_mode = 'r+'
 
@@ -163,6 +215,20 @@ class NumpyArrayWrapper(object):
         # update the offset so that it corresponds to the end of the read array
         unpickler.file_handle.seek(offset + marray.nbytes)
 
+        if (numpy_array_alignment_bytes is None and
+                current_pos % NUMPY_ARRAY_ALIGNMENT_BYTES != 0):
+            message = (
+                f'The memmapped array {marray} loaded from the file '
+                f'{unpickler.file_handle.name} is not not bytes aligned. '
+                'This may cause segmentation faults if this memmapped array '
+                'is used in some libraries like BLAS or PyTorch. '
+                'To get rid of this warning, regenerate your pickle file '
+                'with joblib >= 1.2.0. '
+                'See https://github.com/joblib/joblib/issues/563 '
+                'for more details'
+            )
+            warnings.warn(message)
+
         return marray
 
     def read(self, unpickler):
@@ -239,9 +305,17 @@ class NumpyPickler(Pickler):
         order = 'F' if (array.flags.f_contiguous and
                         not array.flags.c_contiguous) else 'C'
         allow_mmap = not self.buffered and not array.dtype.hasobject
+
+        kwargs = {}
+        try:
+            self.file_handle.tell()
+        except io.UnsupportedOperation:
+            kwargs = {'numpy_array_alignment_bytes': None}
+
         wrapper = NumpyArrayWrapper(type(array),
                                     array.shape, order, array.dtype,
-                                    allow_mmap=allow_mmap)
+                                    allow_mmap=allow_mmap,
+                                    **kwargs)
 
         return wrapper
 


Test patch:
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index 8019ca5..cce16e3 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -370,7 +370,7 @@ def test_compressed_pickle_dump_and_load(tmpdir):
             assert result == expected
 
 
-def _check_pickle(filename, expected_list):
+def _check_pickle(filename, expected_list, mmap_mode=None):
     """Helper function to test joblib pickle content.
 
     Note: currently only pickles containing an iterable are supported
@@ -390,17 +390,36 @@ def _check_pickle(filename, expected_list):
                 warnings.filterwarnings(
                     'ignore', module='numpy',
                     message='The compiler package is deprecated')
-                result_list = numpy_pickle.load(filename)
+                result_list = numpy_pickle.load(filename, mmap_mode=mmap_mode)
             filename_base = os.path.basename(filename)
-            expected_nb_warnings = 1 if ("_0.9" in filename_base or
-                                         "_0.8.4" in filename_base) else 0
+            expected_nb_deprecation_warnings = 1 if (
+                "_0.9" in filename_base or "_0.8.4" in filename_base) else 0
+
+            expected_nb_user_warnings = 3 if (
+                re.search("_0.1.+.pkl$", filename_base) and
+                mmap_mode is not None) else 0
+            expected_nb_warnings = \
+                expected_nb_deprecation_warnings + expected_nb_user_warnings
             assert len(warninfo) == expected_nb_warnings
-            for w in warninfo:
-                assert w.category == DeprecationWarning
+
+            deprecation_warnings = [
+                w for w in warninfo if issubclass(
+                    w.category, DeprecationWarning)]
+            user_warnings = [
+                w for w in warninfo if issubclass(
+                    w.category, UserWarning)]
+            for w in deprecation_warnings:
                 assert (str(w.message) ==
                         "The file '{0}' has been generated with a joblib "
                         "version less than 0.10. Please regenerate this "
                         "pickle file.".format(filename))
+
+            for w in user_warnings:
+                escaped_filename = re.escape(filename)
+                assert re.search(
+                    f"memmapped.+{escaped_filename}.+segmentation fault",
+                    str(w.message))
+
             for result, expected in zip(result_list, expected_list):
                 if isinstance(expected, np.ndarray):
                     expected = _ensure_native_byte_order(expected)
@@ -467,6 +486,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, mmap_mode='r')
+
+
 @with_numpy
 def test_numpy_array_byte_order_mismatch_detection():
     # List of numpy arrays with big endian byteorder.
@@ -916,6 +956,17 @@ def test_pickle_in_socket():
 
     np.testing.assert_array_equal(array_reloaded, test_array)
 
+    # Check that a byte-aligned numpy array written in a file can be send over
+    # a socket and then read on the other side
+    bytes_to_send = io.BytesIO()
+    numpy_pickle.dump(test_array, bytes_to_send)
+    server.send(bytes_to_send.getvalue())
+
+    with client.makefile("rb") as cf:
+        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):
@@ -1056,3 +1107,65 @@ def test_lz4_compression_without_lz4(tmpdir):
     with raises(ValueError) as excinfo:
         numpy_pickle.dump(data, fname + '.lz4')
     excinfo.match(msg)
+
+
+protocols = [pickle.DEFAULT_PROTOCOL]
+if pickle.HIGHEST_PROTOCOL != pickle.DEFAULT_PROTOCOL:
+    protocols.append(pickle.HIGHEST_PROTOCOL)
+
+
+@with_numpy
+@parametrize('protocol', protocols)
+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 = {
+        'a0': np.arange(2, dtype=np.uint8),
+        'a1': np.arange(3, dtype=np.uint8),
+        'a2': np.arange(5, dtype=np.uint8),
+        'a3': np.arange(7, dtype=np.uint8),
+        'a4': np.arange(11, dtype=np.uint8),
+        'a5': np.arange(13, dtype=np.uint8),
+        'a6': np.arange(17, dtype=np.uint8),
+        'a7': np.arange(19, dtype=np.uint8),
+        'a8': np.arange(23, dtype=np.uint8),
+    }
+
+    # 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