秋山 翔SHO AKIYAMA / DIRECTOR
An illustrated late-night workspace overlooking a remote server room
Production notebook

Local VSCode Development with Cloud GPUs using Vertex AI Notebooks + VSCode SSH plugin

October 16, 2023

A hands-on route from a familiar local editor to GPU-backed notebooks running in the cloud.

EngineeringGPUVertex AIVS Code

Original article on Medium (opens in a new tab)

Credit: Joanna Kosinska (Unsplash)

There are times when you wish you had your own GPU setup on your local environment to easily develop your deep learning models or test inferences with speed.

This tutorial shows how to use SSH on VSCode connected to a vertex AI notebook to simulate a local VSCode environment while being able to use a cloud instance’s GPU.

Note that this concept works with any GPU VM instance from any cloud provider, but for the purpose of this tutorial and since the ML environment (CUDA, conda, etc.) is out-of-the-box in GCP Vertex AI notebooks, we will be using it for this example.

Prerequisites

  • GCP account with SSH permissions to compute engine instances on GCP iAM ( Editor role )
  • Authenticated gcloudwith gcloud auth login
  • CUDA-enabled Vertex AI notebook instance
  • Remote Development extension pack on VSCode

Steps

1. Get the external IP address of the notebook in Compute Engine
ℹ️ NOTE: You may want to reserve this as a Static IP Address to prevent it from changing when restarting the instance

2. Add an entry to your~/.ssh/config with the name you’d like to use and the external IP address from step 1:

Host <name you would like to use, ex. - research-notebook>
   HostName <the external IP address from step 1>
   UseKeychain yes
   AddKeysToAgent yes
   IdentityFile ~/.ssh/google_compute_engine
   User jupyter

ℹ️ NOTE: The IdentityFile here is the default SSH identity file for gcloud . If you use or have set up a different one, you will need to change the value for IdentityFile to point to your gcloud identity file.

3. On VSCode, open the Command Palette (default Ctrl/⌘ + Shift + P) and search for Remote SSH: Connect to Host

4. Select the hostname you would like to connect to from the list.
ℹ️ NOTE: The entry you added in ~/.ssh/config is automatically shown in the list , ex. - remote-notebook

5. In the notebook instance’s remote SSH window, open the Command Palette (default Ctrl/⌘ + Shift + P) and search for Remote: Install Remote Development Extensions

6. Install the Python and Jupyter extensions

7. You should now be able to use the notebook instance with GPU, as well as run terminal commands, as if it were a directory opened locally on VSCode!

Ending Notes

This is very useful if you want to leverage the power of GPUs without too much setup, since the workflow is similar to a local environment but being able to use GPUs present on the notebook instance.