Repository: joblib/joblib
Pull request #1561 (2024-04-04 17:20:53)

FIX Disable endianness alteration on unserialization of numpy arrays in `joblib.Parallel`
Fixes https://github.com/joblib/joblib/issues/1545

Still todo:

- [x] add non regression test from min reproducer in https://github.com/joblib/joblib/issues/1545
- [x] ~~discuss where exactly we should enable / disable this feature ? to what extent should the issues in https://github.com/joblib/joblib/issues/1123 also be considered for unserialization of auto memmaped array in `joblib.Parallel` ?~~
- [x] ~~discuss potential breaking change ?~~

fcharras: Regarding the second point, in fact it's not an issue because `_ensure_native_byte_order` is only used in `joblib.Parallel` calls when auto-memmapping is triggered. Byteorder for small arrays has always been preserved from main process to child process. I will add a unit test that show that in a separate PR.

fcharras: @alexisthual do you want to try this branch on your usecase and confirm that it fixes it ?
ogrisel: Also we should not forget to document the fix in the changelog.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/joblib/joblib/pull/1561?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
Attention: Patch coverage is `83.33333%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 70.09%. Comparing base [(`661a660`)](https://app.codecov.io/gh/joblib/joblib/commit/661a660d41a6c668ee2572c1d4bb6425fb0b81dc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) to head [(`15fa6a9`)](https://app.codecov.io/gh/joblib/joblib/commit/15fa6a96462017491a3fc805f6966109781dd4e7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).
> Report is 1 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/joblib/joblib/pull/1561?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | Patch % | Lines |
|---|---|---|
| [joblib/numpy\_pickle.py](https://app.codecov.io/gh/joblib/joblib/pull/1561?src=pr&el=tree&filepath=joblib%2Fnumpy_pickle.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL251bXB5X3BpY2tsZS5weQ==) | 85.71% | [3 Missing :warning: ](https://app.codecov.io/gh/joblib/joblib/pull/1561?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) |
| [joblib/test/test\_parallel.py](https://app.codecov.io/gh/joblib/joblib/pull/1561?src=pr&el=tree&filepath=joblib%2Ftest%2Ftest_parallel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9wYXJhbGxlbC5weQ==) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/joblib/joblib/pull/1561?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) |

> :exclamation:  There is a different number of reports uploaded between BASE (661a660) and HEAD (15fa6a9). Click for more details.
> 
> <details><summary>HEAD has 26 uploads less than BASE</summary>
>
>| Flag | BASE (661a660) | HEAD (15fa6a9) |
>|------|------|------|
>||27|1|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1561       +/-   ##
===========================================
- Coverage   95.74%   70.09%   -25.66%     
===========================================
  Files          46       46               
  Lines        7799     7837       +38     
===========================================
- Hits         7467     5493     -1974     
- Misses        332     2344     +2012     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/joblib/joblib/pull/1561?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib).

fcharras: @ogrisel thanks for the review.

Regarding the default values for occurences of `ensure_native_byte_order` parameter, I must admit I didn't try to see the big picture and lazily chose to set default values that reproduced the historical behavior, to ensure compatibility with code that could have been currently using those functions and could break if the default behavior changes.

After browsing the code more in depth, I think we can say that all the classes and functions that are changed in this PR (`NumpyPickler`, `NumpyArrayWrapper`, `_unpickle`,...) are private API, that we can break. Do you agree with that ?

Also I grep'ed joblib code and can say that those changes in `numpy_pickly.py` do not even affect other joblib modules.

So I understand that we _can_ break the default behavior of those internal functions. WDYT doing it a bit differently of what you're suggesting, and instead of setting default `ensure_native_byte_order=False`, not setting a default at all, so we have to explicitly set it wherever it's used ? since there is no harm in altering the signatures of internals, it will be even more explicit IMO.

> I have the feeling that it is saner to make it explicit in the code that the case ensure_native_byte_order and mmap_mode is not None is physically not achievable in general.

Can we re-discuss that in light of the previous commit too

Still todo:
- [ ]: document the fix in the changelog
- [ ]: an additional test with `joblib.load(filename, mmap_mode="r")`



Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 881ee56..a06cbbb 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -56,6 +56,18 @@ Maintenance:
   the warnings that are emitted in this case.
   https://github.com/joblib/joblib/pull/1681
 
+- Fix a regression in 1.3 and 1.4 that caused large big endian arrays to trigger
+  a serialization error. https://github.com/joblib/joblib/issues/1545
+
+- Added a ``ensure_native_byte_order`` parameter to ``joblib.load``. When
+  ``True`` and ``mmap_mode`` is ``None``, loaded arrays are automatically coerced
+  to a byte order that matches the endianness of the host system. This behavior
+  has been the default since ``joblib==1.3``, and can now be disabled if the
+  parameter is set to ``False`` instead. Note that setting it to ``True`` will
+  raise an error if ``mmap_mode`` is not null. The default value ``'auto'`` is
+  equivalent to always setting ``True`` if ``mmap_mode`` is ``None``, else always
+  ``False``.  https://github.com/joblib/joblib/pull/1561
+
 - Fix support for python 3.14 in ``hashing``, with the addition of
   an extra argument in ``Pickler._batch_setitems``.
   https://github.com/joblib/joblib/pull/1688
@@ -68,7 +80,6 @@ Maintenance:
   deadlocks.
 
 
-
 Release 1.4.2 -- 2024/05/02
 ---------------------------
 
diff --git a/continuous_integration/install.sh b/continuous_integration/install.sh
index fa8d974..b4b9e86 100755
--- a/continuous_integration/install.sh
+++ b/continuous_integration/install.sh
@@ -11,28 +11,16 @@ set -xe
 
 create_new_conda_env() {
     conda config --set solver libmamba
+    if [[ "$PYTHON_VERSION" == free-threaded* ]]; then
+        PYTHON_VERSION=${PYTHON_VERSION/free-threaded-/}
+        EXTRA_CONDA_PACKAGES="python-freethreading $EXTRA_CONDA_PACKAGES"
+    fi
     to_install="python=$PYTHON_VERSION pip pytest $EXTRA_CONDA_PACKAGES"
     conda create -n testenv --yes -c conda-forge $to_install
     conda activate testenv
 }
 
-create_new_free_threaded_env() {
-    sudo apt-get -yq update
-    sudo apt-get install -yq ccache
-    sudo apt-get install -yq software-properties-common
-    sudo add-apt-repository --yes ppa:deadsnakes/nightly
-    sudo apt-get update -yq
-    sudo apt-get install -yq --no-install-recommends python3.13-dev python3.13-venv python3.13-nogil
-
-    python3.13t -m venv testenv
-    source testenv/bin/activate
-}
-
-if [[ "$PYTHON_VERSION" == free-threaded* ]]; then
-    create_new_free_threaded_env
-else
-    create_new_conda_env
-fi
+create_new_conda_env
 
 # Install pytest timeout to fasten failure in deadlocking tests
 PIP_INSTALL_PACKAGES="pytest-timeout pytest-asyncio==0.21.1 threadpoolctl"
diff --git a/joblib/numpy_pickle.py b/joblib/numpy_pickle.py
index 8dc6b7a..46052c7 100644
--- a/joblib/numpy_pickle.py
+++ b/joblib/numpy_pickle.py
@@ -155,7 +155,7 @@ class NumpyArrayWrapper(object):
             ):
                 pickler.file_handle.write(chunk.tobytes("C"))
 
-    def read_array(self, unpickler):
+    def read_array(self, unpickler, ensure_native_byte_order):
         """Read array from unpickler file handle.
 
         This function is an adaptation of the numpy read_array function
@@ -205,8 +205,11 @@ class NumpyArrayWrapper(object):
             else:
                 array.shape = self.shape
 
-        # Detect byte order mismatch and swap as needed.
-        return _ensure_native_byte_order(array)
+        if ensure_native_byte_order:
+            # Detect byte order mismatch and swap as needed.
+            array = _ensure_native_byte_order(array)
+
+        return array
 
     def read_mmap(self, unpickler):
         """Read an array using numpy memmap."""
@@ -250,9 +253,9 @@ class NumpyArrayWrapper(object):
             )
             warnings.warn(message)
 
-        return _ensure_native_byte_order(marray)
+        return marray
 
-    def read(self, unpickler):
+    def read(self, unpickler, ensure_native_byte_order):
         """Read the array corresponding to this wrapper.
 
         Use the unpickler to get all information to correctly read the array.
@@ -260,6 +263,9 @@ class NumpyArrayWrapper(object):
         Parameters
         ----------
         unpickler: NumpyUnpickler
+        ensure_native_byte_order: bool
+            If true, coerce the array to use the native endianness of the
+            host system.
 
         Returns
         -------
@@ -268,9 +274,13 @@ class NumpyArrayWrapper(object):
         """
         # When requested, only use memmap mode if allowed.
         if unpickler.mmap_mode is not None and self.allow_mmap:
+            assert not ensure_native_byte_order, (
+                "Memmaps cannot be coerced to a given byte order, "
+                "this code path is impossible."
+            )
             array = self.read_mmap(unpickler)
         else:
-            array = self.read_array(unpickler)
+            array = self.read_array(unpickler, ensure_native_byte_order)
 
         # Manage array subclass case
         if hasattr(array, "__array_prepare__") and self.subclass not in (
@@ -395,6 +405,9 @@ class NumpyUnpickler(Unpickler):
         The memorymap mode to use for reading numpy arrays.
     file_handle: file_like
         File object to unpickle from.
+    ensure_native_byte_order: bool
+        If True, coerce the array to use the native endianness of the
+        host system.
     filename: str
         Name of the file to unpickle from. It should correspond to file_handle.
         This parameter is required when using mmap_mode.
@@ -405,7 +418,7 @@ class NumpyUnpickler(Unpickler):
 
     dispatch = Unpickler.dispatch.copy()
 
-    def __init__(self, filename, file_handle, mmap_mode=None):
+    def __init__(self, filename, file_handle, ensure_native_byte_order, mmap_mode=None):
         # The next line is for backward compatibility with pickle generated
         # with joblib versions less than 0.10.
         self._dirname = os.path.dirname(filename)
@@ -415,6 +428,7 @@ class NumpyUnpickler(Unpickler):
         # filename is required for numpy mmap mode.
         self.filename = filename
         self.compat_mode = False
+        self.ensure_native_byte_order = ensure_native_byte_order
         Unpickler.__init__(self, self.file_handle)
         try:
             import numpy as np
@@ -444,7 +458,11 @@ class NumpyUnpickler(Unpickler):
             # the end of the unpickling.
             if isinstance(array_wrapper, NDArrayWrapper):
                 self.compat_mode = True
-            self.stack.append(array_wrapper.read(self))
+                _array_payload = array_wrapper.read(self)
+            else:
+                _array_payload = array_wrapper.read(self, self.ensure_native_byte_order)
+
+            self.stack.append(_array_payload)
 
     # Be careful to register our new method.
     dispatch[pickle.BUILD[0]] = load_build
@@ -593,7 +611,7 @@ def dump(value, filename, compress=0, protocol=None):
     return [filename]
 
 
-def _unpickle(fobj, filename="", mmap_mode=None):
+def _unpickle(fobj, ensure_native_byte_order, filename="", mmap_mode=None):
     """Internal unpickling function."""
     # We are careful to open the file handle early and keep it open to
     # avoid race-conditions on renames.
@@ -601,7 +619,9 @@ def _unpickle(fobj, filename="", mmap_mode=None):
     # the case with the old persistence format, moving the directory
     # will create a race when joblib tries to access the companion
     # files.
-    unpickler = NumpyUnpickler(filename, fobj, mmap_mode=mmap_mode)
+    unpickler = NumpyUnpickler(
+        filename, fobj, ensure_native_byte_order, mmap_mode=mmap_mode
+    )
     obj = None
     try:
         obj = unpickler.load()
@@ -628,14 +648,31 @@ def _unpickle(fobj, filename="", mmap_mode=None):
 def load_temporary_memmap(filename, mmap_mode, unlink_on_gc_collect):
     from ._memmapping_reducer import JOBLIB_MMAPS, add_maybe_unlink_finalizer
 
-    obj = load(filename, mmap_mode)
+    with open(filename, "rb") as f:
+        with _validate_fileobject_and_memmap(f, filename, mmap_mode) as (
+            fobj,
+            validated_mmap_mode,
+        ):
+            # Memmap are used for interprocess communication, which should
+            # keep the objects untouched. We pass `ensure_native_byte_order=False`
+            # to remain consistent with the loading behavior of non-memmaped arrays
+            # in workers, where the byte order is preserved.
+            # Note that we do not implement endianness change for memmaps, as this
+            # would result in inconsistent behavior.
+            obj = _unpickle(
+                fobj,
+                ensure_native_byte_order=False,
+                filename=filename,
+                mmap_mode=validated_mmap_mode,
+            )
+
     JOBLIB_MMAPS.add(obj.filename)
     if unlink_on_gc_collect:
         add_maybe_unlink_finalizer(obj)
     return obj
 
 
-def load(filename, mmap_mode=None):
+def load(filename, mmap_mode=None, ensure_native_byte_order="auto"):
     """Reconstruct a Python object from a file persisted with joblib.dump.
 
     Read more in the :ref:`User Guide <persistence>`.
@@ -653,6 +690,12 @@ def load(filename, mmap_mode=None):
         mode has no effect for compressed files. Note that in this
         case the reconstructed object might no longer match exactly
         the originally pickled object.
+    ensure_native_byte_order: bool, or 'auto', default=='auto'
+        If True, ensures that the byte order of the loaded arrays matches the
+        native byte ordering (or _endianness_) of the host system. This is not
+        compatible with memory-mapped arrays and using non-null `mmap_mode`
+        parameter at the same time will raise an error. The default 'auto'
+        parameter is equivalent to True if `mmap_mode` is None, else False.
 
     Returns
     -------
@@ -672,6 +715,15 @@ def load(filename, mmap_mode=None):
     object might not match the original pickled object. Note that if the
     file was saved with compression, the arrays cannot be memmapped.
     """
+    if ensure_native_byte_order == "auto":
+        ensure_native_byte_order = mmap_mode is None
+
+    if ensure_native_byte_order and mmap_mode is not None:
+        raise ValueError(
+            "Native byte ordering can only be enforced if 'mmap_mode' parameter "
+            f"is set to None, but got 'mmap_mode={mmap_mode}' instead."
+        )
+
     if Path is not None and isinstance(filename, Path):
         filename = str(filename)
 
@@ -679,7 +731,7 @@ def load(filename, mmap_mode=None):
         fobj = filename
         filename = getattr(fobj, "name", "")
         with _validate_fileobject_and_memmap(fobj, filename, mmap_mode) as (fobj, _):
-            obj = _unpickle(fobj)
+            obj = _unpickle(fobj, ensure_native_byte_order=ensure_native_byte_order)
     else:
         with open(filename, "rb") as f:
             with _validate_fileobject_and_memmap(f, filename, mmap_mode) as (
@@ -692,5 +744,14 @@ def load(filename, mmap_mode=None):
                     # Joblib so we load it with joblib compatibility function.
                     return load_compatibility(fobj)
 
-                obj = _unpickle(fobj, filename, validated_mmap_mode)
+                # A memory-mapped array has to be mapped with the endianness
+                # it has been written with. Other arrays are coerced to the
+                # native endianness of the host system.
+                obj = _unpickle(
+                    fobj,
+                    ensure_native_byte_order=ensure_native_byte_order,
+                    filename=filename,
+                    mmap_mode=validated_mmap_mode,
+                )
+
     return obj


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index d774560..ba96689 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -85,7 +85,7 @@ jobs:
             os: ubuntu-latest
             PYTHON_VERSION: "3.9"
             JOBLIB_TESTS_DEFAULT_PARALLEL_BACKEND: "threading"
-          - name: python-nogil
+          - name: python-free-threaded
             os: 'ubuntu-latest'
             PYTHON_VERSION: "free-threaded-3.13"
 
diff --git a/continuous_integration/run_tests.sh b/continuous_integration/run_tests.sh
index 95dab0e..a0c0ba0 100755
--- a/continuous_integration/run_tests.sh
+++ b/continuous_integration/run_tests.sh
@@ -3,16 +3,17 @@
 set -xe
 
 echo "Activating test environment:"
+conda activate testenv
+
 if [[ "$PYTHON_VERSION" == free-threaded* ]]; then
-    source testenv/bin/activate
     # This is needed because for now some C extensions have not declared their
     # thread-safety with free-threaded Python, for example numpy and coverage.tracer
     export PYTHON_GIL=0
-else
-    conda activate testenv
 fi
+
 which python
-python -V
+# Show python version and build information (e.g. free-threaded or not)
+python -VV
 python -c "import multiprocessing as mp; print('multiprocessing.cpu_count():', mp.cpu_count())"
 python -c "import joblib; print('joblib.cpu_count():', joblib.cpu_count())"
 
diff --git a/joblib/test/test_numpy_pickle.py b/joblib/test/test_numpy_pickle.py
index 6c61710..9bbd927 100644
--- a/joblib/test/test_numpy_pickle.py
+++ b/joblib/test/test_numpy_pickle.py
@@ -347,6 +347,46 @@ def test_compressed_pickle_dump_and_load(tmpdir):
             assert result == expected
 
 
+@with_numpy
+def test_memmap_load(tmpdir):
+    little_endian_dtype = np.dtype("<i8")
+    big_endian_dtype = np.dtype(">i8")
+    all_dtypes = (little_endian_dtype, big_endian_dtype)
+
+    le_array = np.arange(5, dtype=little_endian_dtype)
+    be_array = np.arange(5, dtype=big_endian_dtype)
+
+    fname = tmpdir.join("temp.pkl").strpath
+
+    numpy_pickle.dump([le_array, be_array], fname)
+
+    le_array_native_load, be_array_native_load = numpy_pickle.load(
+        fname, ensure_native_byte_order=True
+    )
+
+    assert le_array_native_load.dtype == be_array_native_load.dtype
+    assert le_array_native_load.dtype in all_dtypes
+
+    le_array_nonnative_load, be_array_nonnative_load = numpy_pickle.load(
+        fname, ensure_native_byte_order=False
+    )
+
+    assert le_array_nonnative_load.dtype == le_array.dtype
+    assert be_array_nonnative_load.dtype == be_array.dtype
+
+
+def test_invalid_parameters_raise():
+    expected_msg = (
+        "Native byte ordering can only be enforced if 'mmap_mode' parameter "
+        "is set to None, but got 'mmap_mode=r+' instead."
+    )
+
+    with raises(ValueError, match=re.escape(expected_msg)):
+        numpy_pickle.load(
+            "/path/to/some/dump.pkl", ensure_native_byte_order=True, mmap_mode="r+"
+        )
+
+
 def _check_pickle(filename, expected_list, mmap_mode=None):
     """Helper function to test joblib pickle content.
 
diff --git a/joblib/test/test_parallel.py b/joblib/test/test_parallel.py
index 01da874..0f8a15a 100644
--- a/joblib/test/test_parallel.py
+++ b/joblib/test/test_parallel.py
@@ -421,7 +421,8 @@ def test_parallel_pickling():
 @with_numpy
 @with_multiprocessing
 @parametrize("byteorder", ["<", ">", "="])
-def test_parallel_byteorder_corruption(byteorder):
+@parametrize("max_nbytes", [1, "1M"])
+def test_parallel_byteorder_corruption(byteorder, max_nbytes):
     def inspect_byteorder(x):
         return x, x.dtype.byteorder
 
@@ -429,7 +430,7 @@ def test_parallel_byteorder_corruption(byteorder):
 
     initial_np_byteorder = x.dtype.byteorder
 
-    result = Parallel(n_jobs=2, backend="loky")(
+    result = Parallel(n_jobs=2, backend="loky", max_nbytes=max_nbytes)(
         delayed(inspect_byteorder)(x) for _ in range(3)
     )