Keep import cudf working after a GPU read

src/zarr_kvikio/_pipeline.py

@@ -456,9 +456,7 @@ def release_memory() -> None:
                 # when a synchronize above failed (e.g. after a CUDA error). Otherwise no decode
                 # runs now (new ones wait for the lock), and `_pool` sets it again.
                 if _pinned_pool:
-                    import nvidia.nvcomp as nv
-
-                    nv.set_pinned_allocator(None)
+                    _nvcomp().set_pinned_allocator(None)
                     _pinned_pool.pop().free_all_blocks()
 
 
@@ -560,11 +558,26 @@ def reserve(self, fds: int) -> Iterator[None]:
 _open_files = _FileBudget()
 
 
+def _nvcomp() -> Any:
+    """`nvidia.nvcomp`, imported without breaking a later `import cudf`.
+
+    nvCOMP's `__init__` deletes the `libnvcomp` attribute of the `nvidia` package but leaves the
+    module in `sys.modules`, so libcudf's loader (`import nvidia.libnvcomp`, then
+    `nvidia.libnvcomp.load_library()`) raises AttributeError unless it is put back.
+    """
+    import nvidia
+    import nvidia.nvcomp as nv  # (`as` works even while another thread's import finishes)
+
+    if (lib := sys.modules.get("nvidia.libnvcomp")) is not None:
+        nvidia.libnvcomp = lib
+    return nv
+
+
 def _pool(device_id: int) -> queue.Queue[tuple[Any, Any] | None]:
     """The pool of nvCOMP zstd codecs for a device, each with the CUDA stream it runs on."""
     import cupy as cp
-    import nvidia.nvcomp as nv
 
+    nv = _nvcomp()
     with _pools_lock:
         if device_id not in _pools:
             # (with no codecs, decodes would wait for one forever)
@@ -621,8 +634,7 @@ def _decode_into(
 
     Unless `check` is False (see `_check_frame_sizes`), raises if a frame decodes to another size.
     """
-    import nvidia.nvcomp as nv
-
+    nv = _nvcomp()
     limit = _max_decode_bytes()
     groups: list[tuple[int, int]] = []
     start = size = 0

tests/test_pipeline.py

@@ -377,6 +377,18 @@ def test_release_memory(tmp_path, read_path):
     check(tmp_path / "a.zarr", slice(None))  # a new codec is created on demand
 
 
+def test_reads_keep_cudf_importable(tmp_path, read_path):
+    """nvCOMP's import deletes `nvidia.libnvcomp`, which libcudf's loader (`import cudf`) uses."""
+    if read_path == "host":
+        pytest.skip("host reads don't import nvCOMP")
+    pytest.importorskip("nvidia.libnvcomp")
+    write(tmp_path / "a.zarr", rng.random(100).astype("float32"), (10,))
+    check(tmp_path / "a.zarr", slice(None))
+    import nvidia.libnvcomp  # (as libcudf's loader does)
+
+    nvidia.libnvcomp.load_library()
+
+
 def test_enable(tmp_path, fast_path_calls):
     data = rng.random(1000).astype("float32")
     write(tmp_path / "a.zarr", data, (100,), shards=(500,))