Merge branch 'master' into python_refactor

.github/workflows/docker-build.yml

@@ -47,11 +47,11 @@ jobs:
           else
             export FF_CUDA_ARCH=70
           fi
-          ./docker/build.sh flexflow
+          ./docker/build.sh --image_name flexflow --cuda_version 11.8
 
       - name: Check availability of Python flexflow.core module
         if: ${{ matrix.gpu_backend == 'cuda' }}
-        run: docker run --env CPU_ONLY_TEST=1 --entrypoint /bin/bash flexflow-cuda:latest -c "export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH; sudo ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1; python -c 'import flexflow.core; exit()'"
+        run: docker run --env CPU_ONLY_TEST=1 --entrypoint /bin/bash flexflow-cuda-11.8.0:latest -c "export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH; sudo ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1; python -c 'import flexflow.core; exit()'"
 
       - name: Publish Docker environment image (on push to master)
         if: github.repository_owner == 'flexflow'
@@ -60,8 +60,8 @@ jobs:
           FF_GPU_BACKEND: ${{ matrix.gpu_backend }}
         run: |
           if [[ ( ${{ github.event_name }} == 'push' || ${{ github.event_name }} == 'schedule' ) && ${GITHUB_REF#refs/heads/} == "master" ]]; then
-            ./docker/publish.sh "flexflow-environment-${FF_GPU_BACKEND}"
-            ./docker/publish.sh "flexflow-${FF_GPU_BACKEND}"
+            ./docker/publish.sh --image_name "flexflow-environment-${FF_GPU_BACKEND}" --cuda_version 11.8 
+            ./docker/publish.sh --image_name "flexflow-${FF_GPU_BACKEND}" --cuda_version 11.8
           else
             echo "No need to update Docker containers in ghrc.io registry at this time."
           fi

docker/build.sh

@@ -8,8 +8,55 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 # Cd into $FF_HOME. Assumes this script is in $FF_HOME/docker
 cd "$SCRIPT_DIR/.."
 
-# Get name of desired Docker image as input
-image="${1:-flexflow}"
+cuda_version="empty"
+image="flexflow"
+
+# Parse command-line options and # Get name of desired Docker image and cuda version as input
+while [[ $# -gt 0 ]]; do
+  key="$1"
+
+  case $key in
+    --cuda_version)
+      cuda_version="$2"
+      shift 2
+      ;;
+    --image_name)
+      image="$2"
+      shift 2
+      ;;
+    *)
+      echo "Invalid option: $key"
+      exit 1
+      ;;
+  esac
+done
+
+if [[ $cuda_version == "empty" ]]; then
+  cuda_version=$(command -v nvcc >/dev/null 2>&1 && nvcc --version | grep "release" | awk '{print $NF}')
+  # Change cuda_version eg. V11.7.99 to 11.7
+  cuda_version=${cuda_version:1:4}
+fi
+
+
+if [[ "$cuda_version" != @(11.1|11.3|11.5|11.6|11.7|11.8) ]]; then
+  # validate the verison of CUDA against a list of supported ones
+  # 11.1, 11.3, 11.5, 11.6, 11.7, 11.8
+  # Available versions: 11.1.1 | 11.2.2 | 11.3.1 | 11.5.2 | 11.6.2 | 11.7.1 | 11.8.0
+  echo "cuda_version is not supported, please choose among {11.1,11.3,11.5,11.6,11.7,11.8}"
+  exit 1
+fi
+
+# modify cuda version to available versions
+if [[ "$cuda_version" == @(11.1|11.3|11.7) ]]; then
+  cuda_version=${cuda_version}.1
+elif [[ "$cuda_version" == @(11.2|11.5|11.6) ]]; then 
+  cuda_version=${cuda_version}.2
+elif [[ "$cuda_version" == @(11.8) ]]; then 
+  cuda_version=${cuda_version}.0
+fi
+
+
+
 if [[ "$image" != @(flexflow-environment|flexflow) ]]; then
   echo "Error, image name ${image} is invalid. Choose between 'flexflow-environment' and 'flexflow'."
   exit 1
@@ -26,8 +73,8 @@ else
   echo "Letting FlexFlow build for a default GPU backend: cuda"
 fi
 
-# Build the FlexFlow Enviroment docker image
-docker build --build-arg "FF_GPU_BACKEND=${FF_GPU_BACKEND}" -t "flexflow-environment-${FF_GPU_BACKEND}" -f docker/flexflow-environment/Dockerfile .
+# Build the FlexFlow Enviroment docker image with input cuda version
+docker build --build-arg "FF_GPU_BACKEND=${FF_GPU_BACKEND}" --build-arg "cuda_version=${cuda_version}" -t "flexflow-environment-${FF_GPU_BACKEND}-${cuda_version}" -f docker/flexflow-environment/Dockerfile .
 
 # If the user only wants to build the environment image, we are done
 if [[ "$image" == "flexflow-environment" ]]; then
@@ -85,4 +132,4 @@ fi
 # Set value of BUILD_CONFIGS
 get_build_configs
 
-docker build --build-arg "N_BUILD_CORES=${n_build_cores}" --build-arg "FF_GPU_BACKEND=${FF_GPU_BACKEND}" --build-arg "BUILD_CONFIGS=${BUILD_CONFIGS}" -t "flexflow-${FF_GPU_BACKEND}" -f docker/flexflow/Dockerfile .
+docker build --build-arg "N_BUILD_CORES=${n_build_cores}" --build-arg "FF_GPU_BACKEND=${FF_GPU_BACKEND}" --build-arg "BUILD_CONFIGS=${BUILD_CONFIGS}" --build-arg "cuda_version=${cuda_version}" -t "flexflow-${FF_GPU_BACKEND}-${cuda_version}" -f docker/flexflow/Dockerfile .

docker/flexflow-environment/Dockerfile

@@ -1,4 +1,5 @@
-FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
+ARG cuda_version
+FROM nvidia/cuda:${cuda_version}-cudnn8-devel-ubuntu20.04
 
 LABEL org.opencontainers.image.source=https://github.com/flexflow/FlexFlow
 LABEL org.opencontainers.image.description="FlexFlow environment container"

docker/flexflow/Dockerfile

@@ -1,5 +1,6 @@
 ARG FF_GPU_BACKEND "cuda"
-FROM flexflow-environment-$FF_GPU_BACKEND:latest
+ARG cuda_version
+FROM flexflow-environment-$FF_GPU_BACKEND-$cuda_version:latest
 
 LABEL org.opencontainers.image.source=https://github.com/flexflow/FlexFlow
 LABEL org.opencontainers.image.description="FlexFlow container"

docker/publish.sh

@@ -4,14 +4,59 @@ set -euo pipefail
 # Cd into directory holding this script
 cd "${BASH_SOURCE[0]%/*}"
 
-image=${1:-"flexflow-cuda"}
+cuda_version="empty"
+image="flexflow-cuda"
+
+# Parse command-line options
+while [[ $# -gt 0 ]]; do
+  key="$1"
+
+  case $key in
+    --cuda_version)
+      cuda_version="$2"
+      shift 2
+      ;;
+    --image_name)
+      image="$2"
+      shift 2
+      ;;
+    *)
+      echo "Invalid option: $key"
+      exit 1
+      ;;
+  esac
+done
+
+if [[ $cuda_version == "empty" ]]; then
+  cuda_version=$(command -v nvcc >/dev/null 2>&1 && nvcc --version | grep "release" | awk '{print $NF}')
+  # Change cuda_version eg. V11.7.99 to 11.7
+  cuda_version=${cuda_version:1:4}
+fi
+
+if [[ "$cuda_version" != @(11.1|11.3|11.5|11.6|11.7|11.8) ]]; then
+  # validate the verison of CUDA against a list of supported ones
+  # 11.1, 11.3, 11.5, 11.6, 11.7, 11.8
+  echo "cuda_version is not supported, please choose among {11.1,11.3,11.5,11.6,11.7,11.8}"
+  exit 1
+fi
+
+# modify cuda version to available versions
+if [[ "$cuda_version" == @(11.1|11.3|11.7) ]]; then
+  cuda_version=${cuda_version}.1
+elif [[ "$cuda_version" == @(11.2|11.5|11.6) ]]; then 
+  cuda_version=${cuda_version}.2
+elif [[ "$cuda_version" == @(11.8) ]]; then 
+  cuda_version=${cuda_version}.0
+fi
+
+
 if [[ "${image}" != @(flexflow-environment-cuda|flexflow-environment-hip_cuda|flexflow-environment-hip_rocm|flexflow-environment-intel|flexflow-cuda|flexflow-hip_cuda|flexflow-hip_rocm|flexflow-intel) ]]; then
   echo "Error, image name ${image} is invalid. Choose between 'flexflow-environment-{cuda,hip_cuda,hip_rocm,intel}' and 'flexflow-{cuda,hip_cuda,hip_rocm,intel}'."
   exit 1
 fi
 
 # Check that image exists
-docker image inspect "${image}":latest > /dev/null
+docker image inspect "${image}-${cuda_version}":latest > /dev/null
 
 # Log into container registry
 FLEXFLOW_CONTAINER_TOKEN=${FLEXFLOW_CONTAINER_TOKEN:-}
@@ -21,7 +66,9 @@ echo "$FLEXFLOW_CONTAINER_TOKEN" | docker login ghcr.io -u flexflow --password-s
 # Tag image to be uploaded
 git_sha=${GITHUB_SHA:-$(git rev-parse HEAD)}
 if [ -z "$git_sha" ]; then echo "Commit hash cannot be detected, cannot publish the docker image to ghrc.io"; exit; fi
-docker tag "$image":latest ghcr.io/flexflow/"$image":latest
+
+docker tag "${image}-${cuda_version}":latest ghcr.io/flexflow/"$image-$cuda_version":latest
+
 
 # Upload image
-docker push ghcr.io/flexflow/"$image":latest
+docker push ghcr.io/flexflow/"$image-$cuda_version":latest