Keep codecs and filters registered from Python on one thread

src/blosc2/blosc2_ext.pyx

@@ -3580,7 +3580,26 @@ cdef int64_t _store_data(blosc2_schunk *schunk, _write_op op, int64_t nchunk, da
 def get_blocksize():
     return blosc1_get_blocksize()
 
+cdef inline c_bool _runs_python_callbacks(const blosc2_cparams *cparams) noexcept:
+    """Whether `cparams` name a codec or filter registered from Python."""
+    return _python_codec_ids[cparams.compcode] or _uses_python_filter(cparams)
+
+
+cdef _check_blosc_nthreads():
+    """Codecs and filters registered from Python need one-thread contexts, but C-Blosc2 gives every
+    context BLOSC_NTHREADS threads when that is set (also the ones asking for one), and their callbacks
+    would then run on C-Blosc2's worker threads, without the GIL: raise instead.  Checked where they are
+    registered and where contexts for their data are set up, since the variable can change at run time."""
+    nthreads = context_nthreads(1)
+    if nthreads > 1:
+        raise ValueError(f"BLOSC_NTHREADS={nthreads} is not supported with codecs or filters registered "
+                         f"from Python: C-Blosc2 would run them on {nthreads} threads, but they need one "
+                         "(unset BLOSC_NTHREADS or set it to 1)")
+
+
 cdef _check_cparams(blosc2_cparams *cparams):
+    if _runs_python_callbacks(cparams):
+        _check_blosc_nthreads()
     if cparams.nthreads > 1:
         if _python_codec_ids[cparams.compcode]:
             raise ValueError("Cannot use multi-threading with user defined Python codecs")
@@ -3593,6 +3612,8 @@ cdef _check_cparams(blosc2_cparams *cparams):
 cdef _check_dparams(blosc2_dparams* dparams, blosc2_cparams* cparams=NULL):
     if cparams == NULL:
         return
+    if _runs_python_callbacks(cparams):
+        _check_blosc_nthreads()
     if dparams.nthreads > 1:
         if _python_codec_ids[cparams.compcode]:
             raise ValueError("Cannot use multi-threading with user defined Python codecs")
@@ -3602,6 +3623,31 @@ cdef _check_dparams(blosc2_dparams* dparams, blosc2_cparams* cparams=NULL):
             raise ValueError("`nthreads` must be 1 when a postfilter is set")
 
 
+cdef _one_thread_for_python_callbacks(blosc2_schunk *sc):
+    """Give `sc` one-thread contexts if its codec or a filter is registered from Python, the rule that
+    _check_cparams() and _check_dparams() enforce for new containers.  Super-chunks made from frames or
+    opened from files start with C-Blosc2's global number of threads, which would run the Python
+    callbacks on C-Blosc2's worker threads, without the GIL.  Raises if BLOSC_NTHREADS is above one."""
+    cdef blosc2_context *ctx
+    if not _runs_python_callbacks(sc.storage.cparams):
+        return
+    _check_blosc_nthreads()
+    if sc.storage.cparams.nthreads != 1:
+        sc.storage.cparams.nthreads = 1
+        ctx = blosc2_create_cctx(dereference(sc.storage.cparams))
+        if ctx == NULL:
+            raise RuntimeError("Could not create compression context")
+        blosc2_free_ctx(sc.cctx)
+        sc.cctx = ctx
+    if sc.storage.dparams.nthreads != 1:
+        sc.storage.dparams.nthreads = 1
+        ctx = blosc2_create_dctx(dereference(sc.storage.dparams))
+        if ctx == NULL:
+            raise RuntimeError("Could not create decompression context")
+        blosc2_free_ctx(sc.dctx)
+        sc.dctx = ctx
+
+
 cdef create_cparams_from_kwargs(blosc2_cparams *cparams, kwargs):
     if "compcode" in kwargs:
         raise NameError("`compcode` has been renamed to `codec`. Please go update your code.")
@@ -4154,6 +4200,7 @@ cdef class SChunk:
                 check_schunk_params(self.schunk, kwargs)
             if schunk_is_ndarray(self.schunk):
                 raise ValueError("Cannot open an NDArray as a SChunk. Please use blosc2.open instead")
+            _one_thread_for_python_callbacks(self.schunk)
         else:
             self.schunk = blosc2_schunk_new(&storage)
 
@@ -5882,6 +5929,8 @@ def open(urlpath, mode, offset, **kwargs):
         check_schunk_params(schunk, kwargs)
     cparams = kwargs.get("cparams")
     dparams = kwargs.get("dparams", blosc2.DParams())
+    # As _one_thread_for_python_callbacks(), for the parameters set below: they default to blosc2.nthreads
+    python_callbacks = _runs_python_callbacks(schunk.storage.cparams)
 
     if is_ndarray:
         res = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, <char *> "blosc2_schunk*", NULL),
