Choose the IDE you want and develop on Vertex AI Workbench (Part I)

Notebooks are the de-facto development standard tools for data science, and Google Cloud provides Vertex AI Workbench to make data scientists more productive.
Still, other development IDEs (e.g. VSCode, Pycharm, etc.) offer interesting features (debugging, autocompletion, etc.) and developers feel more productive using those tools, especially at the time of writing libraries and support code.

In this series, we are going to present how two of the most common development IDEs, VSCode and Pycharm, can be configured on your local machine to work with a remote Vertex AI Workbench.

A data scientist using locally installed IDEs to interact with remote Vertex AI Workbench instances

In this article, we start with VSCode and we briefly present how to create a Vertex AI Workbench instance, more precisely an User Managed Notebook. Then we discuss how you can connect to a notebook with VSCode.
We tried to provide a detailed step by step guide. It might seem complex at first since apparently many steps are involved, but it is actually pretty simple.

Create User Managed Notebooks

To illustrate how to use VSCode with Vertex Notebooks, we create two User Managed Notebooks. As we will see there are a few architectural differences characterizing each option.

Providing all the details about how to create a user managed notebook is beyond our objectives. Refer to the documentation to read more information about how to create a notebook. Also, consider that a user managed notebook is created in a subnet of a VPC (see also the documentation): you must ensure that you can reach your notebooks from your laptop.

We create a first notebook using a custom container image

and a second notebook using Google provided images like, for instance, Python:

Once each instance is available, we connect with the JupyterLab installed on each machine and we clone the Vertex AI demos github repository.

In the custom-image notebook you should obtain something like

While in the python-image machine you should get something like

Cloning the sample repository allows us to compare what you see when you connect with the JypyterLab and when you connect with IDEs to the very same instance.

We can now continue configuring the machine where you installed VSCode and that we will use to connect with these instances.

Install GCloud CLI

To work seamlessly with Google Cloud services we recommend installing the GCloud CLI on your working environment. This configuration applies for any IDE of your choice.

Configure Visual Studio Code

To start working with VSCode and Vertex AI Notebooks, we recommend to install the following VSCode extensions:

• Cloud Code: it integrates with Google Cloud services, like for instance Google Compute engine (GCE). At the moment, to work with GCE, you need to install the pre-release version.
• Remote Development: it allows you to open any folder in a container or on a remote machine, and take advantage of VS Code’s full feature set

Configure Your Credentials and Project

To connect with the Vertex AI Notebook machine, we need its IP (depending on your configuration you might use its private or public IP) and configure key certificates. To do so, we can use Cloud Code extension.

• You can login directly from Visual Studio code selecting the Cloud Code extension and following steps suggested by the extension.

• Once logged in, you can select your GCP project using Cloud Code

• All notebooks we have created in the previous step are now displayed in the GCE section of the Cloud Code extension.

• We use the SSH button to connect with each instance

In doing so:

1. cloud connect will configure private keys on your device
2. you can also test your connectivity from your local machine with the remote instance
3. cou can collect a few settings that will be required later. For instance, when connected, the terminal will display something like:

Where

1. admin_vigano is my user
2. python-image is the host we are connected with.

Take note of these settings.

You might experience issues trying to SSH into the remote notebook. In case it does not work, refer to the documentation to identify the issue. A good starting point is represented by the troubleshooting tool and often firewall rules are a common suspect.

Connect to a User Managed Notebook

To illustrate how to connect to a user managed notebook, we start first with the python-image machine (the one created starting from Google provided images).

• Use Ctrl + Shift + P (or VSCode Palette) and select Remote-SSH: Connect to Host

From the following list, select Configure SSH Hosts

and choose the file you want to edit (in our case we edited the first option, home/vigano/.ssh/config)

and add an entry similar to the following:

depending on your network settings, you can use the name of the instance, its private or public IP where

• Host: is the name that VSCode will display later. You can choose what name you like the most (i.e. the name of the Vertex AI Workbench instance)
• HostName: the private or public IP of your notebook (depending on your network settings)
• IdentityFile: is the private key that has been created when we connected with Cloud Code with the instances. You should find it in the .ssh folder of your user directory
• User: corresponds to your Cloud Identify user (with underscores and few changes). We recommended to take note of your user in a previous step, so you should have already at hand.
• Use Ctrl + Shift + P (or VSCode Palette) and select Remote-SSH: Connect to Host. You should see also the host configured in the previous step and you can connect

In case you are requested to confirm the fingerprint of the machine, select “Continue”

On the left bottom corner VSCode will report that you are successfully connected to the instance as follows:

• Once connected, if you open folder /home/jupyter/ you will see the same folder structure we observed when we connected with JupyterLab through the console

• To start working with Python and Jupyter you need to install Python and Jupyter extensions on the remote host.
  Notice that if you open a notebook (e.g. vertex-ai-samples/notebooks/official/automl/automl_forecasting_bqml_arima_plus_comparison.ipynb) and you select the kernel, VSCode suggests you to install Python and Jupyter extensions.

• Install Jupyter and Python extensions on the remote notebook to remotely interact with the Vertex AI Workbench as if you were working on your local machine. For instance, you can interactively execute Python code on remote environments or execute Jupyter notebooks within VSCode leveraging kernels defined on Vertex AI Workbench.
  Remember to select Install in SSH (you want these extension working on the remote instance, not on your local machine).

Connect to a Custom User Managed Notebook

To use VSCode with a Vertex AI Workbench obtained starting from a custom container (e.g. our custom-image machine) a few additional steps are required:

• Connect to the instance following the same procedure described in the previous section.
• If you check /home/jupyter directory you will find the cloned content, but you will not find libraries, environments or code configured on your container.
  In fact, your custom container is executed on a docker instance on the virtual machine:

• To be able to connect to the container with VSCode, your user must belong to the docker group of the remote machine. Add your user as follows:

sudo usermod -aG docker $USER

You need to reset (or stop / start) the Vertex AI Workbench to apply this configuration.

• Connect again with the instance with VSCode. Once connected, select from the command palette ( or Ctrl + Shift + P) Remote-Containers: Attach to Running Container.
  Ensure you are executing this command from a VSCode window connected to the remote machine!

VSCode connects to docker and return the list of containers available on the instance

• select the payload-container and VSCode will update the UI (or open a new window) showing on the bottom left corner that it is connected with the selected container.

• Install Jupyter and Python extensions on the remote container to remotely interact with the Vertex AI Workbench as if you were working on your local machine. For instance, you can interactively execute Python code on remote environments or execute Jupyter notebooks within VSCode leveraging kernels defined within the custom container.
  Remember to select Install in Container (you want these extension working on the remote container, not on your local machine).

Conclusions

In this article we discussed how to connect to Vertex AI Workbench from VSCode, a widely adopted development IDE.In particular, we have shown how to work with standard and custom images, highlighting a few differences.

In subsequent article we will explore how to connect with Vertex AI Workbench with PyCharm

Acknowledgements

I would like to thank Ivan Nardini and Saeed Aghabozorgi for their suggestions and feedbacks.