Address blocker where benchmark runs failed to initialize JAX on Viperlite (vlp) and Ghostlite (glp) pods on Borg.

PiperOrigin-RevId: 945140560

Author: abhishek002002

checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py

@@ -20,14 +20,17 @@
 using standard environment variables like JAX_COORDINATOR_ADDRESS,
 JAX_PROCESS_ID and JAX_NUM_PROCESSES.
 """
-# pylint: disable=g-statement-before-imports,g-import-not-at-top
+import os
 
-try:  # SimDevice import must occur before JAX.
-  import simdevice  # pylint: disable=unused-import
-except (ImportError, FileNotFoundError):
-  pass
+# pylint: disable=g-statement-before-imports,g-import-not-at-top
+if os.environ.get('SIMDEVICE_CONFIG_PATH') or os.environ.get(
+    'ENABLE_SIMDEVICE', ''
+).lower() in ('true', '1'):
+  try:  # SimDevice import must occur before JAX.
+    import simdevice  # pylint: disable=unused-import
+  except (ImportError, FileNotFoundError):
+    pass
 
-import os
 
 from absl import app
 from absl import flags