memory handling with dask-cuda on a windows machine - dask

Question

I am currently exploring how to handle memory in dask-cuda in order to write a function that will interpolate values along lines that cross an image.
My machine is a very basic windows 10 laptop with a single gpu (GeForce GTX 1050 4GB memory) and 16GB of RAM. I am using the following packages:
cupy 10.2.0
cudatoolkit 11.6.0
dask 2022.1.0
dask-cuda 22.2
A minimal version of my code is as follows:
import cupy as cp
import numpy as np
import dask.array as da
from dask_cuda import LocalCUDACluster
from dask.distributed import Client
from dask import compute
from math import pi
def bilinear_interpolate(im, x, y):
x0 = np.floor(x).astype(np.int32)
x1 = x0 + 1
y0 = np.floor(y).astype(np.int32)
y1 = y0 + 1
# keep coordinates within bounds
x0 = np.clip(x0, 0, im.shape[1]-1)
x1 = np.clip(x1, 0, im.shape[1]-1)
y0 = np.clip(y0, 0, im.shape[0]-1)
y1 = np.clip(y1, 0, im.shape[0]-1)
# retrieve values
Ia = im[ y0, x0 ]
Ib = im[ y1, x0 ]
Ic = im[ y0, x1 ]
Id = im[ y1, x1 ]
# calculate weights
wa = (x1-x) * (y1-y)
wb = (x1-x) * (y-y0)
wc = (x-x0) * (y1-y)
wd = (x-x0) * (y-y0)
return (wa*Ia + wb*Ib + wc*Ic + wd*Id).astype(np.float32)
# image used 5000x5000
im = cp.asarray(im)
# Set cuda-dask Client
cluster = LocalCUDACluster(
CUDA_VISIBLE_DEVICES = "0",
device_memory_limit = 0.4,
jit_unspill= True
)
client = Client(cluster)
# location
loc = im.shape[0] // 2, im.shape[1] //2
# radius
radius = [1,3000]
# pts per line
pts_line = radius[1]-radius[0]
# lines per degree factor
factor = 200
# radial chunks
rchunks = 180
# generate lines
r = cp.linspace(radius[0], radius[1], pts_line , dtype=np.int32)
psi = cp.linspace(0, 2*pi, 360*factor, dtype=np.float32)
R, PSI = cp.meshgrid(r, psi, sparse= True)
dR = da.from_array(R, chunks=(-1,pts_line), asarray=False).astype(np.float32)
dPSI = da.from_array(PSI, chunks=(rchunks,-1), asarray=False).astype(np.float32)
# Use polar coordinate to generate lines
rr = loc[0] + dR*np.cos(dPSI)
cc = loc[1] + dR* np.sin(dPSI)
rslt= da.map_blocks(bilinear_interpolate, im, cc, rr)
z = cp.asnumpy(rslt.compute().T)
I have managed, through trial and error, to interpolate +90k lines (each with 3000 points) across a single band image with 5000x5000 in size. However, as mentioned above, I am trying to write this function so that it will compute regardless of image size and number of lines to interpolate.
Much of the information I have found does not refer to windows machines (e.g. it appears that one cannot set rmm_pool_limit in a LocalCUDACluster as RAPIDS rmm package only works in linux). I am also familiar, thanks to the video by Mads Kristensen, with the limitations of the device_memory_limit parameter (i.e., it being a soft target) and the use of jit_unspill = True parameter as a way to minimize GPU memory spikes, etc. Yet, despite maintaining chuck sizes way below what I have available as GPU memory, I am still running into out of memory errors.
What is the best way to determine memory requirements? I had imagine that chuck size was the most crucial aspect (I thought that regardless of the original size of the array provided the chunk size was way below GPU memory limits I would be okay).