Experimental: Build default Docker images for both Ubuntu versions 20.04 and 22.04

Author: peterschmidt85

.github/workflows/docker.yml

@@ -52,6 +52,7 @@ jobs:
     strategy:
       matrix:
         flavor: ["base", "devel", "devel-efa"]
+        ubuntu_version: ["20", "22"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -75,7 +76,7 @@ jobs:
           fi
           docker buildx build \
             --platform linux/amd64 \
-            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:${{ inputs.image_version }}-${{ matrix.flavor }} \
+            --tag dstackai/${{ env.BUILD_DOCKER_REPO }}:${{ inputs.image_version }}-${{ matrix.flavor }}-ubuntu${{ matrix.ubuntu_version }}.04 \
             --build-arg FLAVOR=${{ matrix.flavor }} \
             --provenance=false \
             --push \

docker/base/Dockerfile

@@ -1,7 +1,8 @@
 # syntax = edrevo/dockerfile-plus
+ARG UBUNTU_VERSION
 
 # Build stage
-FROM nvidia/cuda:12.1.1-base-ubuntu20.04 AS builder
+FROM nvidia/cuda:12.1.1-base-ubuntu${UBUNTU_VERSION}.04 AS builder
 
 ENV NCCL_HOME=/opt/nccl
 ENV CUDA_HOME=/usr/local/cuda
@@ -10,7 +11,7 @@ ENV OPEN_MPI_PATH=/usr/lib/x86_64-linux-gnu/openmpi
 # Prerequisites
 
 RUN export DEBIAN_FRONTEND=noninteractive \
-    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}04/x86_64/3bf863cc.pub \
     && apt-get update --fix-missing \
     && apt-get upgrade -y \
     && ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime \

docker/base/Dockerfile.common

@@ -1,4 +1,6 @@
-FROM nvidia/cuda:12.1.1-base-ubuntu20.04
+ARG UBUNTU_VERSION
+
+FROM nvidia/cuda:12.1.1-base-ubuntu${UBUNTU_VERSION}.04
 
 ARG _UV_HOME="/opt/uv"
 
@@ -9,7 +11,7 @@ ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
 ENV PATH="${UV_INSTALL_DIR}:${PATH}"
 
 RUN export DEBIAN_FRONTEND=noninteractive \
-    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}04/x86_64/3bf863cc.pub \
     && apt-get update --fix-missing \
     && apt-get upgrade -y \
     && ln -fs /usr/share/zoneinfo/America/New_York /etc/localtime \

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -999,9 +999,9 @@ def _patch_base_image_for_aws_efa(
     if not image_name.startswith(f"{settings.DSTACK_BASE_IMAGE}:"):
         return image_name
 
-    if image_name.endswith("-base"):
-        return image_name[:-5] + "-devel-efa"
-    elif image_name.endswith("-devel"):
-        return image_name[:-6] + "-devel-efa"
+    if image_name.endswith(f"-base-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"):
+        return image_name[:-17] + f"-devel-efa-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
+    elif image_name.endswith(f"-devel-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"):
+        return image_name[:-18] + f"-devel-efa-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
 
     return image_name

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -58,7 +58,7 @@ def get_default_image(nvcc: bool = False) -> str:
     Args:
         nvcc: If True, returns 'devel' variant, otherwise 'base'.
     """
-    return f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-{'devel' if nvcc else 'base'}"
+    return f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-{'devel' if nvcc else 'base'}-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
 
 
 class JobConfigurator(ABC):

src/dstack/_internal/settings.py

@@ -14,6 +14,9 @@
 
 DSTACK_BASE_IMAGE = os.getenv("DSTACK_BASE_IMAGE", "dstackai/base")
 DSTACK_BASE_IMAGE_VERSION = os.getenv("DSTACK_BASE_IMAGE_VERSION", version.base_image)
+DSTACK_BASE_IMAGE_UBUNTU_VERSION = os.getenv(
+    "DSTACK_BASE_IMAGE_UBUNTU_VERSION", version.base_image_ubuntu_version
+)
 
 
 class FeatureFlags:

src/dstack/version.py

@@ -1,3 +1,4 @@
 __version__ = "0.0.0"
 __is_release__ = False
 base_image = "0.10"
+base_image_ubuntu_version = "20.04"

src/tests/_internal/server/background/tasks/test_process_running_jobs.py

@@ -931,11 +931,11 @@ def _call_patch_base_image_for_aws_efa(
         ],
     )
     def test_patch_aws_efa_instance_with_suffix(self, suffix: str, instance_type: str):
-        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}"
+        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
         result = self._call_patch_base_image_for_aws_efa(
             image_name, BackendType.AWS, instance_type
         )
-        expected = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-devel-efa"
+        expected = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-devel-efa-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
         assert result == expected
 
     @pytest.mark.parametrize("suffix", ["-base", "-devel"])
@@ -954,11 +954,11 @@ def test_patch_aws_efa_instance_with_suffix(self, suffix: str, instance_type: st
         ],
     )
     def test_patch_all_efa_instance_types(self, instance_type: str, suffix: str):
-        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}"
+        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
         result = self._call_patch_base_image_for_aws_efa(
             image_name, BackendType.AWS, instance_type
         )
-        expected = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-devel-efa"
+        expected = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}-devel-efa-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
         assert result == expected
 
     @pytest.mark.parametrize("suffix", ["-base", "-devel"])
@@ -978,7 +978,7 @@ def test_patch_all_efa_instance_types(self, instance_type: str, suffix: str):
     def test_no_patch_non_aws_backends(
         self, backend: BackendType, suffix: str, instance_type: str
     ):
-        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}"
+        image_name = f"{settings.DSTACK_BASE_IMAGE}:{settings.DSTACK_BASE_IMAGE_VERSION}{suffix}-ubuntu{settings.DSTACK_BASE_IMAGE_UBUNTU_VERSION}"
         result = self._call_patch_base_image_for_aws_efa(image_name, backend, instance_type)
         assert result == image_name