Merge pull request AI-Hypercomputer#2908 from AI-Hypercomputer:custom_wheels_mode

PiperOrigin-RevId: 852865112

Author: Google-ML-Automation

RESTRUCTURE.md

@@ -27,7 +27,6 @@ comments, or questions by creating a new
 ├── README.md
 ├── dependencies/
 │   ├── dockerfiles/
-│   │   ├── maxtext_custom_wheels.Dockerfile
 │   │   ├── maxtext_tpu_dependencies.Dockerfile
 │   │   ├── maxtext_gpu_dependencies.Dockerfile
 │   │   └── maxtext_runner.Dockerfile

dependencies/dockerfiles/maxtext_custom_wheels.Dockerfile

@@ -96,18 +96,18 @@ if [[ -z ${JAX_VERSION+x} ]] ; then
 fi
 if [[ -z ${MODE} ]]; then
   export MODE=stable
-# TODO: Remove 'custom_wheels' mode support when tpu-recipes migration is complete.
-elif [[ ${MODE} == "custom_wheels" ]]; then
-  export WORKFLOW=custom-wheels
-  export MODE=nightly
 fi
 if [[ -z ${DEVICE} ]]; then
   export DEVICE=tpu
 fi
+if [[ -z ${WORKFLOW} ]]; then
+  export WORKFLOW=pre-training
+fi
 
 # Create docker build arguments array
 docker_build_args=(
   "DEVICE=${DEVICE}"
+  "WORKFLOW=${WORKFLOW}"
   "MODE=${MODE}"
   "JAX_VERSION=${JAX_VERSION}"
 )
@@ -172,12 +172,6 @@ build_tpu_image() {
   if [[ ${WORKFLOW} == "post-training" || ${WORKFLOW} == "post-training-experimental" ]]; then
     build_post_training_image
   fi
-
-  # TODO: Remove 'custom_wheels' mode support when tpu-recipes migration is complete.
-  if [[ ${WORKFLOW} == "custom-wheels" ]]; then
-    echo "Building custom wheels dependencies."
-    run_docker_build "$MAXTEXT_REPO_ROOT/dependencies/dockerfiles/maxtext_custom_wheels.Dockerfile" "BASEIMAGE=${LOCAL_IMAGE_NAME}"
-  fi
 }
 
 if [[ ${DEVICE} == "gpu" ]]; then