Upgrade Base image to 12.2 (#704)

* Upgrade Base image to 12.2

* Add platform back

* Prep for building images in integration tests

* Drop training image and bump rc version

* Upgrade versions for release

.github/workflows/integration-tests.yml

@@ -33,7 +33,7 @@ jobs:
       matrix:
         python_version: ["3.8", "3.9", "3.10", "3.11"]
         use_gpu: ["y", "n"]
-        job_type: ["server", "training"]
+        job_type: ["server"]
     steps:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v1

.github/workflows/main.yml

@@ -36,7 +36,7 @@ jobs:
       matrix:
         python_version: ["3.8", "3.9", "3.10", "3.11"]
         use_gpu: ["y", "n"]
-        job_type: ["server", "training"]
+        job_type: ["server"]
     steps:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v1

bin/generate_base_images.py

@@ -104,13 +104,18 @@ def _build(
             "docker",
             "buildx",
             "build",
-            "--platform=linux/amd64",
+            "--platform=linux/arm64,linux/amd64",
             ".",
             "-t",
             image_with_tag,
         ]
         if push:
             cmd.append("--push")
+
+        # Needed to support multi-arch build.
+        subprocess.run(
+            ["docker", "buildx", "create", "--use"], cwd=build_ctx_path, check=True
+        )
         subprocess.run(cmd, cwd=build_ctx_path, check=True)
 
 

docker/base_images/base_image.Dockerfile.jinja

@@ -1,19 +1,19 @@
 {% if use_gpu %}
-FROM nvidia/cuda:11.2.1-base-ubuntu20.04
-ENV CUDNN_VERSION=8.1.0.77
-ENV CUDA=11.2
+FROM nvidia/cuda:12.2.2-base-ubuntu20.04
+ENV CUDNN_VERSION=8.9.5.29
+ENV CUDA=12.2
 ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
 
 RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \
     apt-get update && apt-get install -y --no-install-recommends \
         ca-certificates \
-        cuda-command-line-tools-11-2 \
-        libcublas-11-2 \
-        libcublas-dev-11-2 \
-        libcufft-11-2 \
-        libcurand-11-2 \
-        libcusolver-11-2 \
-        libcusparse-11-2 \
+        cuda-command-line-tools-12-2 \
+        libcublas-12-2 \
+        libcublas-dev-12-2 \
+        libcufft-12-2 \
+        libcurand-12-2 \
+        libcusolver-12-2 \
+        libcusparse-12-2 \
         libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA} \
         libgomp1 \
         && \

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.7.14"
+version = "0.7.15rc0"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/contexts/image_builder/util.py

@@ -11,7 +11,7 @@
 # [IMPORTANT] Make sure all images for this version are published to dockerhub
 # before change to this value lands. This value is used to look for base images
 # when building docker image for a truss.
-TRUSS_BASE_IMAGE_VERSION_TAG = "v0.4.9"
+TRUSS_BASE_IMAGE_VERSION_TAG = "v0.7.15"
 
 
 def file_is_empty(path: Path, ignore_hash_style_comments: bool = True) -> bool: