Fixing docker build for DeepSpeed, introducing build for conda environment

docker/deepspeed/Dockerfile

@@ -0,0 +1,95 @@
+FROM continuumio/miniconda3:24.7.1-0
+
+ARG DEEPSPEED_VERSION=0.16.2
+ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
+
+SHELL ["/bin/bash", "-l", "-c"]
+ENV SHELL=/bin/bash
+ENV HOST=0.0.0.0
+ENV DEBIAN_FRONTEND=noninteractive
+ENV CUDA_DOCKER_ARCH=all
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV CONDA_AUTO_UPDATE_CONDA="false"
+
+##############################################################################
+# Directories:
+##############################################################################
+ENV STAGE_DIR=/tmp
+RUN mkdir -p ${STAGE_DIR}
+
+##############################################################################
+# Installation/Basic Utilities
+##############################################################################
+RUN <<EOR
+  apt-get update
+  apt-get upgrade -y
+  apt-get install --no-install-recommends -y \
+      espeak-ng \
+      curl \
+      wget \
+      jq \
+      vim
+
+   apt-get clean && rm -rf /var/lib/apt/lists/*
+EOR
+
+##############################################################################
+# Create a conda environment and install dependencies:
+##############################################################################
+COPY build/environment-*.yml environment.yml
+RUN <<EOR
+    RESULT=$( { conda env create -f environment.yml ; } 2>&1 )
+
+    if echo $RESULT | grep -izq error ; then
+      echo "Failed to install conda dependencies: $RESULT"
+      exit 1
+    fi
+
+    conda clean -a && pip cache purge
+EOR
+
+##############################################################################
+# Install python dependencies
+##############################################################################
+ENV PIP_CACHE_DIR=/tmp/pip_cache
+RUN <<EOR
+  conda activate alltalk
+
+  pip install \
+      deepspeed-kernels \
+      scikit-learn
+EOR
+
+
+##############################################################################
+# Clone DeepSpeed sources
+##############################################################################
+RUN <<EOR
+  git clone https://github.com/microsoft/DeepSpeed.git ${STAGE_DIR}/DeepSpeed
+  cd ${STAGE_DIR}/DeepSpeed
+  git checkout .
+  git checkout "tags/v${DEEPSPEED_VERSION}" -b "v${DEEPSPEED_VERSION}"
+EOR
+
+##############################################################################
+# Create DeepSpeed build file
+##############################################################################
+RUN <<EOR
+  cat << EOF > build_deepspeed.sh
+#!/usr/bin/env bash
+source ~/.bashrc
+mkdir -p /deepspeed
+
+conda activate alltalk
+cd ${STAGE_DIR}/DeepSpeed
+DS_BUILD_OPS=1 python setup.py build_ext -j8 bdist_wheel
+mv ${STAGE_DIR}/DeepSpeed/dist/*.whl /deepspeed/
+EOF
+EOR
+
+RUN chmod +x /build_deepspeed.sh
+
+##############################################################################
+# Build DeepSpeed:
+##############################################################################
+ENTRYPOINT ["/build_deepspeed.sh"]

docker/deepspeed/build-deepspeed.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR
+
+. ${SCRIPT_DIR=}/../variables.sh
+
+# Parse arguments
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --python-version)
+      PYTHON_VERSION="$2"
+      shift
+      ;;
+    --deepspeed-version)
+      DEEPSPEED_VERSION="$2"
+      shift
+      ;;
+    *)
+      # Allow to pass arbitrary arguments to docker as well to be flexible:
+      echo "Unknown argument '$1'"
+      exit 1
+      ;;
+  esac
+  shift
+done
+
+PYTHON_SHORT_VERSION=${PYTHON_VERSION%.*}
+PYTHON_VERSION_NO_DOT=${PYTHON_SHORT_VERSION//./}
+if [[ -n $(find build -name "deepspeed-${DEEPSPEED_VERSION}*-cp${PYTHON_VERSION_NO_DOT}-cp${PYTHON_VERSION_NO_DOT}-*.whl") ]]
+then
+   echo "DeepSpeed was already built - skipping..."
+   exit 0
+fi
+
+CONDA_ENV=${SCRIPT_DIR=}/../conda/build/environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml
+if [ ! -f ${CONDA_ENV} ]; then
+    echo "No conda environment found. Please run 'build-conda-env.sh' first."
+    exit 1
+fi
+
+echo "Building DeepSpeed $DEEPSPEED_VERSION using python ${PYTHON_VERSION}"
+
+rm -rf build # make sure to properly clean up
+mkdir -p build
+cp ${CONDA_ENV} ${SCRIPT_DIR=}/build
+
+docker buildx \
+  build \
+  --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
+  -t deepspeed:$DEEPSPEED_VERSION \
+  .
+
+docker run \
+  --rm \
+  -it \
+  --gpus=all \
+  --name deepspeed \
+  -v $SCRIPT_DIR/build:/deepspeed \
+  deepspeed:$DEEPSPEED_VERSION