Merge pull request #3378 from AI-Hypercomputer:cleanup_post_training_source

PiperOrigin-RevId: 882705751

Author: Google-ML-Automation

.github/workflows/UploadDockerImages.yml

@@ -79,32 +79,18 @@ jobs:
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
 
-  tpu-post-training-stable:
-    name: tpu-post-training-stable
-    needs: setup
-    uses: ./.github/workflows/build_and_push_docker_image.yml
-    with:
-      image_name: maxtext_post_training_stable
-      device: tpu
-      build_mode: stable
-      workflow: post-training
-      dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
-      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-
   tpu-post-training-nightly:
     name: tpu-post-training-nightly
-    needs: [setup, tpu-post-training-stable]
+    needs: [setup]
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: maxtext_post_training_nightly
       device: tpu
       build_mode: nightly
       workflow: post-training
-      dockerfile: ./src/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
+      dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
-      base_image: gcr.io/tpu-prod-env-multipod/maxtext_post_training_stable:${{ needs.setup.outputs.image_date }}
 
   gpu-pre-training:
     name: ${{ matrix.image_name }}

.github/workflows/build_and_push_docker_image.yml

@@ -37,10 +37,6 @@ on:
       image_date:
         required: false
         type: string
-      base_image:
-        required: false
-        type: string
-        default: ''
       workflow:
         required: false
         type: string
@@ -129,7 +125,6 @@ jobs:
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
             INCLUDE_TEST_ASSETS=true
-            ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'

src/dependencies/scripts/docker_build_dependency_image.sh

@@ -46,12 +46,9 @@
 # POST-TRAINING BUILD EXAMPLES
 # ==================================
 
-# Build docker image with stable pre-training dependencies and stable post-training dependencies
+# Build docker image with post-training dependencies
 ## bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training
 
-# Build docker image with stable pre-training dependencies and post-training dependencies from GitHub head
-## bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training POST_TRAINING_SOURCE=local
-
 if [ "${BASH_SOURCE-}" ]; then
   this_file="${BASH_SOURCE[0]}"
 elif [ "${ZSH_VERSION-}" ]; then
@@ -118,24 +115,6 @@ run_docker_build() {
   docker build --network host $(printf -- '--build-arg %q ' "$@") -f "$dockerfile_path" -t "$LOCAL_IMAGE_NAME" .
 }
 
-# Function to build post-training dependencies from local Github head
-build_post_training_deps_from_local_github() {
-  # To install vllm, tunix, tpu-inference from a local path, we copy it into the build context, excluding __pycache__.
-  # This assumes vllm, tunix, tpu-inference is a sibling directory to the current one (maxtext).
-  rsync -a --exclude='__pycache__' ../tpu-inference .
-  rsync -a --exclude='__pycache__' ../vllm .
-  rsync -a --exclude='__pycache__' ../tunix .
-
-  # The cleanup is set to run even if the build fails to remove the copied directory.
-  trap "rm -rf ./tpu-inference ./vllm ./tunix" EXIT INT TERM
-
-  DOCKERFILE_NAME='maxtext_post_training_local_dependencies.Dockerfile'
-  echo "Building local post-training dependencies: $DOCKERFILE_NAME"
-
-  run_docker_build "$MAXTEXT_REPO_ROOT/src/dependencies/dockerfiles/$DOCKERFILE_NAME" \
-    "MODE=${WORKFLOW}" "BASEIMAGE=${LOCAL_IMAGE_NAME}"
-}
-
 # Function to build image for GPUs
 build_gpu_image() {
   if [[ ${MODE} == "pinned" ]]; then
@@ -162,13 +141,6 @@ build_tpu_image() {
 
   echo "Building docker image with arguments: ${docker_build_args[*]}"
   run_docker_build "$MAXTEXT_REPO_ROOT/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile" "${docker_build_args[@]}"
-
-  # Handle post-training workflow if specified
-  if [[ ${WORKFLOW} == "post-training" || ${WORKFLOW} == "post-training-experimental" ]]; then
-    if [[ ${POST_TRAINING_SOURCE} == "local" ]]; then
-      build_post_training_deps_from_local_github
-    fi
-  fi
 }
 
 if [[ ${DEVICE} == "gpu" ]]; then