improve build and init scripts

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

@@ -1,13 +1,17 @@
-#!/bin/bash
-
-# Must exit and fail to build if any command fails
-set -eo pipefail
+#!/bin/false
 
 if [[ -z $ROCM_VERSION ]]; then
  printf "No valid ROCM_VERSION specified\n" >&2
  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}" >> /opt/ai-dock/etc/environment.sh
+printf "export ROCM_LEVEL=\"%s\"\n" "${ROCM_LEVEL}" >> /opt/ai-dock/etc/environment.sh
+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 +43,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/clean.sh

@@ -1,7 +1,5 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
 
 # Tidy up and keep image small
 apt-get clean -y

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

@@ -1,7 +1,9 @@
-#!/bin/bash
+#!/bin/false
 
-# 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}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh
 
 apt-get update
 apt-get upgrade -y --no-install-recommends
@@ -92,3 +94,8 @@ touch /etc/rclone/rclone.conf
 # Git config
 
 git config --global --add safe.directory "*"
+
+# Ensure correct environment for child builds
+
+printf "source /opt/ai-dock/etc/environment.sh\n" >> /etc/profile.d/02-ai-dock.sh
+printf "source /opt/ai-dock/etc/environment.sh\n" >> /etc/bash.bashrc

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

@@ -1,6 +1,3 @@
-#!/bin/bash
+#!/bin/false
 
-# Must exit and fail to build if any command fails
-set -eo pipefail
 
-exit 0

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

@@ -1,6 +1,14 @@
-#!/bin/bash
+#!/bin/false
 
-# 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}" >> /opt/ai-dock/etc/environment.sh
+printf "export CUDNN_VERSION=\"%s\"\n" "${CUDNN_VERSION}" >> /opt/ai-dock/etc/environment.sh
+printf "export CUDA_LEVEL=\"%s\"\n" "${CUDA_LEVEL}" >> /opt/ai-dock/etc/environment.sh
+printf "export LD_LIBRARY_PATH=\"%s\"\n" "${LD_LIBRARY_PATH}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /opt/ai-dock/etc/environment.sh
+printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh

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,11 +391,10 @@ 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" >> /opt/ai-dock/etc/environment.sh
  done < <(env -0)
- 
+
  if [[ -n $MAMBA_DEFAULT_ENV ]]; then
  printf "micromamba activate %s\n" $MAMBA_DEFAULT_ENV >> /root/.bashrc
  fi
@@ -460,6 +454,9 @@ function init_debug_print() {
  printf "authorized_keys...\n\n"
  cat /root/.ssh/authorized_keys
  printf "\n--------------------------------------------\n"
+ printf "/opt/ai-dock/etc/environment.sh...\n\n"
+ cat /opt/ai-dock/etc/environment.sh
+ 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