Peter Davies

Ramblings

GPU signal processing in Docker

I had a new Ubuntu 18.04 workstation with a pair of RTX 3090s and wanted to do DSP on the GPU: filtering, resampling and spectral estimates over long IQ recordings, which is slow work for a CPU. My first attempt was to install CUDA on the host and pip install everything on top. That turned into an afternoon of library conflicts, so I moved everything into containers. This is how it ended up, as of December 2020.

Only the driver goes on the host

The rule that saved me: the host needs the NVIDIA driver and the container toolkit, nothing else. The CUDA toolkit, cuDNN and the Python libraries all live in the image, so each image can have the CUDA version it was built for.

My earlier host install had mixed NVIDIA's cuda packages from their apt repository with Ubuntu's own nvidia-cuda-toolkit package. They overwrite each other's files, and apt ends up refusing to continue:

dpkg -i --force-overwrite /var/cache/apt/archives/nvidia-cuda-dev_9.1.85-3ubuntu1_amd64.deb
apt-get -f install

That gets apt working again, but the real fix is not to install both. If you do want CUDA on the host, pick one source and remove the other.

The RTX 30-series cards are Ampere (compute capability 8.6) and need driver 455 or newer and CUDA 11.1 or newer for full support:

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt-get update
sudo apt-get install nvidia-driver-455
sudo reboot
nvidia-smi

NVIDIA container toolkit

Docker 19.03 and later understand --gpus natively; they need NVIDIA's container toolkit to do the work:

distribution=$(. /etc/os-release; echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list \
    | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

Check a container can see the cards:

docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

RAPIDS with JupyterLab

RAPIDS publishes images with CUDA, CuPy, cuDF and cuSignal already built against each other, and JupyterLab starts automatically:

docker run --gpus all --rm -it \
    -p 8888:8888 -p 8787:8787 -p 8786:8786 \
    -v ~/Projects:/rapids/notebooks/host \
    rapidsai/rapidsai:cuda11.0-runtime-ubuntu18.04-py3.8

The runtime images are what you want for using the libraries. The devel images add compilers and source for building RAPIDS itself, and the nightly ones are for testing the next release; I tried rapidsai/rapidsai-core-dev-nightly:0.18-cuda11.0-devel-ubuntu18.04-py3.8 and went back to the stable runtime.

To add a package without building your own image, the RAPIDS images install anything listed in EXTRA_CONDA_PACKAGES (or EXTRA_PIP_PACKAGES) at start-up:

docker run --gpus all --rm -it -p 8888:8888 \
    -e EXTRA_CONDA_PACKAGES="-c pytorch pytorch" \
    -v ~/Projects:/rapids/notebooks/host \
    rapidsai/rapidsai:cuda11.0-runtime-ubuntu18.04-py3.8

It reinstalls every time the container starts, so once you settle on a set of packages, put them in a two-line Dockerfile FROM the RAPIDS image instead.

cuSignal

cuSignal copies the scipy.signal API but works on CuPy arrays on the GPU. A minimal example: a wanted tone at 250 kHz and an unwanted one at 3 MHz in noise, decimated by 8 and then a Welch power spectral density.

import cupy as cp
import cusignal

fs = 10e6
n = 2**24
t = cp.arange(n) / fs
noise = 0.5 * (cp.random.randn(n) + 1j * cp.random.randn(n))
x = cp.exp(2j * cp.pi * 250e3 * t) + cp.exp(2j * cp.pi * 3e6 * t) + noise

y = cusignal.resample_poly(x, 1, 8, window=("kaiser", 5.0))
f, pxx = cusignal.welch(y, fs=fs / 8, nperseg=4096, return_onesided=False)

print(float(f[cp.argmax(pxx)]))

After decimating by 8 the new sample rate is 1.25 MHz, so anything beyond 625 kHz either side of the centre would alias back into the band. The 3 MHz tone would land on 500 kHz; the low-pass filter inside resample_poly removes it first, and the printed peak is the 250 kHz tone. Swap resample_poly for plain slicing (y = x[::8]), which decimates with no filter, and the alias at 500 kHz appears next to the real tone. It's a quick sanity test that the filter and the PSD are doing what you think.

Two habits matter more than anything else for speed:

TensorFlow GPU images and CUDA mismatches

I also needed some older TensorFlow 1.x code to run on the same box, and that's where most of the pain was.

First mistake: installing a TensorFlow 1.14 wheel into an image that already had TensorFlow 2 in it. pip's resolver warns but carries on:

tensorflow 1.14.1 requires tensorboard<1.15.0,>=1.14.0, but you have tensorboard 2.4.0 which is incompatible.
tensorflow 1.14.1 requires tensorflow-estimator<1.15.0rc0,>=1.14.0rc0, but you have tensorflow-estimator 2.3.0 which is incompatible.

Don't install a different TensorFlow major on top of another one. Use an image built for the version you need, or a clean virtualenv.

Second: each TensorFlow build looks for the exact CUDA library versions it was compiled against. If they're missing, it doesn't fail; it logs a warning per library and quietly runs on the CPU:

Could not load dynamic library 'libcudart.so.10.1'; dlerror: libcudart.so.10.1: cannot open shared object file: No such file or directory
...
Skipping registering GPU devices...

So always check after starting a container:

python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

An empty list means CPU only.

Third, and specific to Ampere: the official TensorFlow 1.x builds stop at CUDA 10, which doesn't support these cards. The answer for 1.x code is NVIDIA's own build of TensorFlow 1.15 against CUDA 11, either in their NGC container:

docker run --gpus all --rm -it -v ~/Projects:/workspace/host \
    nvcr.io/nvidia/tensorflow:20.11-tf1-py3

or from pip inside an Ubuntu 20.04 image with Python 3.8, which those wheels require. They aren't on PyPI, so NVIDIA's package index goes in first:

python3 -m pip install nvidia-pyindex
python3 -m pip install nvidia-tensorflow[horovod]
python3 -m pip install gast==0.2.2

The gast pin is for TensorFlow 1.15, which emits AutoGraph warnings or fails to convert functions with newer releases. For TensorFlow 2 code, 2.4.0 is the first release built against CUDA 11 and its tensorflow/tensorflow:2.4.0-gpu image works on the 30-series.

Building TensorFlow from source

Before finding the NVIDIA builds I tried compiling TensorFlow. It uses Bazel, and each TensorFlow release needs one specific Bazel version, listed in .bazelversion in the source tree. The Bazel apt repository gives you the latest by default:

apt install curl gnupg
curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg
mv bazel.gpg /etc/apt/trusted.gpg.d/
echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" \
    | tee /etc/apt/sources.list.d/bazel.list
apt update && apt install bazel

That installed Bazel 3.7.1, which is far too new for TensorFlow 1.15. The repository also has versioned packages (apt install bazel-0.26.1, for example), or you can use Bazelisk, which reads .bazelversion and fetches the right one. A full build took long enough that the prebuilt NVIDIA packages were the better route.

Today: cuSignal has been retired and its functions moved into CuPy as cupyx.scipy.signal, so the example above works with that import and no extra package. The nvidia/cuda:11.0-base tag has been removed from Docker Hub; use a current full tag such as nvidia/cuda:12.4.1-base-ubuntu22.04, and NVIDIA's container toolkit now installs from its own repository with nvidia-ctk to configure Docker.