Creating a fast.ai learning environment with VS Code + Docker on WSL2

2022-11-16

From the very beginning of the fast.ai book, the authors insists that

To do nearly everything in this book, you’ll need access to a computer with an NVIDIA GPU (unfortunately other brands of GPU are not fully supported by the main deep learning libraries). However, we don’t recommend you buy one; in fact, even if you already have one, we don’t suggest you use it just yet! Setting up a computer takes time and energy, and you want all your energy to focus on deep learning right now.

A: My two cents: heed this advice! If you like computers you will be tempted to set up your own box. Beware! It is feasible but surprisingly involved and distracting.

Fair enough. I heeded their sensible advice and used both Google Colab and Paperspace to work through the first two chapters. Like they said, it was a snap to start training with a GPU and for ease-of-use they are both 10/10.

But as an environment, I felt a bit of friction working in it. They were sometimes a bit unresponsive, sometimes I’d come back and there wouldn’t be a GPU available, it didn’t have the bells and whistles of VSCode that I’d gotten accustomed to like a built-in terminal. First-world problems I know, and it’s not a big deal at all. But I figured that if I was going to spend hours on this course, and I had the privilege of a local GPU, it might be a good investment of time to set up a comfortable local working environment.

This post will be a retrospective on how (after more time than I expected, as they warned me!) I obtained the following setup on my laptop:

A Docker container with GPU support on WSL2 with fastai installed
VS Code connected to said container to use it as a full-featured development environment.

I really enjoy working with this setup, and (despite what the authors said) I’m happy I took the time to work this out.

Remark: These steps are for a local setup, but outside of the VS Code integration I see no reason why one couldn’t do the same on a remote Linux instance. In that case you’d probably use something like VS Code’s integration with ssh.

Why Docker?

This isn’t meant to be, and isn’t going to be, a particularly friendly introduction to Docker. Maybe later! But in the meantime there are many friendly introductions to this magical piece of tech. This is a sliiiightly old but still very good introduction to Docker. My personal introduction to Docker was Nigel Poulton’s book Docker Deep Dive, though nowadays I google things and look at the docs more often.

Anyway, why Docker for machine learning in particular? There are many reasons, but as regards setting up a local environment it’s for the same reasons that everyone recommends using Anaconda: to avoid dependency nightmares. It lets you fix once and for all a working combination of drivers and software versions (e.g. CUDA and PyTorch), without needing to worry the software on your computer as a whole. Want to update? You can simply scrap the container and build a new one: no hassles with dependency conflicts or (god forbid) downgrading drivers.

It’s also important that, if you use VS Code, there is virtually zero usability compromise to developing in a Docker container vs. developing locally. Via VS Code’s incredible integration with Docker containers, I literally cannot tell the difference. Other IDEs might have similar features, but I’m not familiar with any of them.

Preliminaries: GPU support on WSL2

(Skip this if you’re following along on a Linux box.)

Follow the instructions on https://docs.nvidia.com/cuda/wsl-user-guide/index.html to get GPU acceleration on WSL2. It’s very thorough and step-by-step. The download may take a while. The distro you choose to install doesn’t really matter since we’ll be installing everything in a container anyway.
Download Docker Desktop for your Windows machine. By the magic of WSL2, the Docker Desktop software that you install like a piece of Windows software enables Docker on your WSL2 distribution.
Try running docker run --rm --gpus all nvidia/cuda:11.7.1-runtime-ubuntu20.04 nvidia-smi. Note that the latest tag is now deprecated, so you have to specify a tag. If the output looks good, you’re ready to go!

Creating the container image

I took some inspiration from fastai’s docker-container build files to create the Dockerfile to build our container. The commands are relatively self-explanatory - details are in the documentation Our base environment will be Ubuntu 20.04 LTS. The full build file looks like this, with the installation separated into several scripts:

ARG IMAGE=nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04
FROM $IMAGE

ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
ENV HOME=/root
ENV SHELL /bin/bash
WORKDIR /root

COPY ./ubuntu_script.sh ./
COPY ./conda_script.sh ./
COPY ./pytorch.sh ./
COPY ./mamba_script.sh ./

RUN bash ubuntu_script.sh
RUN bash conda_script.sh
RUN bash pytorch.sh

