Add gpu image creation pipeline

Signed-off-by: Kunjan <kunjanp@google.com>

Author: coolkp

.github/workflows/UploadDockerImages.yml

@@ -35,3 +35,16 @@ jobs:
     - name: build maxdiffusion jax nightly image
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly
+
+  build-gpu-image:
+    runs-on: ["self-hosted", "e2", "cpu"]
+    steps:
+    - uses: actions/checkout@v3
+    - name: Cleanup old docker images
+      run: docker system prune --all --force
+    - name: build maxdiffusion jax stable stack gpu image
+      run: |
+        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack_gpu MODE=stable_stack PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack_gpu BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest DEVICE=gpu
+    - name: build maxdiffusion jax nightly image
+      run: |
+         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly_gpu MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly DEVICE=gpu

.github/workflows/build_and_upload_images.sh

@@ -34,13 +34,15 @@ for ARGUMENT in "$@"; do
     echo "$KEY"="$VALUE"
 done
 
+export DEVICE="${DEVICE:-tpu}"
+
 if [[ ! -v CLOUD_IMAGE_NAME ]] || [[ ! -v PROJECT ]] || [[ ! -v MODE ]] ; then
   echo "You must set CLOUD_IMAGE_NAME, PROJECT and MODE"
   exit 1
 fi
 
 gcloud auth configure-docker us-docker.pkg.dev --quiet
-bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE=$MODE
+bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE=$MODE DEVICE=$DEVICE
 image_date=$(date +%Y-%m-%d)
 
 # Upload only dependencies image

maxdiffusion_gpu_dependencies.Dockerfile

@@ -22,8 +22,7 @@ RUN apt-get update && apt-get install -y google-cloud-sdk
 # Set environment variables for Google Cloud SDK
 ENV PATH="/usr/local/google-cloud-sdk/bin:${PATH}"
 
-# Upgrade libcusprase to work with Jax
-RUN apt-get update && apt-get install -y libcusparse-12-3
+
 
 ARG MODE
 ENV ENV_MODE=$MODE
@@ -45,6 +44,7 @@ RUN ls .
 
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
-
+RUN pip install -r requirements.txt
+RUN pip install -U "jax[cuda12]"
 
 WORKDIR /deps