Running VS Code server over SSH using SLURM (without overburdening the login node)

VS Code has an extension, Remote SSH, which makes it possible to develop remotely by establishing and connecting to a remote VS Code server over SSH. You can edit files, open terminal shells, even run jupyter notebooks all in VS Code. It’s really nice.

The tricky part: clusters have login nodes which are the only way to ssh in, but are resource constrained. To avoid using up too many resources on the login node, you shouldn’t be running your VS Code server on it. Also, if you want to run a jupyter notebook in VSCode, then you really need to do it on a compute node.

The typical solution is to get an interactive job for a compute node, and run code in there. But how can you run your VS Code server (and all the accompanying terminals, jupyter notebooks, used when using VS Code) on a compute node? SSH-ing directly into the compute node isn’t possible. Instead, you have to set up a proxy through the login node.

Here’s how to connect VS Code server over SSH to a compute node using SLURM:

Step 1: connect to login node via SSH:

ssh sca63@g2-login-05.coecis.cornell.edu

Step 2: get an interactive job that will run your code session:

srun --nodes=1 --gres=gpu:1 --cpus-per-task=8 --time=04:00:00 --mem=50G --partition=default_partition-interactive --pty bash

Step 3: add the compute node’s name to your list of hosts in your ssh config file (~/.ssh/config) with something like this:

Host <compute-node-name>
  User <ssh username>
  ProxyCommand ssh <ssh username@<ssh login address> -W <compute-node-name>:%p

For example, it might look like:

Host ellis-compute-01
  User sca63
  ProxyCommand ssh sca63@g2-login-05.coecis.cornell.edu -W ellis-compute-01:%p

Step 4: In VS Code, open the command palette, select Remote-SSH: Connect to Host..., choose your compute node name. Then you should be good to go!

Addendum: Making things easier

If you’re like me, you’re on a cluster with a bunch of nodes, and you get a new node each time you get an interactive job. 

Solution 1: make a shortcut for adding a new node to the file. 

addnode() {
    NODE_NAME=$1
    CONFIG_FILE="$HOME/.ssh/config"

    # Create config entry
    CONFIG_ENTRY="Host $NODE_NAME
ProxyCommand ssh sca63@g2-login-05.coecis.cornell.edu -W $NODE_NAME:%p
User sca63"

    # Add entry to config file
    echo "$CONFIG_ENTRY" >> "$CONFIG_FILE"
    echo "SSH config entry added for $NODE_NAME"
}

 Solution 2: get all the node names and add entries for all of them in advance

First get the names of all the nodes from the cluster with sinfo -N -h -o '%N

Then add the config lines to the SSH config file:

# convert to sorted list of unique names
names = sorted(list(set(names.split())))

# get the config lines for each name
def lines(name):
    s = """Host {name}
  ProxyCommand ssh sca63@g2-login-05.coecis.cornell.edu -W {name}:%p
  User sca63

"""
    return s.format(name=name)

s = ''.join([lines(name) for name in names])

with open('/Users/simon/.ssh/config', 'a') as f:
    f.write(s)