Repository: mosaicml/composer

Issue:
Handle checkpointing crashing if the checkpoint file does not exist
I'm seeing situations where, if the checkpoint file does not exist, an exception will be raised about some unrelated temporary directory. Furthermore, the other ranks will not be killed, causing the process to be stuck (and cloud provider instances continue to run)

Traceback:
```
Traceback (most recent call last):
  File "/workspace/src/composer/composer/trainer/checkpoint.py", line 198, in _download_checkpoint
    with tarfile.open(rank_zero_checkpoint_archive_filepath) as tarball:
  File "/usr/lib/python3.8/tarfile.py", line 1603, in open
    return func(name, "r", fileobj, **kwargs)
  File "/usr/lib/python3.8/tarfile.py", line 1667, in gzopen
    fileobj = GzipFile(name, mode + "b", compresslevel, fileobj)
  File "/usr/lib/python3.8/gzip.py", line 173, in __init__
    fileobj = self.myfileobj = builtins.open(filename, mode or 'rb')
FileNotFoundError: [Errno 2] No such file or directory: '/tmp/tmpj8wdttfd/rank_0.it14000_8979.tar'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/workspace/src/composer/examples/run_mosaic_trainer.py", line 40, in <module>
    main()
  File "/workspace/src/composer/examples/run_mosaic_trainer.py", line 35, in main
    trainer = Trainer.create_from_hparams(hparams=hparams)
  File "/workspace/src/composer/composer/trainer/trainer.py", line 407, in create_from_hparams
    trainer = cls(
  File "/workspace/src/composer/composer/trainer/trainer.py", line 345, in __init__
    restored_seed = self.checkpoint_loader.load_checkpoint(state=self.state)
  File "/workspace/src/composer/composer/trainer/checkpoint.py", line 292, in load_checkpoint
    mosaic_checkpoint_filepath, extracted_checkpoint_folder = self._download_checkpoint(node_checkpoint_folder)
  File "/workspace/src/composer/composer/trainer/checkpoint.py", line 202, in _download_checkpoint
    raise RuntimeError(f"Unable to retrieve checkpoint {checkpoint_name}") from e
RuntimeError: Unable to retrieve checkpoint bert_checkpoints/bert_baseline/rank_0/bert_checkpoints/it14000_8979.tar

wandb: Waiting for W&B process to finish, PID 1838... (failed 1). Press ctrl-c to abort syncing.
wandb:
wandb: Run history:
wandb:   ActFnSearch/Event.INIT ▁
wandb:
wandb: Run summary:
wandb:   ActFnSearch/Event.INIT 1
wandb:
wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Synced Rank 0: https://wandb.ai/mosaic-ml/moin-data-bert-base-cola-aws-v100/runs/26cn3s85
wandb: Find logs at: ./wandb/run-20220128_033239-26cn3s85/logs/debug.log
```

To be clear, this is user error -- we should:
1) Ensure the checkpoint file exists, rather than say "unable to retrieve it"
2) Handle crashing more gracefully such that the Kubernetes job is also killed


Discussion:

