improve build and init scripts

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/amd.sh

@@ -8,6 +8,13 @@ if [[ -z $ROCM_VERSION ]]; then
     exit 1
 fi
 
+export ROCM_VERSION="$ROCM_VERSION"
+export ROCM_LEVEL="$ROCM_LEVEL"
+export PATH=/opt/rocm/bin:$PATH
+printf "export ROCM_VERSION=\"%s\"\n" "${ROCM_VERSION}" >> /etc/bash.bashrc
+printf "export ROCM_LEVEL=\"%s\"\n" "${ROCM_LEVEL}" >> /etc/bash.bashrc
+printf "export PATH=\"%s\"\n" "${PATH}" >> /etc/bash.bashrc
+
 curl -Ss https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null
 
 echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} jammy main" \
@@ -39,4 +46,4 @@ elif [[ "${ROCM_LEVEL}" == "devel" ]]; then
 else
     printf "No valid ROCM_LEVEL specified\n" >&2
     exit 1
-fi
+fi

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh

@@ -3,6 +3,11 @@
 # Must exit and fail to build if any command fails
 set -eo pipefail
 
+export MAMBA_CREATE="micromamba create --always-softlink -y -c conda-forge"
+export MAMBA_INSTALL="micromamba install --always-softlink -y -c conda-forge"
+printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /etc/bash.bashrc
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /etc/bash.bashrc
+
 apt-get update
 apt-get upgrade -y --no-install-recommends
 

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh

@@ -3,19 +3,17 @@
 # Must exit and fail to build if any command fails
 set -eo pipefail
 
