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.
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:
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
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
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!