Add workflow for unit tests with Pathways backend

.github/workflows/RunTests.yml

@@ -90,7 +90,21 @@
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
 
+  tpu_pathways_unit_tests:
+    needs: tpu_image
+    uses: ./.github/workflows/run_pathways_tests_internal.yml
+    with:
+      device_type: tpu
+      device_name: v4-8
+      # cloud_runner: linux-x86-ct4p-240-4tpu
+      cloud_runner: linux-x86-ct5lp-112-4tpu-test
+      pytest_marker: 'not cpu_only and not gpu_only and not integration_test'
+      xla_python_client_mem_fraction: 0.75
+      tf_force_gpu_allow_growth: false
+      container_resource_option: "--privileged"
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+
   tpu_integration_tests:
     needs: tpu_image
     uses: ./.github/workflows/run_tests_internal.yml
     with:
@@ -131,7 +145,7 @@
 
   clean_up:
     if: ${{ always() }}
-    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests]
+    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests]
     name: "Clean up"
     runs-on: ["self-hosted"]
     permissions:
@@ -150,7 +164,7 @@
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests]
+    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
@@ -182,7 +196,7 @@
       name: Close issue after 3 successful builds
       # This job runs only if all the preceding test jobs succeeded
       if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
-      needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests]
+      needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests]
       runs-on: ubuntu-latest
       permissions:
         issues: write

.github/workflows/run_pathways_tests_internal.yml

@@ -0,0 +1,102 @@
+# 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
+      image_type:
+        required: false
+        type: string
+      pytest_marker:
+        required: true
+        type: string
+      is_scheduled_run:
+        required: true
+        type: string
+      xla_python_client_mem_fraction:
+        required: true
+        type: string
+      tf_force_gpu_allow_growth:
+        required: true
+        type: string
+      container_resource_option:
+        required: true
+        type: string
+      cloud_runner:
+        required: false
+        type: string
+
+jobs:
+  run:
+    runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
+    container:
+      image: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ inputs.image_type != '' && inputs.image_type || inputs.device_type }}
+      env:
+        XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
+        TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
+        TPU_SKIP_MDS_QUERY: ${{ inputs.image_type == 'tpu' && inputs.device_type != 'tpu' && '1' || '' }}
+        JAX_PLATFORMS: "proxy"
+        JAX_BACKEND_TARGET: "grpc://proxy:29000"
+      options: ${{ inputs.container_resource_option }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Run Tests
+        run: |
+          if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then
+            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
+          else
+            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only"
+          fi
+          python3 -m pip install -e . --no-dependencies &&
+          python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0
+
+    services:
+      resource_manager:
+        image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
+        ports:
+          - "29001:29001"
+          - "29002:29002"
+        options:
+          --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29001, --node_type=resource_manager, --instance_count=1, --instance_type=tpuv4:2x2x1, --gcs_scratch_location=gs://cloud-pathways-staging/tmp]
+        env:
+          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"
+        options:
+          --privileged
+          --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29005, --resource_manager_address=resource_manager:29001, --gcs_scratch_location=gs://cloud-pathways-staging/tmp]
+
+      proxy:
+        image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
+        ports:
+          - "29000:29000"
+        options:
+          --entrypoint=[/usr/pathways/run/cloud_proxy_server_sanitized, --server_port=29000, --resource_manager_address=resource_manager:29001, --gcs_scratch_location=gs://cloud-pathways-staging/tmp]