Maxtext unit tests with Pathways backend.

RoshaniN · RoshaniN · commit 1755edb8501d · 2025-03-06T18:32:17.000Z
Modified Pathways workflow to run specifically on Pathways runner.

Final tests on Pathways.

New changes to help run tests.

Trial to get Pathways working.

run Pathways containers and Maxtext as part of the same job.

Move installation to script

More changes.

Installing docker also as part of the script.

Simplified flow test.

Few more changes.

New way of installation

Docker compose with Maxtext and Pathways containers.

Other changes to use maxtext container in docker compose YAML.

Trying to merge all services together.

Adding everything as services again.

Directly running docker compose.

Elegant solution with Pathways containers as services.

Elegant solution with Pathways containers as services.

Reverting to Maxtext in docker compose.

Changes to run with latest JAX SS image.

Run tests with correct markers.

Move installation to Github runner step, simplify command.

Formalize the changes.

Formalize the changes.

Formalize the changes.

Notify on Pathways test failures, modify test script.

Error handling.
diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml
@@ -67,6 +67,15 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
 
+  tpu_pathways_unit_tests:
+    needs: [tpu_image]
+    uses: ./.github/workflows/run_pathways_tests_internal.yml
+    with:
+      device_type: tpu
+      device_name: v4-8
+      pytest_marker: 'not gpu_only and not integration_test'
+      test_directory: 'tests'
+
   tpu_integration_tests:
     needs: tpu_image
     uses: ./.github/workflows/run_tests_internal.yml
@@ -105,7 +114,7 @@ jobs:
 
   clean_up:
     if: ${{ always() }}  # always execute, regardless of previous jobs or steps.
-    needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests]
+    needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_pathways_unit_tests, tpu_integration_tests]
     name: "Clean up"
     runs-on: ["self-hosted"]
     permissions:
@@ -119,7 +128,7 @@ jobs:
 
   notify:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests]
+    needs: [gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_pathways_unit_tests, tpu_integration_tests]
     runs-on: ["self-hosted"]
     steps:
     - name: Check whether one of the jobs failed
diff --git a/.github/workflows/run_pathways_tests_internal.yml b/.github/workflows/run_pathways_tests_internal.yml
@@ -0,0 +1,50 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This file runs unit tests with Pathways backend.
+
+name: Run Pathways Tests
+
+on:
+  workflow_call:
+    inputs:
+      device_type:
+        required: true
+        type: string
+      device_name:
+        required: true
+        type: string
+      pytest_marker:
+        required: true
+        type: string
+      test_directory:
+        required: true
+        type: string
+
+jobs:
+  start_pathways:
+    name: Start Pathways tests on (${{ inputs.device_name }})
+    runs-on: ["self-hosted", "${{ inputs.device_type }}", "${{ inputs.device_name }}"]
+    steps:
+      - name: Install prerequisites on the self-hosted runner for Pathways
+        run: |
+          echo "Setting up the prerequisites"
+          apt-get install -y docker --quiet
+          apt-get install -y docker-compose-plugin --quiet
+          docker compose version # To ensure docker compose is installed
+          gcloud auth configure-docker us-docker.pkg.dev --quiet
+      - uses: actions/checkout@v4
+      - name: Start Pathways and Maxtext containers using script
+        run: |
+          bash docker_run_pathways_containers.sh maxtext_image=gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu command="cd MaxText ; python3 -m pytest ${{ inputs.test_directory }} -m '${{ inputs.pytest_marker }}' -s"
diff --git a/docker_run_pathways_containers.sh b/docker_run_pathways_containers.sh
@@ -0,0 +1,48 @@
+#!/bin/bash
+
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Examples -
+# TPU unit tests - bash docker_run_pathways_containers.sh maxtext_image=<your Maxtext image> command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
+# Subset of unit tests - bash docker_run_pathways_containers.sh maxtext_image=<your Maxtext image> command="cd MaxText ; python3 -m pytest tests/train_tests.py -m 'not gpu_only and not integration_test' -s"
+
+echo "Running docker_run_pathways_containers.sh"
+
+# Stop execution if any command exits with error
+set -e
+
+# Parse input variables
+for ARGUMENT in "$@"; do
+    IFS='=' read -r KEY VALUE <<< "$ARGUMENT"
+    export "$KEY"="$VALUE"
+    echo "$KEY"="$VALUE"
+done
+
+cd utils_pathways
+
+# Error handling - Run 'docker compose down' if 'docker compose up' errors.
+# Exit the script with the exit code of docker compose up.
+clean_up () {
+    overall_test_exit_code=$?
+    echo "Tests are not successful, exit code is $overall_test_exit_code"
+    MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose down
+    exit $overall_test_exit_code
+
+}
+trap clean_up EXIT
+
+# Setting up and tearing down the test setup -
+MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose up --exit-code-from maxtext
+MAXTEXT_IMAGE=${maxtext_image} COMMAND=${command} docker compose down
diff --git a/utils_pathways/docker-compose.yml b/utils_pathways/docker-compose.yml
@@ -0,0 +1,78 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+services:
+
+  resource_manager:
+    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
+    ports:
+      - "29001:29001"
+      - "29002:29002"
+    entrypoint:
+      - /usr/pathways/run/cloud_pathways_server_sanitized
+      - --server_port=29001
+      - --node_type=resource_manager
+      - --instance_count=1
+      - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
+      - --instance_type=tpuv4:2x2x1
+    environment:
+      - HOST_ADDRESS=resource_manager
+      - TPU_SKIP_MDS_QUERY=true
+
+
+  worker:
+    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
+    ports:
+      - "29005:29005"
+      - "29006:29006"
+      - "8471:8471"
+      - "8080:8080"
+    entrypoint:
+      - /usr/pathways/run/cloud_pathways_server_sanitized
+      - --alsologtostderr
+      - --server_port=29005
+      - --resource_manager_address=resource_manager:29001
+      - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
+    privileged: true
+    depends_on:
+      - resource_manager
+
+
+  proxy:
+    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
+    ports:
+      - "29008:29008"
+    entrypoint:
+      - /usr/pathways/run/cloud_proxy_server_sanitized
+      - --server_port=29008
+      - --resource_manager_address=resource_manager:29001
+      - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
+    depends_on:
+      - worker
+      - resource_manager
+
+
+  maxtext:
+    image: ${MAXTEXT_IMAGE}
+    privileged: true
+    volumes:
+      - ../:/deps
+    environment:
+      - JAX_PLATFORMS=proxy
+      - JAX_BACKEND_TARGET=grpc://proxy:29008
+    entrypoint: ["bash", "-c"]
+    command:
+      - ${COMMAND}
+    depends_on:
+      - proxy
