Configurable load planner and storage reader in the dcp_saver restore API (#824)

Summary:
Pull Request resolved: #824

Configurable load planner and storage reader in the dcp_saver restore API

# This Stack
DCP saver is the TorchTNT callback which allows checkpointing via the Distributed Checkpointing APIs. Current implementation for the Restore API doesn't expose the Load Planner and Storage Reader in the API for clients to plug in their implementations. It enforces the default planner and FsspecReader.

# This diff
DCP saver Restore API now supports planner and storage reader allowing clients to plug in their implementations.

Reviewed By: anshulverma

Differential Revision: D56922769

fbshipit-source-id: 9965c3e78463f494b92196d48c506acc39fcb04d

Author: saumishr

tests/framework/callbacks/test_dcp_saver.py

@@ -18,9 +18,12 @@
 
 import torch
 from torch import nn
-from torch.distributed.checkpoint import FileSystemWriter
-from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
-from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.checkpoint.default_planner import (
+    DefaultLoadPlanner,
+    DefaultSavePlanner,
+)
+from torch.distributed.checkpoint.metadata import Metadata, STATE_DICT_TYPE
 from torch.utils.data import DataLoader
 from torchsnapshot.test_utils import assert_state_dict_eq, check_state_dict_eq
 from torchtnt.framework._test_utils import (
@@ -348,6 +351,42 @@ def test_save_planner_storage_components(self, mock_dist_cp: MagicMock) -> None:
             self.assertIsInstance(planner, DummySavePlanner)
             self.assertIsInstance(storage_writer, DummyStorageWriter)
 
+    @patch("torchtnt.framework.callbacks.dcp_saver.dcp")
+    def test_restore_default_planner_storage_components(
+        self, mock_dist_cp: MagicMock
+    ) -> None:
+        from torch.distributed.checkpoint._fsspec_filesystem import FsspecReader
+
+        my_unit = DummyTrainUnit(input_dim=2)
+        restore_options = RestoreOptions(restore_optimizers=False)
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=restore_options,
+        )
+        planner = mock_dist_cp.load.call_args[1]["planner"]
+        storage_reader = mock_dist_cp.load.call_args[1]["storage_reader"]
+
+        self.assertIsInstance(planner, DefaultLoadPlanner)
+        self.assertIsInstance(storage_reader, FsspecReader)
+
+    @patch("torchtnt.framework.callbacks.dcp_saver.dcp")
+    def test_restore_planner_storage_components(self, mock_dist_cp: MagicMock) -> None:
+        my_unit = DummyTrainUnit(input_dim=2)
+        restore_options = RestoreOptions(restore_optimizers=False)
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=restore_options,
+            planner=DummyLoadPlanner(),
+            storage_reader=DummyStorageReader(path="path/to/snapshot"),
+        )
+        planner = mock_dist_cp.load.call_args[1]["planner"]
+        storage_reader = mock_dist_cp.load.call_args[1]["storage_reader"]
+
+        self.assertIsInstance(planner, DummyLoadPlanner)
+        self.assertIsInstance(storage_reader, DummyStorageReader)
+
 
 class DummyStatefulDataLoader:
     def __init__(self, dataloader: DataLoader) -> None:
@@ -375,9 +414,30 @@ def set_up_planner(self, state_dict: STATE_DICT_TYPE, is_coordinator: bool) -> N
         super().set_up_planner(state_dict, is_coordinator)
 
 
+class DummyLoadPlanner(DefaultLoadPlanner):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Metadata,
+        is_coordinator: bool,
+    ) -> None:
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+
 class DummyStorageWriter(FileSystemWriter):
     def __init__(self, path: str) -> None:
         super().__init__(path)
 
     def set_up_storage_writer(self, is_coordinator: bool) -> None:
         pass
+
+
+class DummyStorageReader(FileSystemReader):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    def set_up_storage_writer(self, is_coordinator: bool) -> None:
+        pass

torchtnt/framework/callbacks/dcp_saver.py

@@ -14,9 +14,12 @@
 import torch
 import torch.distributed as dist
 from torch.distributed import checkpoint as dcp
-from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
-from torch.distributed.checkpoint.planner import SavePlanner
-from torch.distributed.checkpoint.storage import StorageWriter
+from torch.distributed.checkpoint.default_planner import (
+    DefaultLoadPlanner,
+    DefaultSavePlanner,
+)
+from torch.distributed.checkpoint.planner import LoadPlanner, SavePlanner
+from torch.distributed.checkpoint.storage import StorageReader, StorageWriter
 
 from torchtnt.framework.callbacks._checkpoint_utils import (
     _prepare_app_state_for_checkpoint,
@@ -255,6 +258,8 @@ def restore(
         process_group: Optional[dist.ProcessGroup] = None,
         restore_options: Optional[RestoreOptions] = None,
         knob_options: Optional[KnobOptions] = None,
+        planner: Optional[LoadPlanner] = None,
+        storage_reader: Optional[StorageReader] = None,
     ) -> None:
         """Utility method to restore dcp checkpoint from a path.
 
@@ -269,8 +274,15 @@ def restore(
                 If torch.distributed is available and a process group is initialized, dcp assumes the intention is to save/load checkpoints in distributed fashion.
             restore_options: Controls what to  filter when restoring the state.
             knob_options: Additional keyword options for StorageWriter and StorageReader
+            planner: Instance of LoadPlanner. If this is not specificed, the default planner will be used. (Default: ``None``)
+            storage_reader: Instance of StorageReader used to perform reads. If this is not specified, it will automatically infer
+                            the reader based on the checkpoint_id. If checkpoint_id is also None, an exception will be raised. (Default: ``None``)
         """
-        storage_reader = Reader(path)
+        if planner is None:
+            planner = DefaultLoadPlanner()
+
+        if storage_reader is None:
+            storage_reader = Reader(path)
 
         restore_options = restore_options or RestoreOptions()
         app_state = _prepare_app_state_for_restore(unit, restore_options)
@@ -309,13 +321,15 @@ def restore(
                 {"app_state": MultiStateful(app_state)},
                 checkpoint_id=path,
                 storage_reader=storage_reader,
+                planner=planner,
                 process_group=process_group,
             )
         except AttributeError:
             dcp.load_state_dict(
                 {"app_state": MultiStateful(app_state)},
                 storage_reader=storage_reader,
                 process_group=process_group,
+                planner=planner,
             )
         rank_zero_info(f"Restored snapshot from path: {path}", logger=logger)