stabilize GHA test workflows and use local checkpoints

.github/workflows/run_tests_against_package.yml

@@ -162,13 +162,15 @@ jobs:
             # Dynamically discover the 'nvidia' folder and prepend all its sub-library
             # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent
             # JAX from partially loading incompatible system-level CUDA libraries.
-            NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1)
-            if [ -n "${NVIDIA_DIR}" ]; then
-              for dir in "${NVIDIA_DIR}"/*; do
-                if [ -d "$dir/lib" ]; then
-                  export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH}
-                fi
-              done
+            if [ -d ".venv/lib" ]; then
+              NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1)
+              if [ -n "${NVIDIA_DIR}" ]; then
+                for dir in "${NVIDIA_DIR}"/*; do
+                  if [ -d "$dir/lib" ]; then
+                    export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH}
+                  fi
+                done
+              fi
             fi
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then

tests/integration/checkpoint_compatibility_test.py

@@ -49,6 +49,7 @@ def run_checkpoint_compatibility(hardware, attention_type):
       "grain_worker_count=0",
       "grain_train_files=/tmp/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*",
   ]
+  local_ckpt_dir = "/tmp/maxtext_local_output"
 
   # Run training using grain input pipeline
   train_main(
@@ -60,6 +61,7 @@ def run_checkpoint_compatibility(hardware, attention_type):
           attention_type=attention_type,
           dataset_type="grain",
           dataset_path="/tmp/gcsfuse",
+          base_output_directory=local_ckpt_dir,
       )
       + grain_command
   )
@@ -74,6 +76,7 @@ def run_checkpoint_compatibility(hardware, attention_type):
           attention_type=attention_type,
           dataset_type="tfds",
           dataset_path="/tmp/gcsfuse",
+          base_output_directory=local_ckpt_dir,
       )
   )
 

tests/integration/checkpointing_test.py

@@ -41,7 +41,9 @@
 )
 
 
-def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention_type, dataset_type, dataset_path):
+def get_checkpointing_command(
+    run_date, hardware, steps, metrics_file, attention_type, dataset_type, dataset_path, base_output_directory=None
+):
   """Generates a command list for a checkpointing test run.
 
   Args:
@@ -56,7 +58,8 @@ def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention
   Returns:
     A list of strings representing the command line arguments.
   """
-  base_output_directory = get_test_base_output_directory()
+  if base_output_directory is None:
+    base_output_directory = get_test_base_output_directory()
   model_params = [
       "base_emb_dim=128",
       "base_num_query_heads=2",
@@ -148,6 +151,7 @@ def run_checkpointing(hardware, attention_type):
       "grain_worker_count=0",
       f"grain_train_files={selected_pattern}",
   ]
+  local_ckpt_dir = "/tmp/maxtext_local_output"
   train_main(
       get_checkpointing_command(
           run_date,
@@ -157,6 +161,7 @@ def run_checkpointing(hardware, attention_type):
           attention_type=attention_type,
           dataset_type="grain",
           dataset_path=dataset_path,
+          base_output_directory=local_ckpt_dir,
       )
       + grain_command
   )
@@ -170,6 +175,7 @@ def run_checkpointing(hardware, attention_type):
           attention_type=attention_type,
           dataset_type="grain",
           dataset_path=dataset_path,
+          base_output_directory=local_ckpt_dir,
       )
       + grain_command
   )