Refactor CPU vs TPU tests.

TPU tests are marked by the tag `requires-tpu`.  CPU tests are anything else.

CPU tests need to be configured to disable all devices except CPU, otherwise
the test may hang when searching for TPUs.

For TPU tests, we add the required LIBTPU flags in .bazelrc.

PiperOrigin-RevId: 715081513

Author: cantonios

.bazelrc

@@ -101,7 +101,10 @@ build:clang --copt=-Qunused-arguments
 ##############################################################################
 # Test configurations.
 ##############################################################################
-test:cpu --test_env=JAX_PLATFORMS=cpu --test_tag_filters=cpu
+# Configure TPU for SparseCore usage.
+test --test_env="LIBTPU_INIT_ARGS=--2a886c8_chip_config_name=megachip_tccontrol"
+# Show output for failing tests.
+test --test_output=errors
 
 #############################################################################
 # Some configs to make getting some forms of debug builds. In general, the

.github/workflows/build_and_test.yml

@@ -70,4 +70,4 @@ jobs:
 
       - name: Run CPU tests
         run: |
-          bazel test --config=cpu --test_output=errors --keep_going //...
+          bazel test --test_tag_filters=-requires-tpu --test_output=errors --keep_going //...

jax_tpu_embedding/sparsecore/examples/shakespeare/BUILD

@@ -34,6 +34,7 @@ pytype_strict_contrib_test(
     main = "jax_sc_shakespeare_jit.py",
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/examples/models/shakespeare:dataset",
@@ -105,6 +106,7 @@ pytype_strict_contrib_test(
     main = "jax_sc_shakespeare_pmap.py",
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/examples/models/shakespeare:dataset",

jax_tpu_embedding/sparsecore/lib/checkpointing/BUILD

@@ -31,7 +31,7 @@ pytype_strict_library(
 pytype_strict_contrib_test(
     name = "checkpoint_utils_test",
     srcs = ["checkpoint_utils_test.py"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         ":checkpoint_utils",
         "//jax_tpu_embedding/sparsecore/lib/nn:embedding",

jax_tpu_embedding/sparsecore/lib/core/BUILD

@@ -36,7 +36,7 @@ cc_library(
 cc_test(
     name = "input_preprocessing_threads_test",
     srcs = ["input_preprocessing_threads_test.cc"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         ":input_preprocessing_threads",
         "@com_google_googletest//:gtest_main",
@@ -65,7 +65,7 @@ cc_library(
 cc_test(
     name = "input_preprocessing_util_test",
     srcs = ["input_preprocessing_util_test.cc"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         ":input_preprocessing_util",
         "@com_google_googletest//:gtest_main",
@@ -118,9 +118,9 @@ pytype_strict_contrib_test(
         "input_preprocessing_test.py",
     ],
     env = {
+        "JAX_PLATFORMS": "cpu",
         "XLA_FLAGS": "--xla_dump_to=sponge",
     },
-    tags = ["cpu"],
     deps = [
         ":constants",
         ":input_preprocessing",
@@ -136,9 +136,9 @@ pytype_strict_contrib_test(
         "input_preprocessing_cc_test.py",
     ],
     env = {
+        "JAX_PLATFORMS": "cpu",
         "XLA_FLAGS": "--xla_dump_to=sponge",
     },
-    tags = ["cpu"],
     deps = [
         ":constants",
         ":input_preprocessing_cc",

jax_tpu_embedding/sparsecore/lib/core/primitives/tests/BUILD

@@ -33,6 +33,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
@@ -59,6 +60,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
@@ -85,6 +87,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_sgd_with_mini_batching",
@@ -111,6 +114,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_adagrad_with_mini_batching",
@@ -138,6 +142,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
@@ -164,6 +169,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
@@ -188,6 +194,7 @@ pytype_strict_contrib_test(
     ],
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",

jax_tpu_embedding/sparsecore/lib/fdo/BUILD

@@ -39,7 +39,7 @@ pytype_strict_library(
 pytype_strict_contrib_test(
     name = "file_fdo_client_test",
     srcs = ["file_fdo_client_test.py"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         ":file_fdo_client",
         pypi_requirement("absl/testing:absltest"),

jax_tpu_embedding/sparsecore/lib/flax/tests/BUILD

@@ -30,6 +30,7 @@ pytype_strict_contrib_test(
     ],
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/flax:embed",
@@ -61,6 +62,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/examples/models/shakespeare:dataset",

jax_tpu_embedding/sparsecore/lib/nn/tests/BUILD

@@ -25,7 +25,7 @@ pytype_strict_contrib_test(
     srcs = [
         "preprocess_sparse_dense_matmul_input_test.py",
     ],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing",
         "//jax_tpu_embedding/sparsecore/lib/nn:embedding",
@@ -64,7 +64,7 @@ pytype_strict_library(
 pytype_strict_contrib_test(
     name = "test_utils_test",
     srcs = ["test_utils_test.py"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         ":test_utils",
         pypi_requirement("absl/logging"),
@@ -89,6 +89,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         ":test_utils",
@@ -118,6 +119,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         ":test_utils",
@@ -143,6 +145,7 @@ pytype_strict_contrib_test(
     ],
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/nn:embedding",
@@ -168,6 +171,7 @@ pytype_strict_contrib_test(
     ],
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         ":test_utils",
@@ -185,7 +189,7 @@ pytype_strict_contrib_test(
 pytype_strict_contrib_test(
     name = "embedding_spec_test",
     srcs = ["embedding_spec_test.py"],
-    tags = ["cpu"],
+    env = {"JAX_PLATFORMS": "cpu"},
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/nn:embedding_spec",
         pypi_requirement("absl/testing:absltest"),

jax_tpu_embedding/sparsecore/tests/BUILD

@@ -31,6 +31,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/examples/models/shakespeare:dataset",
@@ -60,6 +61,7 @@ pytype_strict_contrib_test(
     },
     tags = [
         "exclusive-if-local",
+        "requires-tpu",
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/examples/models/shakespeare:dataset",