chore: import upstream snapshot with attribution

deepspeed/runtime/checkpoint_engine/datastates_checkpoint_engine.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# Apache-2.0 License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
+
+# DeepSpeed Team
+
+from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
+    CheckpointEngine, CheckpointCommitInfo
+
+ENGINE_NAME = "DataStatesCheckpointEngine"
+
+
+class DataStatesCheckpointEngine(CheckpointEngine):
+
+    def __init__(self, deepspeed_config, rank):
+        super().__init__(deepspeed_config)
+        self.commit_info = None
+        self.ckpt_engine = None
+        try:
+            from datastates import CheckpointEngine as DataStatesEngine
+            self.ckpt_engine = DataStatesEngine(deepspeed_config, rank)
+        except ImportError:
+            raise RuntimeError("Please install DataStates from https://github.com/DataStates/datastates-llm.")
+        except Exception as e:
+            raise RuntimeError(f"An error occurred while initializing DataStates Checkpoint Engine: {e}")
+
+    def __del__(self):
+        self.cleanup()
+
+    def create(self, info: CheckpointCommitInfo):
+        self.commit_info = info
+        return None
+
+    def save(self, state_dict, path: str):
+        return self.ckpt_engine.save(state_dict, path)
+
+    def load(self, path: str, map_location=None):
+        return self.ckpt_engine.load(path, map_location)
+
+    def commit(self, info: CheckpointCommitInfo):
+        if info is None:
+            return
+        assert info == self.commit_info
+        self.ckpt_engine.wait(persist=True)
+        self.commit_info = None
+        return True
+
+    def cleanup(self):
+        self.commit(self.commit_info)
+        if self.ckpt_engine:
+            self.ckpt_engine.wait(persist=True)
+            del self.ckpt_engine
+
+    def is_decoupled(self):
+        return True
+
+    def preserves_storage_sharing(self):
+        return False