How to connect from VSCode to Vertex Notebooks directly

I have seen many tutorials on how to connect from VScode using Jupyter extension to Vertex Notebooks. The mainstream solution is mainly opening an SSH tunnel, which works, but not available in all scenarios: No Public IP, Organization restricts SSH access via Firewall rules or OS Login. Please take a look at this alternative solution using Jupyter Gateway Client.

Note: This is not supported by GCP, if you encounter any issues ping me here or via StackOverflow.

1. Create a Workbench Instance or User Managed Notebook. Pass the following metadata to allow cross-site requests.

disable-check-xsrf=True

2. Install JupyterLab in your local workstation.

pip install jupyter-lab

3. Create a new Jupyter config file `jupyter_notebook_config.py`

Modify PROXY_URL with your Proxy URL. This can be obtained by clicking Open JupyterLab in Google Cloud Console.

"""VScode support for Workbench Instances."""
import datetime
import subprocess
import hashlib
import json

PROXY_URL = "https://<Your endpoint>-dot-us-central1.notebooks.googleusercontent.com"

c.NotebookApp.open_browser = False
c.ServerApp.token = ""
c.ServerApp.password = ""
c.ServerApp.port = 8080
c.ServerApp.allow_remote_access = True
c.ServerApp.disable_check_xsrf = True

def get_gcloud_token():
  """Helper method to get an OAuth token from gcloud."""
  p = subprocess.run(
      ["gcloud", "auth", "print-access-token"],
      capture_output=True,
      check=True,
      encoding="UTF-8",
  )
  bearer_token = p.stdout.strip()
  return bearer_token


c.GatewayClient.auth_scheme = "Bearer"
c.GatewayClient.auth_token = get_gcloud_token()
c.GatewayClient.headers = json.dumps({"Origin": PROXY_URL})
c.GatewayClient.validate_cert = False
c.GatewayClient.url = PROXY_URL
c.Application.log_level = 0

4. Start JupyterLab locally

jupyter lab --port-retries=0 --ip 0.0.0.0 --allow-root --config=jupyter_notebook_config.py --debug

5. Start VSCode with JupyterLab extension pre-installed and connnect to the local Jupyter server. (`http://127.0.0.1:8080`). Local Jupyter server via Gateway client will connect to remote server and you should be able to access remote Vertex instance.

6. Verify you are connected to remote instance.

The next step is to add support to refresh token via `JUPYTER_GATEWAY_TOKEN_RENEWER_CLASS` Reference: https://github.com/jupyter-server/jupyter_server/blob/main/jupyter_server/gateway/gateway_client.py#L51

For Colab Enterprise, you can use same configuration. (Check: https://github.com/jupyter-server/jupyter_server/pull/1440)

Video here: https://www.youtube.com/watch?v=8o0Jo6KfuAU