Feature Request: Running multiple jobs across GPUs on a single node

[Andrew-S-Rosen]

I am looking to run jobs on Perlmutter at NERSC, which has GPU nodes with 4x A100s. I would like to run each Jobflow Job on a single A100, ideally in a batch mode that continually pulls in new work over the course of the job. Is there a mechanism in place to distribute the jobs over the GPU node? In practice, this would involve separate srun calls as described here. I'm not dealing with anything MPI or OpenMP related here since this is all PyTorch code.

Is this practical with the parallel batch mode? The parallel batch mode seems like it distributes work over individual nodes within a multi-node allocation, but I'm not sure if it can be made to distribute work over multiple GPUs within each node as well.

[gpetretto]

Hi @Andrew-S-Rosen, I am not sure if I understand correctly your use case. As an example, I suppose you would like to reserve a full node and use a batch mode to run 4 Jobs simultaneously, each using a different GPU.
I never tried such a case myself, but I think that this may be already feasible with the current batch mode. How to do that depends on how the pytorch code is called.
In the case of a VASP job, the python code calls an external executable. In that case one can set atomate2_VASP_CMD="srun --nodes 1 -n 24 --exclusive vasp_std" (see https://matgenix.github.io/jobflow-remote/user/advancedoptions.html#parallel-batch). So, if you are calling the pytorch code as an external executable, I suppose that you could have an equivalent call with srun --gpus 1 EXECUTABLE. I don't see any particular limitation related to the fact of using a single node instead of multiple ones. As long as srun can distribute the jobs correctly to different GPUs, this should not be a problem. When I used it the --exclusive flag was necessary to avoid all the jobs being assigned to the same node. I don't know if slurm will be smart enough to keep the jobs in separate GPUs or if there is a similar flag that could be set. But this would be independent from jobflow-remote.

If instead the pytorch code is already included in the Job execution, this is a bit trickier, but still quite easy to do. As you may have seen, when executed in batch mode, in the Job execution folder a submit.sh is still created and it is then called as an external subprocess:

jobflow-remote/src/jobflow_remote/jobs/run.py

Line 244 in a9346a9

["bash", "submit.sh"], # noqa: S603, S607

The content of the script is the same as any other job and it contains the line:

jf -fe execution run /path/to/run/jobflow/81/76/86/817686b0-6ffe-4d8f-bb98-7afdb5a8a952_1

For completely different purposes (i.e. containers), it has been introduced the option to customize this command. In the worker options you can set the execution_cmd

jobflow-remote/src/jobflow_remote/config/base.py

Line 173 in a9346a9

execution_cmd: Optional[str] = Field(

I think this could be reused in your case by setting it to: srun --gpus 1 jf -fe execution run {}.
I cannot really guarantee that it will work, as I never tested something like this. It may also depend on the code you are using, but I believe it is worth a try. (the same consideration about whether the --gpus 1 flag is enough to keep the jobs on separate gpus are valid here as well.)

[Andrew-S-Rosen]

@gpetretto: Thank you for the super helpful reply. Really appreciate it!

I am not sure if I understand correctly your use case. As an example, I suppose you would like to reserve a full node and use a batch mode to run 4 Jobs simultaneously, each using a different GPU.

You are correct. That's what I am describing. However, because I am not launching an external executable (e.g. VASP), I am not making a specific srun call like in ATOMATE2_VASP_CMD. If I were, then indeed exactly what you wrote would make sense and "should just work." Instead, I am wrapping a pure Python function with a @job decorator.

But exactly what you suggested --- setting the execution_cmd --- seems like it would work. Here is how it is normally done for context. I will definitely have my group give this a go, although probably not for a few weeks while we get everything ready.

Much appreciated!

[gpetretto]

Good to hear that this may be useful.
Please, let me know if the solution actually works, so that I can add this as a use case in the documentation.

[Andrew-S-Rosen]

FYI: We are still definitely still going to get to checking this out since it is needed for running our jobs on Perlmutter, but we are running a bit behind schedule. No timeline to promise but hopefully soon!