@@ -5896,9 +5945,14 @@ def open(urlpath, mode, offset, **kwargs):
         # A super-chunk opened in mode "w" starts empty (an NDArray keeps its chunks, so its layout)
         schunk_obj.cparams = cparams if isinstance(cparams, blosc2.CParams) else blosc2.CParams(**cparams)
     else:
-        schunk_obj.cparams = _opened_cparams(schunk_obj.cparams, cparams)
+        opened_cparams = _opened_cparams(schunk_obj.cparams, cparams)
+        if python_callbacks:
+            opened_cparams = dataclasses.replace(opened_cparams, nthreads=1)
+        schunk_obj.cparams = opened_cparams
     if dparams is not None:
-        schunk_obj.dparams = dparams if isinstance(dparams, blosc2.DParams) else blosc2.DParams(**dparams)
+        dparams = dparams if isinstance(dparams, blosc2.DParams) else blosc2.DParams(**dparams)
+        schunk_obj.dparams = dataclasses.replace(dparams, nthreads=1) if python_callbacks else dparams
+    _one_thread_for_python_callbacks((<SChunk> schunk_obj).schunk)  # also for dparams=None
 
     return res
 
@@ -6020,6 +6074,7 @@ def schunk_from_cframe(cframe, copy=False):
     PyBuffer_Release(&buf)
     if not copy:
         schunk._avoid_cframe_free(True)
+    _one_thread_for_python_callbacks(schunk_)
     return schunk
 
 
@@ -6075,6 +6130,8 @@ def register_codec(codec_name, id, encoder=None, decoder=None, version=1):
 
     if (encoder is None and decoder is not None) or (encoder is not None and decoder is None):
         raise ValueError("both encoder and decoder must be given, or none")
+    if encoder is not None:
+        _check_blosc_nthreads()
 
     cdef blosc2_codec codec
     codec.compcode = id
@@ -6151,6 +6208,8 @@ def register_filter(id, forward, backward, filter_name):
                          " and ", BLOSC2_USER_REGISTERED_FILTERS_STOP)
     if (forward is None and backward is not None) or (forward is not None and backward is None):
         raise ValueError("both encoder and decoder must be given, or none")
+    if forward is not None:
+        _check_blosc_nthreads()
 
     cdef blosc2_filter filter
     filter.id = id
@@ -7948,6 +8007,7 @@ def ndarray_from_cframe(cframe, copy=False):
     PyBuffer_Release(&buf)
     if not copy:
         ndarray._schunk._avoid_cframe_free(True)
+    _one_thread_for_python_callbacks(array.sc)
     return ndarray
 
 

src/blosc2/core.py

