Merge pull request #3300 from AI-Hypercomputer:post_train_docker

PiperOrigin-RevId: 878136407

Author: Google-ML-Automation

dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile

@@ -23,6 +23,9 @@ ENV PATH="/usr/local/google-cloud-sdk/bin:/usr/local/bin/python3.12:${PATH}"
 ARG MODE
 ENV ENV_MODE=$MODE
 
+ARG WORKFLOW
+ENV ENV_WORKFLOW=$WORKFLOW
+
 ARG JAX_VERSION
 ENV ENV_JAX_VERSION=$JAX_VERSION
 
@@ -43,14 +46,15 @@ WORKDIR /deps
 # Copy setup files and dependency files separately for better caching
 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 src/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/
+COPY src/MaxText/integration/vllm/ src/MaxText/integration/vllm/
 
 # 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_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}
+RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW 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} WORKFLOW=${ENV_WORKFLOW} 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 . .

dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -118,25 +118,20 @@ run_docker_build() {
   docker build --network host $(printf -- '--build-arg %q ' "$@") -f "$dockerfile_path" -t "$LOCAL_IMAGE_NAME" .
 }
 
-# Function to build post-training image
-build_post_training_image() {
-  DOCKERFILE_NAME=""
-  if [[ ${POST_TRAINING_SOURCE} == "local" ]] ; then
-    # 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"
-  else
-    DOCKERFILE_NAME='maxtext_post_training_dependencies.Dockerfile'
-    echo "Building remote post-training dependencies: $DOCKERFILE_NAME"
-  fi
+# 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/dependencies/dockerfiles/$DOCKERFILE_NAME" \
     "MODE=${WORKFLOW}" "BASEIMAGE=${LOCAL_IMAGE_NAME}"
 }
@@ -170,7 +165,9 @@ build_tpu_image() {
 
   # Handle post-training workflow if specified
   if [[ ${WORKFLOW} == "post-training" || ${WORKFLOW} == "post-training-experimental" ]]; then
-    build_post_training_image
+    if [[ ${POST_TRAINING_SOURCE} == "local" ]]; then
+      build_post_training_deps_from_local_github
+    fi
   fi
 }
 

dependencies/scripts/docker_build_dependency_image.sh

@@ -154,6 +154,11 @@ if [[ -z "$MODE" ]]; then
   export MODE=stable
 fi
 
+# Set default value for $WORKFLOW
+if [[ -z "$WORKFLOW" ]]; then
+  export WORKFLOW=pre-training
+fi
+
 # Unset optional variables if set to NONE
 unset_optional_vars() {
     local optional_vars=("JAX_VERSION" "LIBTPU_VERSION" "LIBTPU_GCS_PATH")
@@ -185,6 +190,19 @@ install_custom_libtpu() {
     gsutil cp "$LIBTPU_GCS_PATH" "$libtpu_path"
 }
 
+install_maxtext_package_without_deps() {
+    # The MaxText package is installed separately from its dependencies to optimize
+    # docker image rebuild times by leveraging docker's layer caching.
+    # Dependencies are installed in a separate step before MaxText code is
+    # copied. This means that if MaxText code changes, but the
+    # dependencies do not, docker can reuse the cached dependency layer, leading
+    # to significantly faster image builds.
+    if [ -f 'pyproject.toml' ]; then
+        echo "Installing MaxText package without installing the dependencies (already installed)"
+        python3 -m uv pip install --no-deps -e .
+    fi
+}
+
 install_maxtext_with_deps() {
     if [[ "$DEVICE" != "tpu" && "$DEVICE" != "gpu" ]]; then
       echo -e "\n\nError: DEVICE must be either 'tpu' or 'gpu'.\n\n"
@@ -200,18 +218,31 @@ install_maxtext_with_deps() {
     python3 -m uv pip install --resolution=lowest -r "$dep_name" \
         -r 'src/install_maxtext_extra_deps/extra_deps_from_github.txt'
 
-    # The MaxText package is installed separately from its dependencies to optimize
-    # docker image rebuild times by leveraging docker's layer caching.
-    # Dependencies are installed in a separate step before MaxText code is
-    # copied. This means that if MaxText code changes, but the
-    # dependencies do not, docker can reuse the cached dependency layer, leading
-    # to significantly faster image builds.
-    if [ -f 'pyproject.toml' ]; then
-        echo "Installing MaxText package without installing the dependencies (already installed)"
-        python3 -m uv pip install --no-deps -e .
+    install_maxtext_package_without_deps
+}
+
+install_post_training_deps() {
+    if [[ "$DEVICE" != "tpu" ]]; then
+      echo -e "\n\nError: DEVICE must be 'tpu'.\n\n"
+      exit 1
     fi
+    echo "Setting up MaxText post-training workflow for $DEVICE device"
+    dep_name='dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt'
+    echo "Installing requirements from $dep_name"
+    python3 -m uv pip install --resolution=lowest -r "$dep_name"
+    python3 -m src.install_maxtext_extra_deps.install_post_train_extra_deps
 }
 
+# ---------- Post-Training workflow installation ----------
+
+if [[ "$WORKFLOW" == "post-training" ]]; then
+    install_post_training_deps
+    install_maxtext_package_without_deps
+    exit 0
+fi
+
+# ---------- Pre-Training workflow installation ----------
+
 # stable mode installation
 if [[ "$MODE" == "stable" ]]; then
     install_maxtext_with_deps