Simplify async wait hooks in DCPSaver (#888)

Summary: Pull Request resolved: #888

Reviewed By: anshulverma

Differential Revision: D61739503

fbshipit-source-id: a0e570f28769e098aa82634c559fd5e3169f551a

Author: diego-urgell

torchtnt/framework/callbacks/base_checkpointer.py

@@ -9,7 +9,7 @@
 import abc
 import logging
 from datetime import timedelta
-from typing import Any, cast, Dict, Iterable, List, Literal, Optional, Union
+from typing import Any, cast, Iterable, List, Literal, Optional, Union
 
 import torch.distributed as dist
 from pyre_extensions import none_throws
@@ -21,7 +21,6 @@
 from torchtnt.utils.checkpoint import (
     BestCheckpointConfig,
     CheckpointManager,
-    CheckpointPath,
     get_best_checkpoint_path,
     get_latest_checkpoint_path,
     MetricData,
@@ -172,7 +171,7 @@ def _generate_checkpoint_and_upkeep(
                 value=metric_value,
             )
 
-        checkpoint_path = self._generate_checkpoint_path(
+        checkpoint_path = self._checkpoint_manager.generate_checkpoint_path(
             epoch,
             step_mapping,
             metric_data,
@@ -185,7 +184,9 @@ def _generate_checkpoint_and_upkeep(
 
         if hook == "on_train_end":
             # 2.1) Make sure that last checkpoint does not already exist
-            if self._does_checkpoint_exist(checkpoint_path, self._process_group):
+            if self._checkpoint_manager.does_checkpoint_exist(
+                checkpoint_path, self._process_group
+            ):
                 rank_zero_warn(
                     "Final checkpoint already exists, skipping.", logger=logger
                 )
@@ -220,30 +221,6 @@ def _generate_checkpoint_and_upkeep(
 
         return True
 
-    def _does_checkpoint_exist(
-        self,
-        checkpoint_path: CheckpointPath,
-        process_group: Optional[dist.ProcessGroup] = None,
-    ) -> bool:
-        # Only keep this function as a hook for downstream checkpointer
-        return self._checkpoint_manager.does_checkpoint_exist(
-            checkpoint_path, process_group
-        )
-
-    def _generate_checkpoint_path(
-        self,
-        epoch: int,
-        step_mapping: Union[int, Dict[Phase, int]],
-        metric_data: Optional[MetricData] = None,
-        process_group: Optional[dist.ProcessGroup] = None,
-    ) -> CheckpointPath:
-        return self._checkpoint_manager.generate_checkpoint_path(
-            epoch,
-            step_mapping,
-            metric_data,
-            process_group=process_group,
-        )
-
     def _get_tracked_metric_value(
         self, unit: Union[TTrainUnit, TEvalUnit]
     ) -> Optional[float]:

torchtnt/framework/callbacks/dcp_saver.py

@@ -39,12 +39,7 @@
     TTrainUnit,
 )
 from torchtnt.framework.utils import get_timing_context
-from torchtnt.utils.checkpoint import (
-    BestCheckpointConfig,
-    CheckpointPath,
-    MetricData,
-    Phase,
-)
+from torchtnt.utils.checkpoint import BestCheckpointConfig
 from torchtnt.utils.optimizer import init_optim_state
 from torchtnt.utils.rank_zero_log import rank_zero_info, rank_zero_warn
 from torchtnt.utils.stateful import MultiStateful, Stateful
@@ -194,6 +189,7 @@ def _async_save(
 
         if self._prev_snapshot is not None:
             if not self._prev_snapshot.done():
+                # TODO this is unreachable at this point, since we are waiting on other functions called before _checkpoint_impl.
                 rank_zero_warn(
                     (
                         "Waiting on previous checkpoint to finish... Consider modifying checkpointing "
@@ -401,36 +397,14 @@ def restore_with_id(
             f"Restored snapshot for checkpoint_id: {checkpoint_id}", logger=logger
         )
 
-    def _does_checkpoint_exist(
-        self,
-        checkpoint_path: CheckpointPath,
-        process_group: Optional[dist.ProcessGroup] = None,
+    def _generate_checkpoint_and_upkeep(
+        self, state: State, unit: Union[TTrainUnit, TEvalUnit], hook: str
     ) -> bool:
-        # if we are still checkpointing, this might cause a collective hang.
-        # so wait here instead
+        # if we are still checkpointing, this might cause a collective hang, since several
+        # operations in the base class use the process group. So wait here instead.
         self._wait()
 
-        return super()._does_checkpoint_exist(
-            checkpoint_path=checkpoint_path, process_group=process_group
-        )
-
-    def _generate_checkpoint_path(
-        self,
-        epoch: int,
-        step_mapping: Union[int, Dict[Phase, int]],
-        metric_data: Optional[MetricData] = None,
-        process_group: Optional[dist.ProcessGroup] = None,
-    ) -> CheckpointPath:
-        # if we are still checkpointing, this might cause a collective hang.
-        # so wait here instead
-        self._wait()
-
-        return super()._generate_checkpoint_path(
-            epoch=epoch,
-            step_mapping=step_mapping,
-            metric_data=metric_data,
-            process_group=process_group,
-        )
+        return super()._generate_checkpoint_and_upkeep(state, unit, hook)
 
     @property
     def default_writer_options(self) -> Dict[str, Any]:

torchtnt/utils/checkpoint.py

@@ -206,7 +206,7 @@ def _is_phase_aware(self) -> bool:
     def newer_than(self, other: "CheckpointPath") -> bool:
         """
         Given another CheckpointPath instance, determine if this checkpoint is strictly newer than the other.
-        Note that recency is determine in terms of the epoch, phase, and number of steps. It is NOT related to
+        Note that recency is determined in terms of the epoch, phase, and number of steps. It is NOT related to
         the timestamp the checkpoint was saved.
 
         Returns: