Can a user restrict a slurm job to only run on specific GPUS?

Viewed 136

I have a workstation with 4 GPUs. I've setup my machine as a single-node slurm instance. I use slurm to automate complex job pipelines in the background while I work in the foreground. So I'm not doing much HPC, I'm just using SLURM to handle job DAGs.

I have a use-case where sometimes I would like to reserve a specific GPU for my own interactive use. I could do this by simply removing one of the GPUs from my /etc/slurm-llnl/gres.conf, but I'd like to avoid this. I don't like using sudo on my system in my day-to-day.

Instead, I'd like a way to either:

  1. Tell sbatch exactly which GPUs it is allowed to use.
  2. Tell sbatch which GPUs it is not allowed to use.

Either one of those would be fine and solve my problem.

Research I've done so far:

I know there is an option in sbatch for --constraint, but that seems to be a node-level feature, and all of my jobs necessarily will execute on my host machine because it is the only slurm node in the network.

I see there is a --gpu-bind, option in the slurm manpage. The bit at the end ("Not supported unless the entire node is allocated to the job.") makes me think this isn't what I want. But I did try to test it out.

I submitted several copies of this job quickly in a row so the queue was full of at least 6 of them:

sbatch \
    --job-name="tester1" \
    --output="test-job-%j-%x.out" \
    --cpus-per-task=1 --mem=1000 --gres="gpu:1" \
    --gpu-bind "map_gpu:2,3" \
    --wrap="python -c \"import torch, os; print(os.getenv('CUDA_VISIBLE_DEVICES', 'x')) and torch.rand(1000).to(0)\""

But the output indicated that GPU 0 and 1 were still allocated to the job, so either I'm doing it wrong, or this is not the way to go.

Is there any way for me to request that slurm jobs only run on GPU's 2 and 3, even though 0, and 1 are also in the gres?

For reference my gres.conf file is:

NAME=gpu File=/dev/nvidia0
NAME=gpu File=/dev/nvidia1
NAME=gpu File=/dev/nvidia2
NAME=gpu File=/dev/nvidia3

And I would be ok with a solution where I had to edit this file once to add some "tag" property to each GPU and then request that sbatch runs the jobs on a GPU with that tag. Not sure how to do that though, or if that is possible.

0 Answers
Related