Develope Vertex AI notebooks in VS Code using GitHub codespaces and remote tunnels

Michael Homme
4 min readApr 14, 2024

--

There are several reasons why you might want to use VS Code to develop your Jupyter Lab notebooks in Vertex AI, such as using your favorite VS Code extensions that aren’t available in Jupyter Lab. Github Copilot, data wrangler are two such extensions that come to mind.

While setting up a remote development environment from your local VS Code to the running Vertex AI is feasible, there are several downsides to it, such as, managing SSH keys and accessing remote running Conda environment with the necessary permissions…

In contrast, developing with remote tunnels gives the ability to connect to a remote machine, via a secure tunnel without the requirement of SSH.

Vertex AI Workbench instances are development environments for the entire data science workflow. Vertex AI Workbench instances combine the best of both worlds of managed instances and user-managed instances (the latter two are being deprecated).

Here are some of the benefits of using Vertex AI Workbench instances:

  • Ease of use: Vertex AI Workbench instances are managed by Google, so you don’t have to worry about setting up and maintaining the underlying infrastructure.
  • Flexibility: Vertex AI Workbench instances allow you to customize your environment with your own packages, libraries and hardware.
  • Workflow-oriented integrations: Vertex AI Workbench instances include a number of workflow-oriented integrations, such as support for JupyterLab, Git, direct data access to BigQuery and Cloud Storage, notebook scheduling and much more.

If you’re looking for a platform that makes it easy to develop and train machine learning models, then Vertex AI Workbench instances are a great option.

Start by creating a Vertex AI Workbench instance.

Choose your settings for:

  • Instance name and location (select Instance in the dropdown under Workbench type)
  • Environment
  • Machine type
  • Disks
  • Networking
  • IAM and security (select “Root access to the instance” to be able to sudo)
  • System heath

Then create the instance.

When the instance is created, open Jupyter Lab, and install VS Code. Alternatively, Code Cli could be installed. However, for the purpose of this article, we will install VS Code through snap repository manager.

Once inside the running instance, select a terminal and run the following commands:

sudo -i
sudo apt-get update
sudo apt-get install snapd
sudo snap install code - -classic

Now, the environment is ready to create a tunnel. In the terminal, run the command:

code tunnel

This command starts a VS Code Server on the running instance and creates a tunnel to it.

You will be prompted to log into your Github account to authenticate you to the tunneling service.

Open the link (marked in blue in the image above) provided in the terminal and copy and paste the code (marked in yellow) to login and grant access to the server:

Back in the instance, when the connection is authenticated, a link will be provided (marked in blue).

Open the link, and select the notebook from VS Code in the browser.

Select your notebooks and select the kernel.

You will then be prompted to install the necessary extensions, such as the Jupyter notebook extension.

Select your running environment by selecting e.g. “Python environments..”, then selecting Conda env.

If successful, the running environment will be shown in the lower left corner of VS Code..

Any additional extensions can now easily be added.

In This article we explored the advantages of using VS Code with remote tunnels for developing Jupyter Lab notebooks in Vertex AI.

The article then provides a step-by-step guide on how to set up a Vertex AI Workbench instance and install VS Code on the instance. It also demonstrates how to create a tunnel using the VS Code terminal and how to access the notebooks in VS Code from the browser.

--

--