fix(test): use local checkpoints in integration tests

Author: darisoy

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,7 @@
 )
 
 
-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 +56,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 +149,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 +159,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 +173,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
   )