Resolving patch:
diff --git a/composer/trainer/checkpoint.py b/composer/trainer/checkpoint.py
--- a/composer/trainer/checkpoint.py
+++ b/composer/trainer/checkpoint.py
@@ -182,37 +182,45 @@ def _download_checkpoint(self, node_checkpoint_folder: str) -> Tuple[str, Option
             extracted_checkpoint_folder = os.path.join(node_checkpoint_folder, "checkpoint")
             mosaic_checkpoint_filepath = os.path.join(extracted_checkpoint_folder, _MOSAIC_STATES_FILENAME)
 
-        if dist.get_local_rank() == 0:
-            # every NODE needs the GLOBAL rank zero checkpoint
-            self._retrieve_checkpoint(destination_filepath=rank_zero_checkpoint_archive_filepath,
-                                      rank=dist.get_global_rank(),
-                                      ignore_not_found_errors=False)
-            if extracted_checkpoint_folder is not None:
-                try:
-                    with tarfile.open(rank_zero_checkpoint_archive_filepath) as tarball:
+        try:
+            if dist.get_local_rank() == 0:
+                # every NODE needs the GLOBAL rank zero checkpoint
+                self._retrieve_checkpoint(destination_filepath=rank_zero_checkpoint_archive_filepath,
+                                          rank=dist.get_global_rank(),
+                                          ignore_not_found_errors=False)
+                if extracted_checkpoint_folder is not None:
+                    try:
+                        with tarfile.open(rank_zero_checkpoint_archive_filepath) as tarball:
+                            tarball.extractall(extracted_checkpoint_folder)
+                    except FileNotFoundError:
+                        checkpoint_name = self.path.format(rank=dist.get_global_rank())
+                        # Not re-raising the file-not-found error as that is irrelevant;
+                        # the underlying issue is that the checkpoint file does not exist on the disk
+                        # or could not be downloaded
+                        raise RuntimeError(f"Checkpoint {checkpoint_name} does not exist")
+
+            if rank_zero_checkpoint_archive_filepath != rank_n_checkpoint_archive_filepath:
+                # every RANK needs ITS OWN checkpoint.
+                # But, the  global rank zero is a special case -- these files are the same!
+                assert dist.get_global_rank() != 0, "invariant violation"
+
+                # Allowing not-found errors to be ignored as sometimes there won't be rank-local checkpoints
+                # (e.g. when not using deepspeed)
+                self._retrieve_checkpoint(destination_filepath=rank_n_checkpoint_archive_filepath,
+                                          rank=dist.get_global_rank(),
+                                          ignore_not_found_errors=True)
+
+                if extracted_checkpoint_folder is not None:
+                    # it's an archive and needs to be extracted
+                    with tarfile.open(rank_n_checkpoint_archive_filepath) as tarball:
                         tarball.extractall(extracted_checkpoint_folder)
-                except FileNotFoundError as e:
-                    checkpoint_name = self.path.format(rank=dist.get_global_rank())
-                    raise RuntimeError(f"Unable to retrieve checkpoint {checkpoint_name}") from e
-
-        if rank_zero_checkpoint_archive_filepath != rank_n_checkpoint_archive_filepath:
-            # every RANK needs ITS OWN checkpoint.
-            # But, the  global rank zero is a special case -- these files are the same!
-            assert dist.get_global_rank() != 0, "invariant violation"
-
-            # Allowing not-found errors to be ignored as sometimes there won't be rank-local checkpoints
-            # (e.g. when not using deepspeed)
-            self._retrieve_checkpoint(destination_filepath=rank_n_checkpoint_archive_filepath,
-                                      rank=dist.get_global_rank(),
-                                      ignore_not_found_errors=True)
-
-            if extracted_checkpoint_folder is not None:
-                # it's an archive and needs to be extracted
-                with tarfile.open(rank_n_checkpoint_archive_filepath) as tarball:
-                    tarball.extractall(extracted_checkpoint_folder)
-
-        # Wait for all checkpoints on the node to finish downloading
-        dist.barrier()
+        finally:
+            # Wait for all checkpoints on the node to finish downloading
+            # Putting the barrier in a finally so the rank will always block on the barrier,
+            # even if it has an exception.
+            # Any exception will be re-raised after the barrier passes. The launcher script
+            # will detect the process crash and terminate the other ranks
+            dist.barrier()
 
         return mosaic_checkpoint_filepath, extracted_checkpoint_folder
 
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -778,7 +778,6 @@ def _train_loop(self) -> None:
 
                         # total_loss can be None if gradient scaling failed
                         dist.all_reduce(total_loss, reduce_operation="SUM")
-                        dist.barrier()
                         full_loss = total_loss.cpu().item()
                         self.logger.metric_batch({'loss/train': full_loss / dist.get_world_size()})
 


Test patch: