Maxtext unit tests with Pathways backend.

.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

.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"

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

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