Here's a complete toy example (tested on JupyterLab using a SLURM cluster).
The example compiles with Cython a trivial function that sums two integers, but of course one can apply the same technique to complex (and more useful) code.
The key trick here is that one has to set up the Workers to find and import the Cython library.
This requires importing pyximport
, calling pyximport.install()
, and then importing the Cython-generated module on each Worker. This is done using register_worker_callback()
.
Note that the Cython-generated module is placed in the <IPYTHONCACHEDIR/cython
directory (IPYTHONCACHEDIR
can be found by calling IPython.paths.get_ipython_cache_dir()
). The directory must be added to the paths where Python looks for modules, so that the Cython-generated module can be loaded.
This example assumes SLURM, but it is just for my convenience.
The dask.distributed "network" could be set up with any other method (see for instance http://distributed.dask.org/en/latest/setup.html).
from dask import delayed
%load_ext cython
# Create a toy Cython function and put it into a module named remoteCython
%%cython -n remoteCython
def cython_sum(int a, int b):
return a+b
# Set up a distributed cluster (minimal, just for illustration)
# I use SLURM.
from dask_jobqueue import SLURMCluster
from distributed import Client
cluster = SLURMCluster(memory="1GB",
processes=1,
cores=1,
walltime="00:10:00")
cluster.start_workers(1) # Start as many workers as needed.
client = Client(cluster)
def init_pyx(dask_worker):
import pyximport
pyximport.install()
import sys
sys.path.insert(0,'<IPYTHONCACHEDIR>/cython/') # <<< replace <IPYTHONCACHEDIR> as appropriate
import remoteCython
client.register_worker_callbacks(init_pyx) # This runs init_pyx() on any Worker at init
import remoteCython
# ASIDE: you can find where the full path of Cython-generated library by
# looking at remoteCython.__file__
# The following creates a task and submits to the scheduler.
# The task computes the sum of 123 and 321 via the Cython function defined above
future = client.compute(delayed(remoteCython.cython_sum)(123,321))
# The task is executed on the remote worker
# We fetch the result from the remote worker
print(future.result()) # This prints 444
# We're done. Let's release the SLURM jobs.
cluster.close()