Get a shared file system safe signal file name (mosaicml#3485)

mvpatel2000 · Jul 23, 2024 · 29401b8 · 29401b8
1 parent 6f8831d
commit 29401b8
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
@@ -1986,7 +1986,7 @@ def _get_autoresume_checkpoint(
 
             signal_file_path = os.path.join(
                 os.path.dirname(latest_checkpoint_path),
-                f'.node_{dist.get_node_rank()}_local_rank0_completed_autoresume',
+                dist.get_node_signal_file_name(),
             )
             if dist.get_local_rank() == 0:
                 os.makedirs(os.path.dirname(signal_file_path), exist_ok=True)

diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -819,7 +819,7 @@ def download_checkpoint(
         if not checkpoint_is_sharded:
             signal_file_path = os.path.join(
                 node_checkpoint_folder,
-                f'.node_{dist.get_node_rank()}_local_rank0_completed',
+                dist.get_node_signal_file_name(),
             )
             if dist.get_local_rank() == 0:
                 with open(signal_file_path, 'wb') as f:

diff --git a/composer/utils/dist.py b/composer/utils/dist.py
@@ -644,7 +644,7 @@ def get_node_signal_file_name(rng: Optional[random.Random] = None) -> str:
     random_string = ''.join(rng.choices(string.ascii_letters + string.digits, k=6))
     node_rank = get_node_rank()
     file_name_list = [f'._signal_file_node{node_rank}_{random_string}']
-    dist.broadcast_object_list(file_name_list, src=0)
+    broadcast_object_list(file_name_list, src=0)
     return file_name_list[0]