Refactor dockerfiles and deprecate stable_stack and jax_ai_image mode

.github/workflows/UploadDockerImages.yml

@@ -65,11 +65,11 @@ jobs:
           - device: tpu
             build_mode: stable
             image_name: maxtext_jax_stable
-            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
           - device: tpu
             build_mode: nightly
             image_name: maxtext_jax_nightly
-            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}

.github/workflows/build_and_push_docker_image.yml

@@ -113,7 +113,7 @@ jobs:
             DEVICE=${{ inputs.device }}
             MODE=${{ inputs.build_mode }}
             JAX_VERSION=NONE
-            LIBTPU_GCS_PATH=NONE
+            LIBTPU_VERSION=NONE
             BASEIMAGE=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ inputs.image_date }}
 
       - name: Add tags to Docker image

PREFLIGHT.md

@@ -26,7 +26,7 @@ bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m
 ```
 
 For GKE,
-`numactl` should be built into your docker image from [maxtext_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/dependencies/dockerfiles/maxtext_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
+`numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
 
 ```
 bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m MaxText.train src/MaxText/configs/base.yml run_name=$YOUR_JOB_NAME

RESTRUCTURE.md

@@ -27,14 +27,9 @@ comments, or questions by creating a new
 ├── README.md
 ├── dependencies/
 │   ├── dockerfiles/
-│   │   ├── jetstream_pathways.Dockerfile
-│   │   ├── maxengine_server.Dockerfile
 │   │   ├── maxtext_custom_wheels.Dockerfile
-│   │   ├── maxtext_db_dependencies.Dockerfile
-│   │   ├── maxtext_dependencies.Dockerfile
+│   │   ├── maxtext_tpu_dependencies.Dockerfile
 │   │   ├── maxtext_gpu_dependencies.Dockerfile
-│   │   ├── maxtext_jax_ai_image.Dockerfile
-│   │   ├── maxtext_libtpu_path.Dockerfile
 │   │   └── maxtext_runner.Dockerfile
 │   ├── requirements/
 │   │   └── requirements.txt

dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG BASEIMAGE
+ARG BASEIMAGE=maxtext_base_image
 FROM ${BASEIMAGE}
-ARG MODE
 
+ARG MODE
 ENV MODE=$MODE
 
 RUN echo "Installing Post-Training dependencies (vLLM, tpu-inference, tunix) with MODE=${MODE}"

dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG BASEIMAGE
+ARG BASEIMAGE=maxtext_base_image
 FROM ${BASEIMAGE}
+
 ARG MODE
 ENV MODE=$MODE
 
@@ -31,11 +32,9 @@ COPY tunix /tunix
 RUN pip uninstall -y google-tunix
 RUN pip install -e /tunix --no-cache-dir
 
-
 COPY vllm /vllm
 RUN VLLM_TARGET_DEVICE="tpu" pip install -e /vllm --no-cache-dir
 
-
 COPY tpu-inference /tpu-inference
 RUN pip install -e /tpu-inference --no-cache-dir
 

dependencies/dockerfiles/maxtext_db_dependencies.Dockerfile → dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -1,6 +1,7 @@
 # syntax=docker/dockerfile:experimental
-# Copy benchmark-db
-FROM gcr.io/tpu-prod-env-one-vm/benchmark-db:2025-02-14
+
+ARG BASEIMAGE=python:3.12-slim-bullseye
+FROM $BASEIMAGE
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y curl gnupg
@@ -25,8 +26,8 @@ ENV ENV_MODE=$MODE
 ARG JAX_VERSION
 ENV ENV_JAX_VERSION=$JAX_VERSION
 
-ARG LIBTPU_GCS_PATH
-ENV ENV_LIBTPU_GCS_PATH=$LIBTPU_GCS_PATH
+ARG LIBTPU_VERSION
+ENV ENV_LIBTPU_VERSION=$LIBTPU_VERSION
 
 ARG DEVICE
 ENV ENV_DEVICE=$DEVICE
@@ -44,9 +45,12 @@ COPY tools/setup tools/setup/
 COPY dependencies/requirements/ dependencies/requirements/
 COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
 
+# Copy the custom libtpu.so file if it exists inside maxtext repository
+COPY libtpu.so* /root/custom_libtpu/
+
 # Install dependencies - these steps are cached unless the copied files change
-RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION LIBTPU_GCS_PATH=${ENV_LIBTPU_GCS_PATH} DEVICE=${ENV_DEVICE}"
-RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_GCS_PATH=${ENV_LIBTPU_GCS_PATH} DEVICE=${ENV_DEVICE}
+RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
+RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
 
 # Now copy the remaining code (source files that may change frequently)
 COPY . .