-source /opt/ai-dock/bin/build/layer0/config.sh
-
-/opt/ai-dock/bin/build/layer0/common.sh
+source /opt/ai-dock/bin/build/layer0/common.sh
 
 if [[ "$XPU_TARGET" == "NVIDIA_GPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/nvidia.sh
+    source /opt/ai-dock/bin/build/layer0/nvidia.sh
 elif [[ "$XPU_TARGET" == "AMD_GPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/amd.sh
+    source /opt/ai-dock/bin/build/layer0/amd.sh
 elif [[ "$XPU_TARGET" == "CPU" ]]; then
-    /opt/ai-dock/bin/build/layer0/cpu.sh
+    source /opt/ai-dock/bin/build/layer0/cpu.sh
 else
     printf "No valid XPU_TARGET specified\n" >&2
     exit 1
 fi
 
-/opt/ai-dock/bin/build/layer0/clean.sh
+source /opt/ai-dock/bin/build/layer0/clean.sh

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh

@@ -3,4 +3,15 @@
 # Must exit and fail to build if any command fails
 set -eo pipefail
 
-exit 0
+export CUDA_VERSION=${CUDA_VERSION}
+export CUDNN_VERSION=${CUDNN_VERSION}
+export CUDA_LEVEL=${CUDA_LEVEL}
+export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
+export MAMBA_CREATE="micromamba create --always-softlink -y -c nvidia -c conda-forge"
+export MAMBA_INSTALL="micromamba install --always-softlink -y -c nvidia -c conda-forge"
+printf "export CUDA_VERSION=\"%s\"\n" "${CUDA_VERSION}" >> /etc/bash.bashrc
+printf "export CUDNN_VERSION=\"%s\"\n" "${CUDNN_VERSION}" >> /etc/bash.bashrc
+printf "export CUDA_LEVEL=\"%s\"\n" "${CUDA_LEVEL}" >> /etc/bash.bashrc
+printf "export LD_LIBRARY_PATH=\"%s\"\n" "${LD_LIBRARY_PATH}" >> /etc/bash.bashrc
+printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /etc/bash.bashrc
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /etc/bash.bashrc

build/COPY_ROOT/opt/ai-dock/bin/init.sh

@@ -84,11 +84,6 @@ function init_set_envs() {
     while IFS='=' read -r -d '' key val; do
         export "${key}"="$(init_strip_quotes "${val//___/' '}")"
     done < <(env -0)
-
-    # TODO: branch init.sh into common,nvidia,amd,cpu
-    if [[ $XPU_TARGET == "AMD_GPU" ]]; then
-        export PATH=$PATH:/opt/rocm/bin
-    fi
 }
 
 function init_set_ssh_keys() {
@@ -396,16 +391,15 @@ function init_source_preflight_script() {
 
 function init_write_environment() {
     # Ensure all variables available for interactive sessions
-    env > /etc/environment
     while IFS='=' read -r -d '' key val; do
-        printf "export %s=\"%s\"\n" "$key" "$val" >> /root/.bashrc
+        printf "export %s=\"%s\"\n" "$key" "$val" >> /etc/bash.bashrc
     done < <(env -0)
-    
+
     if [[ -n $MAMBA_DEFAULT_ENV ]]; then
-      printf "micromamba activate %s\n" $MAMBA_DEFAULT_ENV >> /root/.bashrc
+      printf "micromamba activate %s\n" $MAMBA_DEFAULT_ENV >> /etc/bash.bashrc
     fi
 
-    printf "cd %s\n" "$WORKSPACE" >> /root/.bashrc
+    printf "cd %s\n" "$WORKSPACE" >> /etc/bash.bashrc
 }
 
 function init_get_provisioning_script() {
@@ -460,6 +454,9 @@ function init_debug_print() {
         printf "authorized_keys...\n\n"
         cat /root/.ssh/authorized_keys
         printf "\n--------------------------------------------\n"
+        printf "/etc/bash.bashrc...\n\n"
+        cat /etc/bash.bashrc
+        printf "\n--------------------------------------------\n"
         printf ".bashrc...\n\n"
         cat /root/.bashrc
         printf "\n---------- END DEBUG INFO---------- \n\n\n"

build/Dockerfile

@@ -18,33 +18,23 @@ SHELL ["/bin/bash", "-c"]
 ENV SHELL="/bin/bash"
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PATH=/opt/ai-dock/bin:/opt/micromamba/bin:/opt/caddy/bin:$PATH
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
 ENV OPT_SYNC=
 ENV PYTHONUNBUFFERED=true
 ENV APT_INSTALL="apt-get install -y --no-install-recommends"
 ENV PIP_INSTALL="pip install --no-cache-dir"
-ENV MAMBA_CHANNELS="-c nvidia -c conda-forge"
-ENV MAMBA_CREATE="micromamba create --always-softlink -y ${MAMBA_CHANNELS}"
-ENV MAMBA_INSTALL="micromamba install --always-softlink -y ${MAMBA_CHANNELS}"
 ENV MAMBA_ROOT_PREFIX=/opt/micromamba
 ENV MAMBA_EXE="/opt/micromamba/bin/micromamba"
 ENV RCLONE_CONFIG="/etc/rclone/rclone.conf"
+ENV IMAGE_SLUG="base-image"
 # Copy early so we can use scripts in the build - Changes to these files will invalidate the cache and cause a rebuild.
 COPY ./COPY_ROOT/ /
 
-# Env these args - We need them to ease software installation in extended images
 ARG CUDA_VERSION
-ENV CUDA_VERSION=${CUDA_VERSION}
 ARG CUDNN_VERSION
-ENV CUDNN_VERSION=${CUDNN_VERSION}
 ARG CUDA_LEVEL
-ENV CUDA_LEVEL=${CUDA_LEVEL}
 ARG ROCM_VERSION
-ENV ROCM_VERSION=${ROCM_VERSION}
 ARG ROCM_LEVEL
-ENV ROCM_LEVEL=${ROCM_LEVEL}
 
-ENV IMAGE_SLUG="base-image"
 # Use build scripts to ensure we can build all targets from one Dockerfile in a single layer.
 # Don't put anything heavy in here - We can use multi-stage building above if necessary.
 RUN set -eo pipefail && /opt/ai-dock/bin/build/layer0/init.sh | tee /var/log/build.log