diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml index 389510bcf9..eb0c472f2c 100644 --- a/.github/workflows/UploadDockerImages.yml +++ b/.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 }} diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml index 67c2bee7e3..dda2c08c59 100644 --- a/.github/workflows/build_and_push_docker_image.yml +++ b/.github/workflows/build_and_push_docker_image.yml @@ -37,10 +37,6 @@ on: image_date: required: true type: string - base_image: - required: false - type: string - default: '' workflow: required: false type: string @@ -125,7 +121,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' diff --git a/src/dependencies/scripts/docker_build_dependency_image.sh b/src/dependencies/scripts/docker_build_dependency_image.sh index 1321e77d9f..406794a49f 100644 --- a/src/dependencies/scripts/docker_build_dependency_image.sh +++ b/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