Assign gpu pools only to test actions (#418)

* Assign gpu pools only to test actions

* Ignore flaky tests

Author: alekstheod

build_tools/rocm/run_xla_multi_gpu.sh

@@ -60,12 +60,17 @@ if [ ! -d /tf/pkg ]; then
 fi
 
 EXCLUDED_TESTS=(
-  CollectiveOpsTestE2E.MemcpyP2pLargeMessage
-  RaggedAllToAllTest/RaggedAllToAllTest.RaggedAllToAll_8GPUs_2ReplicasPerGroups/sync_decomposer
-  RaggedAllToAllTest/RaggedAllToAllTest.RaggedAllToAll_8GPUs_2ReplicasPerGroups/async_decomposer
+  # //xla/tests:collective_ops_test_amdgpu_any
+  RaggedAllToAllTest*
+  AsyncCollectiveOps*
+  AsyncMemcpyCollectiveOps*
+  CollectiveOpsTest*
+  AllReduceTest*
+  Fp8CollectiveOpsTest*
   # //xla/backends/gpu/codegen/triton:fusion_emitter_parametrized_legacy_test_amdgpu_any
   ElementwiseTestSuiteF32/BinaryElementwiseTest.ElementwiseFusionExecutesCorrectly/f32_atan2
   # //xla/tests:collective_ops_e2e_test_amdgpu_any
+  CollectiveOpsTestE2E.MemcpyP2pLargeMessage
   CollectiveOpsTestE2EPipelinedNonPipelined.CollectivePipelinerBackward
   CollectiveOpsTestE2EPipelinedNonPipelined.CollectivePipelinerBackwardStartFromOne
   # //xla/tools/multihost_hlo_runner:functional_hlo_runner_test
@@ -87,20 +92,6 @@ elif [[ $1 == "tsan" ]]; then
     SANITIZER_ARGS+=("--run_under=//build_tools/rocm:sanitizer_wrapper")
     SANITIZER_ARGS+=("--config=tsan")
     TAG_FILTERS="$TAG_FILTERS,-notsan"
-    # excluded from tsan
-    EXCLUDED_TESTS+=(
-        CollectiveOpsTest*
-        Fp8CollectiveOpsTest.AllGather_8BitFloat
-        Fp8CollectiveOpsTest.CollectivePermute_8BitFloat
-        Fp8CollectiveOpsTest.AllToAll_8BitFloat
-        AsyncCollectiveOps*
-        AllReduceTest*
-        RaggedAllToAllTest*
-        AsyncCollectiveOps*
-        AsyncMemcpyCollectiveOps*
-        RaggedAllToAllTest*
-    )
-
     #  tsan tests appear to be flaky in rbe due to the heavy load
     #  force them to run locally
     RBE_OPTIONS+=(

xla/tests/BUILD

@@ -2806,7 +2806,10 @@ xla_test(
         ],
     },
     backends = ["gpu"],
-    tags = ["test_migrated_to_hlo_runner_pjrt"],
+    tags = [
+        "test_migrated_to_hlo_runner_pjrt",
+        "local",
+    ],
     deps = [
         ":hlo_pjrt_test_base",
         ":literal_test_util",

xla/tsl/platform/default/build_config_root.bzl

@@ -20,11 +20,11 @@ GPU_TEST_PROPERTIES = {
 }
 
 ROCM_SINGLE_GPU_TEST_PROPERTIES = {
-    "Pool": "linux_x64_gpu",
+    "test.Pool": "linux_x64_gpu",
 }
 
 ROCM_MULTI_GPU_TEST_PROPERTIES = {
-    "Pool": "linux_x64_multigpu",
+    "test.Pool": "linux_x64_multigpu",
 }
 
 def tf_gpu_tests_tags():