How to limit the usage of GPU to only one parallel worker?

Since the GPU memory is very limited, it can only undertake one job at a time, but much faster than the same job on CPU. Now I have 100 jobs to do, and want them to be done by the GPU as many as possible, and make full use of the slower but memory-rich CPUs.

Currently I've got no ways to ensure there being enough GPU memory free. I have to use try-catch block to gather the data back to CPU memory if an exception occurs. However, once an out-of-GPU-memory exception occurs on a worker, it will never free the GPU memory it consumed. As a result, if there're enough many jobs to do, each worker will hold its own part of useless GPU memory but not be able to do any actual GPU job because the GPU memory is full of useless occupancies.

That's why I start to consider how to fix the GPU to only one specific worker or two (if the GPU can hold 2 jobs in its memory). Or at least, I need a way to fully release GPU memory consumption if a worker decides not to use GPU for its job.

Answer 1

[Edric Ellis]

You could perhaps do something like this. This assumes your code is running in a parallel pool, and that you have only 1 GPU.

spmd
    if labindex == 1
        % worker 1 can use the GPU
        gpuDevice(1);
    else
        % All other workers must not use the GPU
        gpuDevice([]);
    end
end

% Now, use a parfor loop to do stuff.
parfor i = 1:100
    % Ask the GPUDeviceManager if any device is already selected
    isGPUSelected = ~isempty(parallel.gpu.GPUDeviceManager.instance.SelectedDevice);
    if isGPUSelected
        out{i} = runOnGPU(i);
    else
        % This worker has no GPU, run on CPU.
        out{i} = runOnCPU(i);
    end
end