RawKernel does not work with dask.distributed

When I call a RawKernel using on a remote worker using distributed, e.g. ddf.map_partitions(kernel_wrapper), the worker crashes with a segfault. I suspect this is because of the following:

If I first execute the kernel on the client side, it compiles and runs successfully. After that, I submit the kernel for execution by the workers, at which point the NVRTC compiled code gets serialized. Once it gets sent to the other host, the CUDA context is different and the data structures of the compiled kernel are no longer valid. When I do not execute the kernel locally before submitting it to the workers (i.e. do not trigger compilation), everything runs fine.

• Conditions (you can just paste the output of python -c 'import cupy; cupy.show_config()')
  • CuPy version: master, commit:dd2fb09f63fce8b7ddb58deb4b91f0b3669d1a0d
  • OS/Platform: Linux PPCle64
  • CUDA version: 10.1.243
  • cuDNN/NCCL version N/A
• Code to reproduce

(sketch, I did not actually try to run the below code exactly as is)

import dask
import cupy as cp

kernel = cp.RawKernel(r'''

extern "C" __global__ void test_kernel()
    {
    printf("Test\n");
    }
''', 'test_kernel')

def launch_kernel():
    kernel((1,),(128,),())

# comment out the below line and it should run on the remote workers
launch_kernel() # should print out Test

from dask.distributed import Client
client = Client(scheduler_file='my-scheduler.json')
ddf = dask.datasets.timeseries()
ddf.map_partitions(launch_kernel) # Segfault on the workers

• Error messages, stack traces, or logs

distributed.nanny - INFO - Worker process 56138 was killed by signal 11
distributed.nanny - WARNING - Restarting worker

[leofang]

That is exactly my plan 👍

[jakirkham]

Thanks all! 😄