From cc96baefabfdb7fbd5be5a5f2f19410aef9d97f0 Mon Sep 17 00:00:00 2001 From: Abhishek Agrawal Date: Thu, 9 Jul 2026 08:48:47 -0700 Subject: [PATCH] Address blocker where benchmark runs failed to initialize JAX on Viperlite (vlp) and Ghostlite (glp) pods on Borg. PiperOrigin-RevId: 945140560 --- .../_src/testing/benchmarks/run_benchmarks.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py index ab93d150c..7f64c89a8 100644 --- a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/run_benchmarks.py +++ b/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