Add: Added checkpoint cleanup after elastic event

Author: camiloCienet

axlearn/common/launch_trainer.py

@@ -145,6 +145,8 @@ def get_trainer_config(
         )
     trainer_config: SpmdTrainer.Config = trainer_config_fn()
     trainer_config.dir = trainer_config.dir or flag_values.trainer_dir
+
+    print(f"Trainer Config Dir: {trainer_config.dir} by Camilo")
     if flag_values.mesh_selector is not None:
         select_mesh_config(trainer_config, mesh_selector=flag_values.mesh_selector)
     trainer_config.mesh_axis_names = trainer_config.mesh_axis_names or ("data", "model")

axlearn/common/launch_trainer_main.py

@@ -3,6 +3,7 @@
 """Main function for launching the trainer."""
 
 import pathwaysutils
+import functools
 from absl import app, flags
 from pathwaysutils.elastic import elastic, manager
 
@@ -17,10 +18,10 @@
 def main(_):
     measurement.initialize(flags.FLAGS)
     launch.setup()
-    # trainer_config = launch_trainer.get_trainer_config()
+    trainer_config = launch_trainer.get_trainer_config()
     # trainer_config.set(recorder=config_for_function(lambda: measurement.global_recorder))
     # measurement.start_monitoring()
-
+    clean_up_checkpoints = functools.partial(utils.clean_up_checkpoints, checkpoint_dir=trainer_config.dir)
     if pathwaysutils.is_pathways_backend_used() and enable_elastic_training:
 
         def train():
@@ -69,6 +70,7 @@ def pre_callback():
                 max_resizes=10,  # Handle up to 10 slice up or slice down transitions
                 poll_interval=30,  # Monitor thread checks inactive slice health every 30 seconds
                 pre_callback=pre_callback,
+                on_elastic_event_callback=clean_up_checkpoints,
             )(train)
 
         train()

axlearn/common/utils.py

@@ -28,6 +28,7 @@
 from collections.abc import Mapping, Sequence
 from enum import Enum
 from functools import cache
+import subprocess
 from typing import (
     Any,
     Callable,
@@ -109,6 +110,39 @@ def live_devices():
 def live_slice_indices() -> set[int]:
     return {d.slice_index for d in live_devices()}
 
+def clean_up_checkpoints(checkpoint_dir: str):
+
+  print(f"Checking for incomplete checkpoint after an elastic event...Check dir: {checkpoint_dir}")
+
+  # 1. List the directory
+  new_checkpoint_dir = f"{checkpoint_dir}/checkpoints/"
+  result = subprocess.run(['gsutil', 'ls', new_checkpoint_dir], capture_output=True, text=True)
+
+  if result.returncode != 0:
+    print("Failed to inspect checkpoint dir. Continuing")
+    return
+  print(f"Checkpoints==> {[line for line in result.stdout.splitlines()]}")
+  checkpoints = [line for line in result.stdout.splitlines()]
+
+  if not checkpoints:
+    print("Found no existing checkpoints. Continuing")
+    return
+
+  # Sort naturally (Version sort) and get the last one
+  checkpoints.sort(key=lambda x: [int(c) if c.isdigit() else c for c in re.split(r'(\d+)', x)])
+  latest_checkpoint = checkpoints[-1]
+
+  print(f"Checking latest checkpoint: {latest_checkpoint}")
+
+  # 3. Check for commit_success file
+  # gsutil -q stat returns 0 if found, non-zero if not
+  stat_check = subprocess.run(['gsutil', '-q', 'stat', f"{latest_checkpoint}commit_success*"])
+
+  if stat_check.returncode != 0:
+    print(f"No commit_success file found. Deleting {latest_checkpoint}...")
+    subprocess.run(['gsutil', '-m', 'rm', '-rf', latest_checkpoint])
+  else:
+    print(f"Found commit_success file. Keeping {latest_checkpoint}.")
 
 @dataclasses.dataclass
 class HybridMeshShape: