Pickle round trip of a compiled ufunc unloads the original's kernel module (CUDA_ERROR_INVALID_HANDLE) #10334

Description

After a ufunc (or an ElementwiseKernel) has run once, a pickle round trip of the object, followed by freeing the copy, breaks the original: every later call raises

cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_INVALID_HANDLE: invalid resource handle

from cupy.cuda.function._launch (driver.launchKernel). When the original's module is freed in turn (in the runs below, for some of the objects before the process exits), CuPy also reports

Exception ignored in: 'cupy.cuda.function.Module.__dealloc__'
cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_INVALID_HANDLE: invalid resource handle

from cuModuleUnload. The same round trip before the first call is harmless. The expected behaviour is that pickling an object and discarding the copy leaves the original unchanged, as it does for NumPy's ufuncs.

The ufunc pickles through Cython's automatic pickling, by value, including its per-type kernel cache (_kernel_memo), whose values are Function objects holding a Module. cupy.cuda.function.Module is also auto-pickled and pickles as its raw CUmodule handle (ptr). The unpickled copy therefore owns the same module as the original, and its __dealloc__ unloads it when the copy is freed. cupy/_core/_kernel.pyx and cupy/cuda/function.pyx on main still define no pickling of their own for these classes.

This shows up in practice through dask: dask.base.tokenize pickles and unpickles every callable object that is not a plain function (dask/tokenize.py, _normalize_pickle). So dask.array.map_blocks(cupy.abs, x) breaks cupy.abs for the rest of the process once it has run on that dtype. Merely building the graph is enough, whatever the scheduler.

To Reproduce

import pickle
import cupy as cp

x = cp.arange(4, dtype=cp.complex64)
cp.abs(x)                                  # compiles and caches the kernel
copy = pickle.loads(pickle.dumps(cp.abs))  # e.g. what dask.base.tokenize does
del copy                                   # the copy unloads the shared CUmodule
cp.abs(x)                                  # CUDADriverError: CUDA_ERROR_INVALID_HANDLE

Each case below ran in a fresh process (the kernel caches are per process): call the object, round trip, call again ("compile first"), or round trip, then call ("round trip first").

Rows are the kernel objects; the columns are the two orders; each cell is the outcome of the last call. The table holds for CuPy 14.2.0 on an NVIDIA A100 (environment below), where the failing call raises CUDA_ERROR_INVALID_HANDLE and, for exp and the ElementwiseKernel, the destructor error also appears in the same process. It holds identically for CuPy 13.5.1 on ROCm 7.2 (Radeon 8050S, gfx1151), where the failing call raises hipErrorInvalidDevice: invalid device ordinal, and a dask graph over the same kernel segfaults the process. The dask route (map_blocks(cupy.abs, x) after cupy.abs(x), graph built but not computed) fails the same way on CUDA with CuPy 14.2.0 and 13.5.1.

Installation

Wheel (pip install cupy-***)

Environment

OS                           : Linux-6.4.0-150600.23.125_15.0.29-cray_shasta_c-x86_64-with-glibc2.38
Python Version               : 3.14.7
CuPy Version                 : 14.2.0
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 2.5.3
SciPy Version                : None
Cython Build Version         : 3.2.4
Cython Runtime Version       : None
CUDA Root                    : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/cuda/12.9/targets/x86_64-linux
NVCC PATH                    : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/cuda/12.9/bin/nvcc
CUDA Build Version           : 12090
CUDA Driver Version          : 13000
CUDA Runtime Version         : 12090 (linked to CuPy) / 12090 (locally installed)
CUDA Runtime Path            : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/cuda/12.9/lib64/libcudart.so.12 (locally installed)
CUDA Extra Include Dirs      : ['/opt/nvidia/hpc_sdk/Linux_x86_64/25.5/cuda/12.9/include']
cuBLAS Version               : 120900
cuBLAS Path                  : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcublas.so.12
cuFFT Version                : 11400
cuFFT Path                   : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcufft.so.11
cuRAND Version               : 10310
cuRAND Path                  : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcurand.so.10
cuSOLVER Version             : (11, 7, 4)
cuSOLVER Path                : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcusolver.so.11
cuSPARSE Version             : 12509
cuSPARSE Path                : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcusparse.so.12
NVRTC Version                : (12, 9)
NVRTC Path                   : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/cuda/12.9/lib64/libnvrtc.so.12
Thrust Version               : 300102
CUB Build Version            : 300102
Jitify Build Version         : <unknown>
NCCL Build Version           : None
NCCL Runtime Version         : None
NCCL Path                    : None
cuTENSOR Version             : 20200
cuTENSOR Path                : /opt/nvidia/hpc_sdk/Linux_x86_64/25.5/math_libs/12.9/lib64/libcutensor.so.2
cuSPARSELt Build Version     : None
cuSPARSELt Path              : None
Device 0 Name                : NVIDIA A100-SXM4-80GB
Device 0 Compute Capability  : 80
Device 0 PCI Bus ID          : 0000:03:00.0

Additional Information

Possible fixes: give ufunc (and ElementwiseKernel) a __reduce__ that pickles by reference, as NumPy's ufuncs do (the unpickled object is the same one), or at least leave the kernel caches out of the pickled state; and make Module/Function non-picklable, since a raw CUmodule handle has no meaning outside the process that loaded it.

Found through abTEM (https://github.com/abTEM/abTEM), whose lazy element-wise methods passed cupy.abs to map_blocks; it now passes the function's name and resolves it per block.

🤖 Written by Claude Code — Paul reviewed and posted it