Merge pull request #3242 from ROCm:tpu-backend-skips

PiperOrigin-RevId: 875996905

Author: Google-ML-Automation

pytest.ini

@@ -24,6 +24,7 @@ addopts =
     --ignore=tests/unit/engram_vs_reference_test.py
 markers =
     tpu_only: marks tests to be run on TPUs only
+    tpu_backend: marks tests that require a TPU-enabled JAX install (TPU PJRT plugin), but not TPU hardware
     gpu_only: marks tests to be run on GPUs only
     cpu_only: marks tests to be run on CPUs only
     decoupled: tests that validate offline / DECOUPLE_GCLOUD=TRUE mode.

tests/conftest.py

@@ -23,6 +23,7 @@
 import pytest
 from maxtext.common.gcloud_stub import is_decoupled
 import jax
+import importlib.util
 
 # Configure JAX to use unsafe_rbg PRNG implementation to match main scripts.
 if is_decoupled():
@@ -42,6 +43,28 @@
 GCP_MARKERS = {"external_serving", "external_training"}
 
 
+def _has_tpu_backend_support() -> bool:
+  """Whether JAX has TPU backend support installed (PJRT TPU plugin).
+
+  This is intentionally *not* the same as having TPU hardware available.
+  """
+  try:
+    if importlib.util.find_spec("jaxlib") is None:
+      return False
+  except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught
+    return False
+
+  # Heuristic: TPU-enabled jaxlib exposes a TPU client/extension module.
+  # This check avoids initializing any backend at collection time.
+  for mod in ("jaxlib.tpu_client", "jaxlib._src.tpu_client", "jaxlib.tpu_extension"):
+    try:
+      if importlib.util.find_spec(mod) is not None:
+        return True
+    except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught
+      continue
+  return False
+
+
 def pytest_collection_modifyitems(config, items):
   """Customize pytest collection behavior.
 
@@ -55,12 +78,21 @@ def pytest_collection_modifyitems(config, items):
 
   skip_no_tpu = None
   skip_no_gpu = None
+  skip_no_tpu_backend = None
   if not _HAS_TPU:
     skip_no_tpu = pytest.mark.skip(reason="Skipped: requires TPU hardware, none detected")
 
   if not _HAS_GPU:
     skip_no_gpu = pytest.mark.skip(reason="Skipped: requires GPU hardware, none detected")
 
+  if not _has_tpu_backend_support():
+    skip_no_tpu_backend = pytest.mark.skip(
+        reason=(
+            "Skipped: requires a TPU-enabled JAX install (TPU PJRT plugin). "
+            "Install a TPU-enabled jax/jaxlib build to run this test."
+        )
+    )
+
   for item in items:
     # Iterate thru the markers of every test.
     cur_test_markers = {m.name for m in item.iter_markers()}
@@ -76,6 +108,11 @@ def pytest_collection_modifyitems(config, items):
       remaining.append(item)
       continue
 
+    if skip_no_tpu_backend and "tpu_backend" in cur_test_markers:
+      item.add_marker(skip_no_tpu_backend)
+      remaining.append(item)
+      continue
+
     if decoupled and (cur_test_markers & GCP_MARKERS):
       # Deselect tests marked as external_serving/training entirely.
       deselected.append(item)
@@ -98,6 +135,7 @@ def pytest_configure(config):
   for m in [
       "gpu_only: tests that require GPU hardware",
       "tpu_only: tests that require TPU hardware",
+      "tpu_backend: tests that require a TPU-enabled JAX install (TPU PJRT plugin), but not TPU hardware",
       "external_serving: JetStream / serving / decode server components",
       "external_training: goodput integrations",
       "decoupled: marked on tests that are not skipped due to GCP deps, when DECOUPLE_GCLOUD=TRUE",

tests/unit/sharding_compare_test.py

@@ -33,6 +33,8 @@
 
 Transformer = models.transformer_as_linen
 
+pytestmark = [pytest.mark.cpu_only, pytest.mark.tpu_backend]
+
 
 def compute_checksum(d: dict) -> str:
   """Compute a checksum (SHA256) of a dictionary."""

tests/unit/train_compile_test.py

@@ -28,7 +28,7 @@
 from maxtext.trainers.pre_train.train_compile import main as train_compile_main
 from tests.utils.test_helpers import get_test_config_path
 
-pytestmark = [pytest.mark.external_training]
+pytestmark = [pytest.mark.external_training, pytest.mark.tpu_backend]
 
 
 class TrainCompile(unittest.TestCase):