Add integration tests with Pathways backend

Author: ycchenzheng

.github/workflows/RunTests.yml

@@ -121,6 +121,19 @@ jobs:
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
 
+  tpu_pathways_integration_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
+      pytest_marker: 'not cpu_only and not gpu_only and 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' }}
+
   gpu_unit_tests:
     needs: gpu_image
     uses: ./.github/workflows/run_tests_internal.yml
@@ -151,7 +164,7 @@ jobs:
 
   clean_up:
     if: ${{ always() }}
-    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests]
+    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests, tpu_pathways_integration_tests]
     name: "Clean up"
     runs-on: ["self-hosted"]
     permissions:
@@ -170,7 +183,7 @@ jobs:
 
   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, tpu_pathways_unit_tests]
+    needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests, tpu_pathways_integration_tests]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
@@ -202,7 +215,7 @@ jobs:
       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, tpu_pathways_unit_tests]
+      needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests, tpu_pathways_integration_tests]
       runs-on: ubuntu-latest
       permissions:
         issues: write

.github/workflows/run_pathways_tests_internal.yml

@@ -64,6 +64,9 @@ jobs:
         IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: true
         JAX_PLATFORMS: "proxy"
         JAX_BACKEND_TARGET: "grpc://localhost:29000"
+        # JAX_COORDINATOR_ADDRESS: "localhost"
+        # JAX_PROCESS_COUNT: "1"
+        # JAX_PROECESS_ID: "0"
       options: ${{ inputs.container_resource_option }}
     steps:
       - uses: actions/checkout@v4
@@ -77,6 +80,7 @@ jobs:
           python3 -m pip install -e . --no-dependencies &&
           python3 -m pip uninstall -y libtpu &&
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
+          mpirun -n 8 --allow-run-as-root \
           python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not test_default_hlo_match" --durations=0
     
     services: