Merge pull request #2394 from AI-Hypercomputer:rbierneni-move-tunix-installation

PiperOrigin-RevId: 811531652

Author: Google-ML-Automation

.github/workflows/RunTests.yml

@@ -111,6 +111,7 @@ jobs:
       device_name: a100-40gb-4
       cloud_runner: linux-x86-a2-48-a100-4gpu
       pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
+      pytest_addopts: '--ignore=tests/sft_hooks_test.py'
       xla_python_client_mem_fraction: 0.65
       tf_force_gpu_allow_growth: true
       container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
@@ -124,6 +125,7 @@ jobs:
       device_name: a100-40gb-4
       cloud_runner: linux-x86-a2-48-a100-4gpu
       pytest_marker: 'not cpu_only and not tpu_only and integration_test'
+      pytest_addopts: '--ignore=tests/sft_hooks_test.py'
       xla_python_client_mem_fraction: 0.65
       tf_force_gpu_allow_growth: true
       container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"

.github/workflows/run_tests_internal.yml

@@ -31,6 +31,10 @@ on:
       pytest_marker:
         required: true
         type: string
+      pytest_addopts:
+        required: false
+        type: string
+        default: ''
       is_scheduled_run:
         required: true
         type: string
@@ -67,4 +71,4 @@ jobs:
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only"
           fi
           python3 -m pip install -e . --no-dependencies &&
-          LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0
+          LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0

maxtext_jax_ai_image.Dockerfile

@@ -47,6 +47,11 @@ RUN if [ "$DEVICE" = "tpu" ] && [ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pkg
         python3 -m pip install -r /deps/requirements_with_jax_ai_image.txt; \
   fi
 
+# Install tunix at a pinned commit for TPU devices, skip for GPU
+RUN if [ "$DEVICE" = "tpu" ]; then \
+        python3 -m pip install 'google-tunix @ https://github.com/google/tunix/archive/e03f154edd2abfddd6b9babb72a3d0d3c4d81bb2.zip'; \
+  fi
+
 # Now copy the remaining code (source files that may change frequently)
 COPY . .
 
@@ -61,7 +66,7 @@ RUN if [ "$TEST_TYPE" = "xlml" ] || [ "$TEST_TYPE" = "unit_test" ]; then \
     fi
 
 # Run the script available in JAX AI base image to generate the manifest file
-RUN bash /jax-stable-stack/generate_manifest.sh PREFIX=maxtext COMMIT_HASH=$COMMIT_HASH
+RUN bash /jax-ai-image/generate_manifest.sh PREFIX=maxtext COMMIT_HASH=$COMMIT_HASH
 
 # Install (editable) MaxText
 RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .

requirements.txt

@@ -37,7 +37,6 @@ tensorflow-text
 tensorflow
 tiktoken
 transformers
-google-tunix @ https://github.com/google/tunix/archive/e03f154edd2abfddd6b9babb72a3d0d3c4d81bb2.zip
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
 qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

requirements_with_jax_ai_image.txt

@@ -23,5 +23,4 @@ tensorflow-datasets
 tensorflow-text>=2.17.0
 tiktoken
 transformers
-google-tunix @ https://github.com/google/tunix/archive/e03f154edd2abfddd6b9babb72a3d0d3c4d81bb2.zip
 qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

setup.sh

@@ -146,7 +146,7 @@ if [[ "$MODE" == "nightly" ]]; then
     # Remove/update this section based on the pinned github repo commit in requirements.txt
     sed -i -E \
       -e 's|^mlperf-logging @ https?://github.com/mlcommons/logging/archive/.*\.zip$|mlperf-logging@git+https://github.com/mlperf/logging.git|' \
-      -e '/^tunix/!s|^([^ ]*) @ https?://github.com/([^/]*\/[^/]*)/archive/.*\.zip$|\1@git+https://github.com/\2.git|' \
+      -e 's|^([^ ]*) @ https?://github.com/([^/]*\/[^/]*)/archive/.*\.zip$|\1@git+https://github.com/\2.git|' \
       requirements.txt.nightly-temp
 
     echo "--- Installing modified nightly requirements: ---"
@@ -189,14 +189,10 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
             python3 -m uv pip install 'jax[tpu]>0.4' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         fi
 
-        if [[ -n "$LIBTPU_GCS_PATH" ]]; then
-            # Install custom libtpu
-            echo "Installing libtpu.so from $LIBTPU_GCS_PATH to $libtpu_path"
-            # Install required dependency
-            python3 -m uv pip install -U crcmod
-            # Copy libtpu.so from GCS path
-            gsutil cp "$LIBTPU_GCS_PATH" "$libtpu_path"
-        fi
+        # TODO: Once tunix has support for GPUs, move it from here to requirements.txt
+        echo "Installing tunix for stable TPU environment"
+        python3 -m uv pip install 'google-tunix @ https://github.com/google/tunix/archive/e03f154edd2abfddd6b9babb72a3d0d3c4d81bb2.zip'
+
         if [[ -n "$LIBTPU_GCS_PATH" ]]; then
             # Install custom libtpu
             echo "Installing libtpu.so from $LIBTPU_GCS_PATH to $libtpu_path"
@@ -256,6 +252,8 @@ elif [[ $MODE == "nightly" ]]; then
         fi
         echo "Installing nightly tensorboard plugin profile"
         python3 -m uv pip install tbp-nightly --upgrade
+        # Installing tunix
+        python3 -m uv pip install 'git+https://github.com/google/tunix.git'
     fi
     echo "Installing nightly tensorboard plugin profile"
     python3 -m uv pip install tbp-nightly --upgrade

tests/sft_hooks_test.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 """Tests for training and data loading hooks for SFT"""
+import pytest
+
+pytestmark = pytest.mark.tpu_only
 
 import jax