get_state_dict_for_key API with support for flat and nested state dicts for DCP checkpoints (#932)

Summary:
Pull Request resolved: #932

get_state_dict_for_key API with user provided state dict

Reviewed By: diego-urgell, JKSenthil

Differential Revision: D64502115

fbshipit-source-id: 865a043f61f07596fb83aac40025da27aa3d214d

Author: saumishr

torchtnt/framework/_test_utils.py

@@ -297,3 +297,39 @@ def state_dict(self) -> Dict[str, Any]:
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self.sum = state_dict["sum"]
         self.count = state_dict["count"]
+
+
+class DummyStatefulConfig:
+    def __init__(
+        self,
+        storage_path: str,
+        lazy_loading: bool,
+        num_workers_per_gpu: int,
+        max_batch_length: int,
+    ) -> None:
+        self.storage_path = storage_path
+        self.lazy_loading = lazy_loading
+        self.num_workers_per_gpu = num_workers_per_gpu
+        self.max_batch_length = max_batch_length
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {
+            "storage_path": self.storage_path,
+            "data": {
+                "lazy_loading": self.lazy_loading,
+                "train": {
+                    "num_workers_per_gpu": self.num_workers_per_gpu,
+                    "dynamic_batch_config": {
+                        "max_batch_length": self.max_batch_length,
+                    },
+                },
+            },
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        self.storage_path = state_dict["storage_path"]
+        self.lazy_loading = state_dict["data"]["lazy_loading"]
+        self.num_workers_per_gpu = state_dict["data"]["train"]["num_workers_per_gpu"]
+        self.max_batch_length = state_dict["data"]["train"]["dynamic_batch_config"][
+            "max_batch_length"
+        ]