Rebuilding TensorFlow 2.8.4 on Ubuntu 22.04 to patch vulnerabilities
Tags: docker, tensorflow, ai, security
Published: 2024-06-02 06:55:16

Table of contents:

1. Context
2. My process
3. Summary of method & results

TLDR:

| | DeepCell +  tensorflow-2.8.4 | DeepCell + tensorflow-2.8.4-redux | Delta |
| ------ | --- | --- | --- |
| **Compressed size** | 3.2GB | 4.0 GB | **+0.8 GB (+25%)** |
| **VULNs** | 553 | 125 | **-428 (-77%)** |
| Critical | 1 | 1 | 0 (0%) |
| High | 80 | 29 | -51 (-63%) |
| Medium | 349 | 53 | -296 (-85%) |
| Low | 123 | 42 | -81 (-66%) |

Read on for the how, why, wherefore, and finally.

## Context & motivation

Previously we [switched](https://dev.to/dchaley/container-size-analysis-tensorflow-28-base-image-vs-deep-learning-4dbj) from the DeepLearning container to the base TensorFlow container.

Unfortunately the container has 553 security of vulnerabilities according to Google's scanner:

![Screenshot of the Google container vulnerability scanner](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/e7ixtmt1mvja4wzw79s5.png)

The 553 issues break down this way:
- 1 critical [[vuln](https://github.com/advisories/GHSA-gw97-ff7c-9v96)]
- 80 high
- 349 medium
- 123 low

The [official 2.8.4 container](https://hub.docker.com/layers/tensorflow/tensorflow/2.8.4-gpu/images/sha256-4351b59baf4887bcf47eb78b34267786f40460a81fef03c9b9f58e7d58f1c7b7?context=explore) was published in Nov 2022. That's 1.5 years of OS updates **at least**. I looked up the [2.8.4 source](https://github.com/tensorflow/tensorflow/tree/v2.8.4) and found that it's [using Ubuntu 20.04](https://github.com/tensorflow/tensorflow/blob/v2.8.4/tensorflow/tools/dockerfiles/partials/ubuntu/version.partial.Dockerfile) as the base OS. Of note, we're using the x86_64 architecture according to the container image layer: `ENV NVARCH=x86_64`.

So the obvious thing to do is to switch to the most recent Ubuntu version 24.04 right? Well no, that's a short party: NVIDIA doesn't have CUDA packages for 24.04 in their [repository](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/). So it's off to 22.04 – still two years more recent, and more importantly with [CUDA packages](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/).

## My process, as I did it

I wouldn't do it this way again, but this is how I did it.

### Updating the base Ubuntu image + dependencies.

First, I forked the tensorflow repository. I did a master-only clone so I needed to fetch the tag information after clone. Then, I could reset to the 2.8.4 version.

```bash
# Add upstream branch.
git remote add upstream https://github.com/tensorflow/tensorflow.git
git fetch upstream
# Reset master branch to 2.8.4
git checkout master
git reset --hard v2.8.4
git push --force
# Clean out local copy (everything after 2.8.4)
git gc
```

Then, I updated the build steps. Here's what I did, following the instructions in the [containers readme](https://github.com/dchaley/tensorflow-2.8.4-redux/blob/master/tensorflow/tools/dockerfiles/README.md).

**1\. Build the `tf-tools` build tools container:**

```bash
cd tensorflow/tools/dockerfiles
docker build -t tf-tools -f tools.Dockerfile .
```

**2\. Set up aliases:**

```bash
alias asm_dockerfiles="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/tf tf-tools python3 assembler.py "
alias asm_images="docker run --rm -v $(pwd):/tf -v /var/run/docker.sock:/var/run/docker.sock tf-tools python3 assembler.py "
   ```

**3\. Update build settings.** I started with changing the file `partials/ubuntu/version.partial.Dockerfile` to use Ubuntu 22.04.

**4\. Regenerate the dockerfiles.**

```bash
asm_dockerfiles --release dockerfiles --construct_dockerfiles
```

**5\. Rebuild the desired TF-2.8 image.** This builds a container tagged with the `2.8.4-rebuilt` version, which causes the build system to tag the GPU-accelerated container `2.8.4-rebuilt-gpu`.

```bash
asm_images --release versioned --arg _TAG_PREFIX=2.8.4-rebuilt --build_images --only_tags_matching="^2.8.4-rebuilt-gpu$"
```

**6\. Done, or need to fix.** Fix build errors & loop to step 3.

#### Dependency updates

Following this process here's what I [fixed at first](https://github.com/dchaley/tensorflow-2.8.4-redux/commit/7249644d7188d9512986781bd150a1be1af7b87f):

* Downgrade requests & urllib libraries (see [github bug](https://github.com/docker/docker-py/issues/3256))
* Update base Ubuntu to 22.04.
* Update CUDA from 11.2.1 to 11.8.0
* Parameterize CUDA patch level (to support `.0` instead of `.1`)
* Update CUDNN from `8.1.0.77-1` to `8.6.0.163-1`
* Update `libvinfer` from `7.2.2-1` to `8.5.3-1`.
  * I didn't love the major version update. But things seem fine.

At this point the container built, and I could run DeepCell. It output a segmentation image that seems plausible. 

However a new error message popped up in the logs…

```
2024-05-28 19:38:34.423093: W tensorflow/stream_executor/gpu/asm_compiler.cc:80] Couldn't get ptxas version string: INTERNAL: Couldn't invoke ptxas --version
2024-05-28 19:38:34.423903: I tensorflow/core/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
2024-05-28 19:38:34.424006: W tensorflow/stream_executor/gpu/redzone_allocator.cc:314] INTERNAL: Failed to launch ptxas
Relying on driver to perform ptx compilation. 
Modify $PATH to customize ptxas location.
This message will be only logged once.
```

Is this an error? Is it a problem to rely on the driver? I don't know, but I wanted to clear out the error.

#### Finding ptxas

I found a [GitHub issue](https://github.com/google/jax/discussions/6843) that seemed similar (missing ptxas) and saw a suggestion to install nvidia-cuda-toolkit. Alright: but that exploded the container size from 6.5 GB to 12.13 GB … unacceptable 😤    (Incidentally, this is too large for Cloud Shell to build on its limited persistent disk.)

![Screenshot of the Docker desktop app showing two versions of the image](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/rbl7tp58zqrmvyxgwccj.png)

At this point I struggled for a couple hours. The `nvidia-cuda-toolkit` [package info](https://packages.ubuntu.com/jammy/nvidia-cuda-toolkit) says it uses CUDA 11.5. But the [prebuilt containers](https://hub.docker.com/r/nvidia/cuda/tags?page=1&page_size=100&name=ubuntu22.04&ordering=) had 11.7 and 11.8 not 11.5 (I'd previously selected 11.8). The 11.5 packages weren't available in NVIDIA's [Ubuntu 22.04 package repo](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/).

Along the way, I [switched the base container](https://github.com/dchaley/tensorflow-2.8.4-redux/commit/ad79db46b6504c06aa42b03b65c1ba8afd99d7b6) from NVIDIA's [`nvidia/cuda:11.8.0-base-ubuntu22.04`](https://hub.docker.com/layers/nvidia/cuda/11.8.0-base-ubuntu22.04/images/sha256-942f9a2455c62479908e6e6b7fc0eeff38a6daac2f67a0410594e0a04d688db0?context=explore) to [`nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04`](https://hub.docker.com/layers/nvidia/cuda/11.8.0-cudnn8-runtime-ubuntu22.04/images/sha256-8e6794c967a86219264f3ef770036fd7da4f6adcdc0d145074dcd61b2405fda4?context=explore). Rather than pick the versions myself, I figured going with an official NVIDIA container with the files I was installing anyhow made sense.

I eventually found this "[ptxas version issue](https://github.com/google/jax/discussions/10327)" linked from a [TensorFlow discussion](https://discuss.tensorflow.org/t/should-i-worry-about-this-warning-while-training-a-tensorflow-lite-object-detection-model-will-it-affect-my-training/24143) asking whether to worry about a version mismatch warning. Not the same as our message that it's missing, but close enough.

This part caught my eye:

```
You may not need to update to CUDA 11.1; cherry-picking the ptxas binary is often sufficient.
```

Interesting idea. I cherry-picked the binary by launching a container from the rebuilt image, and installing the very large `nvidia-cuda-toolkit`:

```bash
apt-get install nvidia-cuda-toolkit
Need to get 1603 MB of archives.
After this operation, 4505 MB of additional disk space will be used.
Do you want to continue? [Y/n]
```

Gulp. One very long download later, I had a ptxas binary.

```
root@33eda96a19a0:/# which ptxas
/usr/bin/ptxas
```

Now to copy it back to the host, so I can add it to the redux repo for direct insertion into the container. Back on the host:

```
docker cp 33eda96a19a0:/usr/bin/ptxas .
```

Then I [installed it](https://github.com/dchaley/tensorflow-2.8.4-redux/commit/a9ff10438e92d2fd916fee56323922c71cff1bfe) into `/usr/bin/ptxas` in the dockerfile.

Lo and behold: no more ptxas error when running DeepCell.

-----

## Summary

The container was rebuilt by:

- Forking TensorFlow 2.8.4 from source.
- Switching to the `nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04` base image.
- Cherry-picking ptxas from `nvidia-cuda-toolkit` (avoids ~4GB unnecessary files).

The container rebuild yielded these changes:

| | DeepCell +  tensorflow-2.8.4 | DeepCell + tensorflow-2.8.4-redux | Delta |
| ------ | --- | --- | --- |
| **Compressed size** | 3.2GB | 4.0 GB | **+0.8 GB (+25%)** |
| **VULNs** | 553 | 125 | **-428 (-77%)** |
| Critical | 1 | 1 | 0 (0%) |
| High | 80 | 29 | -51 (-63%) |
| Medium | 349 | 53 | -296 (-85%) |
| Low | 123 | 42 | -81 (-66%) |

It's too bad we added 25% to the container size. This may be because I [moved away](https://github.com/dchaley/tensorflow-2.8.4-redux/commit/ad79db46b6504c06aa42b03b65c1ba8afd99d7b6) from the TensorFlow container build's selective dependencies to the full runtime package.

Still, 77% reduction in VULNs (and 63% for the highs) is very good.

The [critical VULN]([[vuln](https://github.com/advisories/GHSA-gw97-ff7c-9v96)]) is in TensorFlow pre 2.11.1. It allows malicious users running custom TensorFlow python code to access memory that's not theirs in some cases. Since we're running our own Python code, and DeepCell's, we're safe as long as nobody sticks in naughty code in those layers. But, we're also stuck with 2.8.4 and can't upgrade to 2.11 so the rationalization is rationalized.

If I were to do it again, I'd skip hand-picking library versions & move straight to an official NVIDIA runtime container.

-----

## Appendix

Helpful command to get into the TF container shell to poke around for files:

```
docker run --user $(id -u):$(id -g) -it -v $(pwd):/tf tensorflow:2.8.4-rebuilt-gpu bash
```

I ran out of disk space on cloud shell a few times. Clear out the docker cache like so:

> ⚠️ Don't run these as-is if you have other containers/images you want to keep!

```bash
docker system prune

# Delete the previously build image
# (make room for new one)
docker image rm tensorflow:2.8.4-rebuilt-gpu
```

In the end, Cloud Shell (which has a limited disk) became a hassle for iterating on builds. I considered a [Cloud Workstation](https://cloud.google.com/workstations/docs/create-workstation) however there's a fixed $0.20/hr cost whether or not you have a workstation running … and I really just need a place to run Docker with disk space, so I used my local computer (a mac). The downloads weren't as fast as on cloud, but hey.

Side note: I'm super impressed with how easy it was to rebuild TF from source. Nice job y'all 🤩

Tools used in the rebuild:

* GCP Cloud Shell
* GCP Artifact Registry container scanner
* Docker (local + cloud shell)
* git & GitHub
* TensorFlow
* apt-file (to look up which package installed a file)