@@ -2521,7 +2521,9 @@ def register_codec(
 
     Notes
     -----
-    * Cannot use multi-threading when using a user-defined codec.
+    * Cannot use multi-threading when using a user-defined codec. A codec with Python
+      functions cannot be registered while the ``BLOSC_NTHREADS`` environment variable is
+      above 1, since C-Blosc2 would use that many threads for it.
 
     * User-defined codecs can only be used inside an `SChunk` instance.
 
@@ -2599,7 +2601,9 @@ def register_filter(
 
     Notes
     -----
-    * Multi-threading cannot be used with a user-defined filter.
+    * Multi-threading cannot be used with a user-defined filter. A filter with Python
+      functions cannot be registered while the ``BLOSC_NTHREADS`` environment variable is
+      above 1, since C-Blosc2 would use that many threads for it.
 
     * User-defined filters can only be used inside an `SChunk` instance.
 

tests/test_python_callbacks_nthreads.py

@@ -0,0 +1,186 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#######################################################################
+
+# Codecs and filters registered from Python need the GIL, so C-Blosc2 must use one thread for them,
+# also where it picks the number of threads itself: containers made from frames or opened from
+# files start with the global number of threads, and BLOSC_NTHREADS replaces the nthreads of every
+# context.  On failure the callbacks run on C-Blosc2's worker threads and crash the interpreter, so
+# every case runs in a child interpreter.
+
+import os
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+import blosc2
+
+pytestmark = pytest.mark.skipif(blosc2.IS_WASM, reason="no subprocesses and threads on emscripten/wasm32")
+
+CHILD = textwrap.dedent(
+    """
+    import os
+    import sys
+    import threading
+
+    import numpy as np
+
+    import blosc2
+
+    kind, source, urlpath = sys.argv[1:4]
+    ID = 170  # of the codec or the filter, in the range for user ids
+    MAIN = threading.get_ident()
+
+
+    def check_thread():
+        # Python code on a C-Blosc2 worker thread runs without the GIL: it often crashes, but not always
+        if threading.get_ident() != MAIN:
+            os._exit(70)
+
+
+    def encoder(input, output, meta, schunk):
+        check_thread()
+        # 2 bytes per item: the values are below 2 ** 16
+        n = input.size // 4
+        output[: 2 * n] = input.view(np.int32).astype(np.uint16).view(np.uint8)
+        return 2 * n
+
+
+    def decoder(input, output, meta, schunk):
+        check_thread()
+        n = output.size // 4
+        output[:] = input[: 2 * n].view(np.uint16).astype(np.int32).view(np.uint8)
+        return output.size
+
+
+    def reverse(input, output, meta, schunk):
+        check_thread()
+        output[:] = input[::-1]
+
+
+    try:
+        if kind == "codec":
+            blosc2.register_codec("nthreads_codec", ID, encoder, decoder)
+            cparams = {"codec": ID, "filters": [blosc2.Filter.NOFILTER], "filters_meta": [0]}
+        else:
+            blosc2.register_filter(ID, reverse, reverse, "nthreads_filter")
+            cparams = {"codec": blosc2.Codec.LZ4, "filters": [ID], "filters_meta": [0]}
+    except ValueError as e:
+        print("refused:", e)
+        sys.exit(0)
+    # 16 blocks per chunk, so that C-Blosc2 would use its worker threads
+    N = 16 * 1000
+    cparams.update(typesize=4, nthreads=1, blocksize=4000, splitmode=blosc2.SplitMode.NEVER_SPLIT)
+    values = (np.arange(3 * N) % 1000).astype(np.int32)
+    if source == "env_later":
+        # BLOSC_NTHREADS set after the registration: C-Blosc2 reads it for every new context
+        dparams = {"nthreads": 1}
+        s = blosc2.SChunk(chunksize=4 * N, data=values, cparams=cparams, dparams=dparams, urlpath=urlpath)
+        cframe = s.to_cframe()
+        os.environ["BLOSC_NTHREADS"] = "4"
+        if blosc2.blosc2_ext.context_nthreads(1) != 4:
+            # This C-Blosc2 does not see variables set at run time (another C runtime), so it cannot use it
+            print("ok")
+            sys.exit(0)
+
+        def set_dparams():
+            s.dparams = blosc2.DParams(nthreads=1)
+            return s
+
+        for make in (lambda: blosc2.SChunk(chunksize=4 * N, data=values, cparams=cparams, dparams=dparams),
+                     lambda: blosc2.schunk_from_cframe(cframe, copy=True),
+                     lambda: blosc2.open(urlpath),
+                     set_dparams):
+            try:
+                make().decompress_chunk(0)
+                print("not refused")
+            except ValueError as e:
+                assert "BLOSC_NTHREADS=4" in str(e), e
+    elif source.startswith("ndarray"):
+        a = blosc2.asarray(values, chunks=(N,), blocks=(N // 16,), cparams=cparams, dparams={"nthreads": 1},
+                           urlpath=urlpath)
+        blosc2.set_nthreads(4)
+        if source == "ndarray_frame":
+            a = blosc2.ndarray_from_cframe(a.to_cframe(), copy=True)
+        else:
+            a = blosc2.open(urlpath, mode="a")
+        assert np.array_equal(a[:], values)
+        a[:N] = values[N : 2 * N]
+        assert np.array_equal(a[:N], values[N : 2 * N])
+    else:
+        s = blosc2.SChunk(chunksize=4 * N, data=values, cparams=cparams, dparams={"nthreads": 1},
+                          urlpath=urlpath)
+        blosc2.set_nthreads(4)
+        if source == "frame":
+            s = blosc2.schunk_from_cframe(s.to_cframe(), copy=True)
+        elif source == "open":
+            s = blosc2.open(urlpath, mode="a")
+        elif source == "read_mode":
+            s = blosc2.SChunk(urlpath=urlpath, mode="r")
+        for i in range(3):
+            chunk = np.frombuffer(s.decompress_chunk(i), np.int32)
+            assert np.array_equal(chunk, values[i * N : (i + 1) * N])
+        if source != "read_mode":  # unrelated bug: that SChunk gets chunksize 2 ** 24, get_slice() fails
+            out = np.empty_like(values)
+            s.get_slice(out=out)
+            assert np.array_equal(out, values)
+        if source in ("frame", "open"):
+            s.append_data(values[:N])
+            assert np.array_equal(np.frombuffer(s.decompress_chunk(3), np.int32), values[:N])
+    print("ok")
+    """
+)
+
+
+@pytest.mark.parametrize(
+    ("kind", "source"),
+    [
+        # Every source with a codec, a filter where the check runs in another place
+        ("codec", "frame"),
+        ("codec", "open"),
+        ("codec", "read_mode"),
+        ("codec", "ndarray_frame"),
+        ("codec", "ndarray_open"),
+        ("filter", "frame"),
+        ("filter", "open"),
+        ("codec", "env_later"),
+        ("filter", "env_later"),
+    ],
+)
+def test_global_nthreads(kind, source, tmp_path):
+    # Containers made from frames get C-Blosc2's global number of threads, opened ones blosc2.nthreads
+    # (here both 4): they must use one
+    env = os.environ.copy()
+    env.pop("BLOSC_NTHREADS", None)
+    result = subprocess.run(
+        [sys.executable, "-c", CHILD, kind, source, str(tmp_path / "nthreads.b2frame")],
+        capture_output=True,
+        text=True,
+        env=env,
+        timeout=300,
+    )
+    assert result.returncode == 0, (result.returncode, result.stderr[-2000:])
+    assert result.stdout.split() == ["ok"]
+
+
+@pytest.mark.parametrize("kind", ["codec", "filter"])
+def test_register_with_blosc_nthreads(kind, tmp_path):
+    # C-Blosc2 gives every context BLOSC_NTHREADS threads, also the ones asking for one: registering a
+    # Python codec or filter must fail then, instead of crashing later (on main: in SChunk(data=...))
+    env = os.environ.copy()
+    env["BLOSC_NTHREADS"] = "4"
+    result = subprocess.run(
+        [sys.executable, "-c", CHILD, kind, "frame", str(tmp_path / "nthreads.b2frame")],
+        capture_output=True,
+        text=True,
+        env=env,
+        timeout=300,
+    )
+    assert result.returncode == 0, (result.returncode, result.stderr[-2000:])
+    assert result.stdout.startswith("refused:"), result.stdout
+    assert "BLOSC_NTHREADS=4" in result.stdout