ENV PATH=/root/miniconda3/bin:${PATH}
SHELL ["/root/miniconda3/bin/conda", "run", "-n", "base", "/bin/bash", "-c"]
RUN bash mamba_script.sh

CMD [ "/bin/bash" ]

Keep this file as well as the scripts below in their own dedicated folder. Here are the scripts, in order:

ubuntu_script.sh

apt update
apt upgrade -y
apt install neovim wget -y

This updates the base distribution and installs a couple of useful programs. If you’d prefer to use another text editor like atom or emacs, you can change the last line accordingly.

conda_script.sh

#!/usr/bin/env bash
set -e

cd

case "$OSTYPE" in
  darwin*)  DOWNLOAD=https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh ;;
  linux*)   DOWNLOAD=https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh  ;;
  *)        echo "unknown: $OSTYPE" ;;
esac

case "$SHELL" in
  /bin/zsh*)   SHELL_NAME=zsh  ;;
  /bin/bash*)  SHELL_NAME=bash ;;
  *)        echo "unknown: $SHELL" ;;
esac

cat << EOF > .condarc
channels:
  - fastai
  - pytorch
  - defaults
channel_priority: strict
EOF

wget --quiet $DOWNLOAD
bash Miniconda3-latest*.sh -b
~/miniconda3/bin/conda init $SHELL_NAME
rm Miniconda3-latest*.sh

perl -n  -e 'print if     />>> conda/../<<< conda/' .bashrc > .condainit
perl -ni -e 'print unless />>> conda/../<<< conda/' .bashrc
echo source ~/.condainit >> .bashrc
echo conda activate >> .condainit
. .condainit

conda install -y python
conda install -y -c conda-forge mamba

conda clean --all -y

This is stolen directly from here. It installs Anaconda and also the mamba package manager.

pytorch.sh

conda install -y pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

This line installs pytorch compiled with CUDA support. The command will be different based on your compute platform - the command for the right installation for you can be copy and pasted from the handy Installation section on the PyTorch website.

Remark. The fastai package comes bundled with PyTorch, but in my experience, it’s important to install PyTorch before fastai. I struggled for a while getting my GPU recognized when training a model even though my container recognized the GPU. My problem in the end was that fastai installed a version of PyTorch compiled without CUDA support. If you try to install PyTorch afterwards, it won’t recognize your new version.

mamba_script.sh

mamba install -y -c fastchan fastai fastbook ipykernel

This installs fastai and other useful packages from the fastchan channel. The lines

SHELL ["/root/miniconda3/bin/conda", "run", "-n", "base", "/bin/bash", "-c"]
RUN bash mamba_script.sh

in the Dockerfile tells docker to switch the shell to the conda shell installed by conda_script.sh.

Installing and Running the Container

Now run docker build . --build-arg IMAGE=nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04 --tag CUSTOM_REPO/CUSTOM_TAG
from inside the directory containing the Dockerfile and the scripts. Change CUSTOM_REPO/CUSTOM_TAG to whatever you want (I chose nakaday/fastbook_base). If there is a newer LTS or a different distribution that you want to use, replace the IMAGE argument with the appropriate image.

Before running the container, I strongly recommend creating a local directory in which to keep all of your fastai files. Anything inside the container which isn’t mounted from an external directory will vanish if the container does! By mounting it as a volume, you can give the container access to those files as though it were a local directory while keeping it separate from the container itself. For my part, I simply cloned the fastbook repository.

If the image looks like it built properly, it’s time to reap the rewards by instantiating a container! Run docker run --gpus all -it -v ~/Documents/Projects/fastbook:/root/fastbook CUSTOM_REPO/CUSTOM_TAG /bin/bash (change ~/Documents/Projects/fastbook to whatever directory you chose to keep your persistent files. Make sure to use an absolute path). You should end up on the command line of your container!

To check that PyTorch is installed with CUDA support, run the following

import torch
torch.cuda.is_available()

Imgur Congratulations! You’ve got a working Docker container with PyTorch installed!

Epilogue: Connecting to VS Code

First install the Remote Development extension inside VS Code. Then, while your container is running, press Ctrl+Shift+p -> “Attach to running container”.

That’s it.

Try opening up a Jupyter notebook and running some models. If everything is installed correctly, it should be indistinguishable from working on a local machine.