From 6bb41872d80c9fad728e69401015b7cee549f1fc Mon Sep 17 00:00:00 2001
From: Edoardo Paone <edoardo.paone@cscs.ch>
Date: Mon, 2 Sep 2024 21:34:23 +0200
Subject: [PATCH] Add arg BASE_IMAGE

---
 ci/base.yml               | 8 +++++++-
 ci/docker/base.Dockerfile | 3 ++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/ci/base.yml b/ci/base.yml
index db16d13e92..79de71ff50 100644
--- a/ci/base.yml
+++ b/ci/base.yml
@@ -26,16 +26,22 @@ variables:
     DOCKERFILE: ci/docker/base.Dockerfile
     # change to 'always' if you want to rebuild, even if target tag exists already (if-not-exists is the default, i.e. we could also skip the variable)
     CSCS_REBUILD_POLICY: if-not-exists
-    DOCKER_BUILD_ARGS: '["ARCH=$ARCH", "HPC_SDK_VERSION=$HPC_SDK_VERSION", "HPC_SDK_NAME=$HPC_SDK_NAME", "CUPY_PACKAGE=$CUPY_PACKAGE", "PYVERSION=$PYVERSION", "CI_PROJECT_DIR=$CI_PROJECT_DIR"]'
+    DOCKER_BUILD_ARGS: '["ARCH=$ARCH", "BASE_IMAGE=$BASE_IMAGE", "HPC_SDK_VERSION=$HPC_SDK_VERSION", "HPC_SDK_NAME=$HPC_SDK_NAME", "CUPY_PACKAGE=$CUPY_PACKAGE", "PYVERSION=$PYVERSION", "CI_PROJECT_DIR=$CI_PROJECT_DIR"]'
 build_baseimage_x86_64:
   extends: [.container-builder-cscs-zen2, .build_baseimage]
   variables:
+    # x86_64 test target is Daint-gpu through Sarus:
+    # the base image does not need to provide the cuda runtime
+    BASE_IMAGE: "ubuntu:20.04"
     HPC_SDK_VERSION: 22.11
     HPC_SDK_NAME: "nvhpc_2022_2211_Linux_${ARCH}_cuda_11.8"
     CUPY_PACKAGE: cupy-cuda11x
 build_baseimage_aarch64:
   extends: [.container-builder-cscs-gh200, .build_baseimage]
   variables:
+    # aarm64 test target is Todi through Container Engine:
+    # the base image should provide the cuda runtime, therefore we use the cuda base image
+    BASE_IMAGE: "docker.io/nvidia/cuda:12.4.1-base-ubuntu20.04"
     HPC_SDK_VERSION: 24.5
     HPC_SDK_NAME: "nvhpc_2024_245_Linux_${ARCH}_cuda_12.4"
     CUPY_PACKAGE: cupy-cuda12x
diff --git a/ci/docker/base.Dockerfile b/ci/docker/base.Dockerfile
index 9f6d438183..8d2fda5fcc 100644
--- a/ci/docker/base.Dockerfile
+++ b/ci/docker/base.Dockerfile
@@ -1,4 +1,5 @@
-FROM ubuntu:20.04
+ARG BASE_IMAGE=ubuntu:20.04
+FROM ${BASE_IMAGE}
 
 ENV LANG C.UTF-8
 ENV LC_ALL C.UTF-8