Using Google Colab with a RunPod A100 Instance

By Saad | Twitter | Discord


If you've ever used Google Colab and said "I wish I could run this notebook with an A100 GPU", this is the guide for you.

RunPod is an amazing service that gives you access to a A100 GPUs for less than $1.50/hr. Unfortunately it's not trivial to use external runtime environments with Colab, but it is possible. We will use SSH Tunneling to trick Colab into thinking our RunLab instance is a local runtime environment.

The A100 comes with 80gb VRAM and is 3-10x faster than the GPUs Colab gives you by default.

Comparisons

For Windows Users: Replace the word 'terminal' with 'PowerShell 7' in this guide. All commands should work the same

Unfortunately the google.colab python package cannot be used in local environments, so mounting gDrive directories will not be possible. You will need to replace any usage of google.colab in your notebook with other methods.

1. Create an SSH key

If you already have an SSH key, you can skip this step

In your terminal, use ssh-keygen to create a public/private key pair.

$ ssh-keygen -t ed25519 -C "[email protected]"

Leave the passphrase option blank

This will create a id_ed25519.pub file inside the ~/.ssh directory.

This file may also be called id_rsa.pub if you generated it a long time ago

2. Create a RunPod A100 Instance

Launch an A100 instance on RunLab. Use the runpod/stack template, and uncheck the "Launch Jupyter Notebook" option. Your configuration should look like the following:

config

After the instance boots up, you can see it in the My Pods section. Click the dropdown on your pod and select Edit, then Edit Job.

Expand the Environment Variables section and replace the PUBLIC_KEY's value with your own public key. To get your public key, run the following:

$ cat ~/.ssh/id_ed25519.pub

public key

Your pod should restart after this

3. Create the SSH tunnel

Once your pod is restarted, click the Connect button on your pod and copy the command under SSH over exposed TCP to your terminal. Add -L 8888:localhost:8888 to this command to set up the SSH tunnel. The final SSH command should look something like this:

$ ssh root@<pod-ip-address> -p 17079 -i ~/.ssh/id_ed25519 -L 8888:localhost:8888

This will setup the SSH tunnel. Essentially any traffic going through your port 8888 will now go through RunPod's port 8888.

You will also now be logged in to your RunLab pod.

4. Setup Jupyter Server and Launch

Once logged in to the pod, we need to set up a few things that Colab requires, and then we can launch Jupyter.

Install the jupyter_http_over_ws extension:

$ pip install jupyter_http_over_ws

Enable the extension in Jupyter:

$ jupyter serverextension enable --py jupyter_http_over_ws

Finally, launch Jupyter notebook, with the following options:

$ jupyter notebook --NotebookApp.allow_origin='https://colab.research.google.com' --port=8888 --NotebookApp.port_retries=0 --allow-root

Once Jupyter is launched, it will give you a few URLs. Copy the localhost one. it should look something like this:
http://localhost:8888/?token=<jupyter-token>
This is the URL that Colab will use to connect to the Jupyter instance running on your pod.

4. Connect Colab to your Pod's Jupyter Instance

Open any Colab notebook and click the dropdown arrow beside Connect. Select Connect to Local Runtime. Paste the URL from the previous step into the Backend URL textbox. And click Connect. Your Colab should now be Connected to RunPod!

Run the !nvidia-smi command to make sure that the GPU is indeed the A100

Colab

As long as Jupyter is running on your pod, you can connect any Colab notebook to this runtime.

Once your done, make sure to shut down your pod and delete all data, or you will be charged extra!

5. Happy Diffusing!

Edit Report
Pub: 12 Sep 2022 02:47 UTC
Edit: 13 Sep 2022 07:42 UTC
Views: 535