test gpu docker images

google · Mar 5, 2024 · 0f33dde · 0f33dde
1 parent 8b793d5
commit 0f33dde
Showing 1 changed file with 36 additions and 35 deletions.
diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
@@ -18,46 +18,47 @@
 name: Build Images
 
 on:
+ pull_request:
  schedule:
  # Run the job daily at 7PM PST (3AM UTC)
  - cron: '0 3 * * *'
 
 jobs:
- tpu:
- build_images:
- strategy:
- fail-fast: false
- matrix:
- device-type: ["v4-8"]
- runs-on: ["self-hosted", "tpu", "${{ matrix.device-type }}"]
- steps:
- - uses: actions/checkout@v3
- - name: build jax stable image
- run : |
- project=tpu-prod-env-multipod
- local_image_name=maxtext_local_jax_stable
- cloud_image_name=maxtext_jax_stable
- bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=stable
- docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
- gcloud auth configure-docker --quiet
- docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
- docker push gcr.io/$project/${cloud_image_name}:latest
- image_date=$(date +%Y-%m-%d)
- docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
- docker push gcr.io/$project/${cloud_image_name}:${image_date}
- - name: build jax nightly image
- run : |
- project=tpu-prod-env-multipod
- local_image_name=maxtext_local_jax_nightly
- cloud_image_name=maxtext_jax_nightly
- bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=nightly
- docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
- gcloud auth configure-docker --quiet
- docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
- docker push gcr.io/$project/${cloud_image_name}:latest
- image_date=$(date +%Y-%m-%d)
- docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
- docker push gcr.io/$project/${cloud_image_name}:${image_date}
+ # tpu:
+ #  build_images:
+ #  strategy:
+ #  fail-fast: false
+ #  matrix:
+ #  device-type: ["v4-8"]
+ #  runs-on: ["self-hosted", "tpu", "${{ matrix.device-type }}"]
+ #  steps:
+ #  - uses: actions/checkout@v3
+ #  - name: build jax stable image
+ #  run : |
+ #  project=tpu-prod-env-multipod
+ #  local_image_name=maxtext_local_jax_stable
+ #  cloud_image_name=maxtext_jax_stable
+ #  bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=stable
+ #  docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
+ #  gcloud auth configure-docker --quiet
+ #  docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
+ #  docker push gcr.io/$project/${cloud_image_name}:latest
+ #  image_date=$(date +%Y-%m-%d)
+ #  docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
+ #  docker push gcr.io/$project/${cloud_image_name}:${image_date}
+ #  - name: build jax nightly image
+ #  run : |
+ #  project=tpu-prod-env-multipod
+ #  local_image_name=maxtext_local_jax_nightly
+ #  cloud_image_name=maxtext_jax_nightly
+ #  bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$local_image_name MODE=nightly
+ #  docker build --build-arg BASEIMAGE=${local_image_name} -f ./maxtext_runner.Dockerfile -t ${local_image_name}_runner .
+ #  gcloud auth configure-docker --quiet
+ #  docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:latest
+ #  docker push gcr.io/$project/${cloud_image_name}:latest
+ #  image_date=$(date +%Y-%m-%d)
+ #  docker tag ${local_image_name}_runner gcr.io/$project/${cloud_image_name}:${image_date}
+ #  docker push gcr.io/$project/${cloud_image_name}:${image_date}
 
  gpu:
  build_images: