Running Custom Docker Containers on Google Colab: Reproducible GPU Environments Without the Base Image Lock-in

The problem: Colab's default runtime pins you to someone else's stack

You open a new notebook, pick "GPU" from the runtime menu, and get a prebuilt environment with a fixed CUDA version, Python release, and a handful of ML libraries. Need CUDA 12.4 but the runtime ships 12.1? Want Python 3.11 instead of 3.10? Need a specific numpy wheel compiled against a particular BLAS? You're stuck — unless you spend the first 15 minutes of every session running pip install and apt-get commands that vanish when the VM recycles.

Google Colab's custom container runtime changes that. You supply a Docker image URI; Colab pulls it, starts the container with GPU/TPU access if requested, and tunnels the Jupyter kernel back to the notebook UI. The result: a reproducible, version-pinned environment that starts in the same UI you already know.

What the image must provide

Colab expects the container to expose a Jupyter server on port 8888 listening on 0.0.0.0 with token authentication disabled (--NotebookApp.token='') because Colab handles auth via its own gRPC proxy. The image also needs:

• jupyter-server and jupyter-server-proxy installed
• A kernel spec registered under /usr/local/share/jupyter/kernels/ (e.g., ipykernel)
• Linux/amd64 architecture — ARM64-only images will fail on Colab's x86 hosts

Minimal base images like nvidia/cuda:12.4-devel-ubuntu22.04 don't include these, so you'll build a derived Dockerfile.

Worked example: a minimal CUDA 12.4 + Python 3.11 image

Below is a complete Dockerfile you can build and push to Docker Hub (or Google Artifact Registry). Replace YOUR_DOCKERHUB_USER with your actual username.

# Dockerfile
FROM nvidia/cuda:12.4-devel-ubuntu22.04

# Install Python 3.11 and Jupyter dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-venv python3.11-dev \
    && rm -rf /var/lib/apt/lists/*

# Make python3.11 the default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1

# Install pip and Jupyter stack
RUN python3 -m ensurepip --upgrade \
    && pip install --no-cache-dir \
        jupyter-server \
        jupyter-server-proxy \
        ipykernel

# Register the kernel so Colab can discover it
RUN python3 -m ipykernel install --sys-prefix --name python3 --display-name "Python 3.11 (CUDA 12.4)"

# Expose the port Colab expects
EXPOSE 8888

# Start Jupyter with token disabled and binding to all interfaces
CMD ["jupyter", "server", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--NotebookApp.token=''", "--NotebookApp.password=''", "--ServerApp.root_dir=/workspace"]

Build and push (run on your local machine or CI with Docker access):

docker build -t YOUR_DOCKERHUB_USER/colab-cuda124-py311:latest .
docker push YOUR_DOCKERHUB_USER/colab-cuda124-py311:latest

Connect from Colab:

1. Open a new notebook at colab.research.google.com.
2. Runtime → Change runtime type → Container.
3. Paste the image URI: YOUR_DOCKERHUB_USER/colab-cuda124-py311:latest.
4. Select GPU (T4, V100, etc.) if you need acceleration.
5. Click Save, then Connect.

Colab will pull the image (cold start: 2–5 minutes on free tier for a ~6 GB image), start the container with --gpus all, inject NVIDIA_VISIBLE_DEVICES and NVIDIA_DRIVER_CAPABILITIES, and attach the kernel. Run !nvidia-smi in a cell to verify GPU visibility.

GPU, secrets, persistence, and the fine print

GPU passthrough

The host driver stays on the VM; your container only needs the matching CUDA toolkit (12.4 in the example). No privileged mode, no host networking — Colab runs containers in a restricted gVisor sandbox.

Secrets

Don't bake API keys into the image. Use Colab's Secret Manager (🔑 icon in the left sidebar), add a secret named WANDB_API_KEY, then read it at runtime:

import os
print(os.environ["WANDB_API_KEY"])  # injected at container start

Persistent data

The container runs as root, so mounting Google Drive is straightforward:

from google.colab import drive
drive.mount('/content/drive')
# Write checkpoints to /content/drive/MyDrive/checkpoints/

Alternatively, mount a GCS FUSE bucket for larger datasets.

Cold-start latency and disk limits

Image size directly affects pull time. A 12 GB image can take 3–5 minutes on the free tier; Colab Pro/Pro+ offers faster throughput and larger ephemeral disk quota (100 GB free, more on paid). Images over ~20 GB risk exceeding the free tier's disk limit during layer extraction.

Idle termination

Free tier: 90 minutes of inactivity kills the runtime. Pro: 4 hours. Long-running training must checkpoint to Drive/GCS and be designed for resumption.

Trade-offs you should weigh

• No docker-compose, privileged mode, or host networking. If your workflow needs multiple coordinated containers or kernel modules, Colab isn't the platform.
• Base image updates are manual. When CUDA 12.5 drops, you rebuild and re-tag; Colab won't auto-rebuild.
• Free-tier GPU quota is shared. Custom containers don't grant dedicated GPU access; you're still subject to preemption.
• ARM64 images won't start. Ensure your build targets linux/amd64 (docker buildx build --platform linux/amd64 ...).

Try it today: a 5-minute verification checklist

1. Build and push the Dockerfile above (or use a known working image like gcr.io/deeplearning-platform-release/pytorch-gpu.2-1.py310:latest for a quick test).
2. Open Colab, switch to Container runtime, paste your image URI, enable GPU, connect.
3. Run !nvidia-smi and import sys; print(sys.version) to confirm CUDA and Python versions.
4. Mount Drive, write a file, disconnect, reconnect — verify the file persists.
5. Add a secret via the 🔑 panel, read it from os.environ.
6. (Optional) Leave the notebook idle for 95 minutes on a free account to observe the timeout.

If all steps succeed, you've got a reproducible, version-controlled Colab environment that survives runtime restarts and keeps secrets out of your images. Next time you need a different CUDA/Python combo, you change the Dockerfile — not the notebook.