How can we specify GPUs for offscreen rendering?

[keiohta]

Hi author, thanks for the continual support!
I'm trying to generate a dataset using multiprocessing, where each process has a SAPIEN engine and renderer and takes a lot of images. I'm using a 32 CPU cores machine with 8 GPUs, but I realized that all processes have access to the same single GPU, and that causes some errors probably related to the GPU memory. Ideally, I want to assign 4 cores to 1 GPU (results in 32 CPU cores to 8 GPUs), but I'm not sure how I can do this. So, could you tell me if SAPIEN has such functionality (e.g., specify GPU id) or any workaround that resolves my situation?

[fbxiang]

The easiest way is to specify CUDA_VISIBLE_DEVICES and SAPIEN renderer will respect it.

[keiohta]

Yes, I agree that explicitly specifying CUDA_VISIBLE_DEVICES is the easiest solution, but I think SAPIEN only respects the CUDA_VISIBLE_DEVICES before calling python, and thus we can't assign different devices on different processes. So, it would be great if we can explicitly assign GPU ids to SAPIEN engines/renderers, like:

def func(gpu_id):
  engine = Engine(gpu_id=gpu_id)

I know the current SAPIEN doesn't support this function, but I'm wondering if there are some workarounds to enable something like the above. For example, pytorch can do this by using torch.cuda.set_device

def func(gpu_id):
  torch.cuda.set_device(gpu_id)

[fbxiang]

First you can assign the environment variable before creating a subprocess. This is how we do things most of the time.
Next you can also choose to dynamically assign the CUDA device by passing device=cuda:x to the VulkanRenderer as an argument. I personally prefer the environment variable as it often ensures consistency across different python packages.

[keiohta]

Thanks for the tips. I'm using a bit old SAPIEN version (1.1) thus no support for passing device, so I'd do with the first one.
Thank you!