Skip to main content

Connect VS Code to SLURM Compute Nodes

How to Connect VS Code to SLURM Compute Nodes for Debugging and Jupyter Notebook Support

Working directly on Slurm login nodes can be cumbersome, particularly when running a debugger or performing interactive development work. The following guide presents a reliable solution developed through research and testing. By following the steps below, VS Code can be connected directly to a compute node, enabling debugging and interactive work within an isolated environment tied directly to your VS Code instance.

Setup

  1. Create SSH public key on your server (not your local machine). You'll be needing this key in the next step
    You can follow the tutorial here for this step.

  2. Setup sshd.job script
    Copy the sshd.job script into your root directory ~.
    Make sure to change the path to your key in the last part of the script ${HOME}/.ssh/<your_key_name>.

  3. Configure known hosts
    Start by opening ~/.ssh/config on your local machine.
    You may have a defined host for your login nodes. If you do, skip to the next part; if not, paste the following into your config file.

Host <slurm_name>
HostName <slurm_hostname>
User <username>
IdentityFile ~/.ssh/<your_key_name>
IdentitiesOnly yes

Then, we will paste the following configuration for going into the compute node. For <the compute_node_name> variable, see the next section.

Host <compute_node_name> !<slurm_hostname>  # notice the "!" sign. Do not remove it, it's not a typo.
User <username>
ProxyJump <slurm_hostname>
PreferredAuthentications publickey
ForwardAgent yes
IdentityFile ~/.ssh/<your_key_name>

An example setup

Host jane # jane is the slurm host's name
HostName jane.gpuserver.com # the ssh url to your server
User janedoe # your username
IdentityFile ~/.ssh/id_rsa_jane # your generated key's name
IdentitiesOnly yes

Host jane-compute-large-1 !jane # jane-compute-large-1 is the assigned compute node name
User janedoe
ProxyJump jane # "Host" from previous configuration
ForwardAgent yes
IdentityFile ~/.ssh/id_rsa_jane # your generated key's name
PreferredAuthentications publickey

Retrieving Compute Nodes

  1. Go into one of your cluster's login nodes. Simply type ssh <slurm_hostname> into your local terminal. Make sure that you are in a directory where sshd.job is visible. If you followed the steps above, it should be in your home directory.

  2. Once inside, use a command that is appropriate for your needs to retrieve a compute node. For example:

sbatch --gpus=rtx_4090:4 --mem-per-cpu=16G --time=48:00:00 --ntasks=8 ./sshd.job
-- or
sbatch --gpus=1 --gres=gpumem:80g --mem-per-cpu=32G --time=120:00:00 --ntasks=8 ./sshd.job

This will submit a job with sshd.job script as the target. You can now simply retrieve the compute node's name by typing squeue in your terminal inside the login node. It should give an output as follows. Simply grab the compute node name from this list and use in your configuration file.

JOBID PARTITION     NAME     USER    ST       TIME         NODES NODELIST(REASON)
<id>  gpuhe.120     sshd     <user>  R        0:00:01      1     <compute_node_name>

Using Copilot When the Network is Blocked on Compute Nodes

Some HPCs block the internet access in the compute nodes. Thus when you connect your vscode into these compute nodes, the copilot will not be able to connect as well. Luckily the solution is easy. Just put the following into your options.json file. Then, it'll magically start working.

"remote.extensionKind": {
    "GitHub.copilot": [
        "ui",
    ],

Working with compute nodes

One way is to directly use the shell by typing ssh <compute_node_name>. Though it's pointless at this point.
You can also use VSCode's remote development tool to connect to your compute node. The name will appear on the known hosts list. Simply click on it and you'll be good to go.

Note: this approach deviates from Slurm's intended usage pattern. However, it simplifies debugging and small-scale testing, particularly when working with VS Code. Accordingly, remember to cancel the job with scancel <job_id> once finished, and use proper sbatch commands for actual training runs, so that compute nodes remain available to everyone and are not held unnecessarily.

 

credit: GitHub