Merge pull request #3956 from AI-Hypercomputer:debug-pathways-devices-fix

PiperOrigin-RevId: 918617201

Author: Google-ML-Automation

.github/workflows/run_pathways_tests.yml

@@ -68,6 +68,7 @@ jobs:
         IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: true
         JAX_PLATFORMS: "proxy"
         JAX_BACKEND_TARGET: "grpc://localhost:29000"
+        TPU_VISIBLE_DEVICES: ""
       options: ${{ inputs.container_resource_option }}
     steps:
       - name: Checkout MaxText
@@ -85,7 +86,6 @@ jobs:
           source .venv/bin/activate
           maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
           uv pip install ${maxtext_wheel}[tpu] --resolution=lowest
-          uv pip uninstall libtpu
           install_tpu_pre_train_extra_deps
           python3 --version
           python3 -m pip freeze

src/maxtext/utils/generate_param_only_checkpoint.py

@@ -137,6 +137,8 @@ def _generate_lora_decode_checkpoints(config, mesh):
         config.enable_checkpointing,
         config.async_checkpointing,
         config.checkpoint_period,
+        use_ocdbt=config.checkpoint_storage_use_ocdbt,
+        use_zarr3=config.checkpoint_storage_use_zarr3,
     )
 
     lora_config, lora_state, lora_state_annotations = lora_utils.setup_initial_lora_state(
@@ -192,6 +194,8 @@ def generate_decode_checkpoint(config):
       config.enable_checkpointing,
       config.async_checkpointing,
       config.checkpoint_period,
+      use_ocdbt=config.checkpoint_storage_use_ocdbt,
+      use_zarr3=config.checkpoint_storage_use_zarr3,
   )
   # Read training state from config.load_paramaters_path
   max_logging.log(f"Read training checkpoint from: {config.load_full_state_path}")

tests/conftest.py

@@ -66,6 +66,12 @@ def _custom_iter(self):
 except AttributeError:
   pass
 
+import os
+
+if os.getenv("JAX_PLATFORMS") == "proxy":
+  # Import maxtext early to register the pathways proxy backend before JAX is queried.
+  import maxtext  # pylint: disable=unused-import
+
 try:
   _HAS_TPU = any(d.platform == "tpu" for d in jax.devices())
 except Exception:  # pragma: no cover  pylint: disable=broad-exception-caught