intelligent-machine-learning
diff --git a/‎dlrover/python/common/comm.py‎
Lines changed: 8 additions & 0 deletions b/‎dlrover/python/common/comm.py‎
Lines changed: 8 additions & 0 deletions
diff --git a/‎dlrover/python/common/constants.py‎
Lines changed: 1 addition & 0 deletions b/‎dlrover/python/common/constants.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎dlrover/python/common/global_context.py‎
Lines changed: 2 additions & 0 deletions b/‎dlrover/python/common/global_context.py‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎dlrover/python/elastic_agent/master_client.py‎
Lines changed: 4 additions & 0 deletions b/‎dlrover/python/elastic_agent/master_client.py‎
Lines changed: 4 additions & 0 deletions
diff --git a/‎dlrover/python/elastic_agent/torch/ckpt_saver.py‎
Lines changed: 147 additions & 10 deletions b/‎dlrover/python/elastic_agent/torch/ckpt_saver.py‎
Lines changed: 147 additions & 10 deletions
@@ -223,6 +223,14 @@ class GlobalStep(Message):
     elapsed_time_per_step: float = 0.0
 
 
+@dataclass
+class RdzvBlocked(Message):
+    """Indicate whether rendezvous completion is blocked."""
+
+    blocked: bool = False
+    reason: str = ""
+
+
 @dataclass
 class HeartBeat(Message):
     timestamp: int = 0
 
@@ -315,6 +315,7 @@ class NodeEnv(object):
     RELAUNCHED_POD = "RELAUNCHED_POD"
     DLROVER_MASTER_ADDR = "DLROVER_MASTER_ADDR"
     DLROVER_MASTER_SERVICE_TYPE = "DLROVER_MASTER_SERVICE_TYPE"
+    DLROVER_TRAINING_ELASTIC_MODE = "DLROVER_TRAINING_ELASTIC_MODE"
     GRPC_ENABLE_FORK = "GRPC_ENABLE_FORK_SUPPORT"
     GRPC_POLL_STRATEGY = "GRPC_POLL_STRATEGY"
     POD_NAME = "POD_NAME"
 
@@ -82,6 +82,7 @@ class DefaultValues(object):
     FIRST_GROUP_IDX = 1000  # group idx initial value for group relaunch
     MAX_RELAUNCH_COUNT = 3  # maximum node relaunch count
     MAX_GROUP_RELAUNCH_COUNT = 3  # maximum node group relaunch count
+    TRAINING_ELASTIC_MODE = "base"
 
 
 class Context(Singleton):
@@ -146,6 +147,7 @@ def __init__(self):
         self.pre_check_operators = DefaultValues.PRE_CHECK_OPS
         self.max_relaunch_count = DefaultValues.MAX_RELAUNCH_COUNT
         self.max_group_relaunch_count = DefaultValues.MAX_GROUP_RELAUNCH_COUNT
+        self.training_elastic_mode = DefaultValues.TRAINING_ELASTIC_MODE
 
     def set_params_from_brain(self):
         self.train_speed_record_num = self.get_param_value_from_brain(
 
@@ -515,6 +515,10 @@ def report_event(
         )
         self._report(message)
 
+    def set_rdzv_blocked(self, blocked, reason=""):
+        message = comm.RdzvBlocked(blocked=blocked, reason=reason)
+        self._report(message)
+
     @classmethod
     def singleton_instance(cls, *args, **kwargs):
         if not cls._instance:
 
@@ -423,19 +423,22 @@ def __init__(
         local_shard_num=1,
         global_shard_num=1,
         save_timeout=CheckpointConstant.SAVE_TIMEOUT,
+        rank=0,
     ) -> None:
         logger.info(
             "Initializing the AsyncSaver with arguments: "
             f"checkpoint_dir={checkpoint_dir}, "
             f"local_shard_num={local_shard_num}, "
             f"global_shard_num={global_shard_num}, "
             f"save_timeout={save_timeout}"
+            f"rank={rank}"
         )
         self.checkpoint_dir = checkpoint_dir
         self.local_shard_num = local_shard_num
         self.global_shard_num = global_shard_num
-        self._node_rank = env_utils.get_node_rank()
-        self._is_agent_rank_0 = self._node_rank == 0
+        self._node_rank = env_utils.get_rank()
+        self._rank = rank
+        self._is_agent_rank_0 = self._rank == 0
         self._shm_handlers: List[SharedMemoryHandler] = []
         self._shm_locks: List[SharedLock] = []
         self._stop_commit = False
@@ -508,8 +511,15 @@ def _factory():
                         and saver_thread
                         and saver_thread.is_alive()
                     ):
-                        logger.info(
-                            "The saver is already created, skip creating the saver."
+                        # bufix: because the checkpoint dir changes every time of training
+                        cls._saver_instance.checkpoint_dir = (
+                            class_meta.kwargs.get("checkpoint_dir")
+                        )
+                        cls._saver_instance._rank = class_meta.kwargs.get(
+                            "rank"
+                        )
+                        cls._saver_instance._is_agent_rank_0 = (
+                            cls._saver_instance._rank == 0
                         )
                         continue
 
@@ -529,6 +539,33 @@ def _factory():
     def get_ckpt_saver(cls):
         return cls._saver_instance
 
+    def ucp(self, input_dir: str, output_dir: str, ucp_device_type: str):
+        """universal checkpoint"""
+        pass
+
+    def get_latest_start_saving_step(self):
+        step = self._latest_step
+        steps = []
+        for shm_handler in self._shm_handlers:
+            default_config = CheckpointConfig()
+            config = shm_handler.get_checkpoint_config(default_config)
+            steps.append(config.step)
+        if len(steps) > 0:
+            step = steps[0]
+        return step
+
+    def get_latest_success_save_dir(self):
+        try:
+            tracker_file = os.path.join(
+                self.checkpoint_dir, CheckpointConstant.TRACER_FILE_NAME
+            )
+            with open(tracker_file, "r") as f:
+                step = int(f.read())
+        except FileNotFoundError:
+            return None, None
+
+        return self.checkpoint_dir, step
+
     @classmethod
     def register_signal_handler(cls):
         sigint_handler = signal.getsignal(signal.SIGINT)
@@ -709,7 +746,7 @@ def _save_shard(
             shm_lock.release()
 
     def _dist_make_dir(self, path, timeout=30):
-        if self._node_rank == 0:
+        if self._rank == 0:
             logger.info(f"Create path by rank0 worker: {path}.")
             self.storage.safe_rmtree(path)
             self.storage.safe_makedirs(path)
@@ -719,7 +756,7 @@ def _dist_make_dir(self, path, timeout=30):
                     return
                 time.sleep(1)
             logger.warning(
-                f"Worker {self._node_rank} can't find path {path} "
+                f"Worker {self._rank} can't find path {path} "
                 f"with timeout {timeout}."
             )
 
@@ -1042,11 +1079,47 @@ def commit_checkpoint(self, step: int, step_done_dir: str, timeout=600):
     def persist_to_storage(
         self, local_shard_id: int, ckpt_config: CheckpointConfig
     ):
+        if ckpt_config is None or not ckpt_config.paths:
+            logger.info(
+                "Skip persisting checkpoint because checkpoint config is missing."
+            )
+            return
         state_dict = self._shm_handlers[local_shard_id].load_state_dict()
+        safe_serialization = None
+        if "safe_serialization" in state_dict:
+            safe_serialization = state_dict.pop("safe_serialization")
         for state_name, sd in state_dict.items():
             if sd and state_name in ckpt_config.paths:
+                from transformers.utils import (
+                    ADAPTER_SAFE_WEIGHTS_NAME,
+                    ADAPTER_WEIGHTS_NAME,
+                    SAFE_WEIGHTS_NAME,
+                    WEIGHTS_NAME,
+                )
+                from safetensors.torch import save_file as safe_save_file
+                import re
+
                 path = ckpt_config.paths[state_name]
-                self.storage.write_state_dict(sd, path, torch.save)
+                reg = re.compile(r"(.*?)-\d{5}-of-\d{5}.bin")
+                match = reg.fullmatch(state_name)
+                if safe_serialization:
+                    if state_name.endswith(ADAPTER_WEIGHTS_NAME):
+                        path = path.replace(
+                            ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME
+                        )
+                        self.storage.write_state_dict(sd, path, safe_save_file)
+                    elif state_name.endswith(WEIGHTS_NAME):
+                        path = path.replace(WEIGHTS_NAME, SAFE_WEIGHTS_NAME)
+                        self.storage.write_state_dict(sd, path, safe_save_file)
+                    elif match:
+                        path = path.replace("pytorch_model", "model").replace(
+                            ".bin", ".safetensors"
+                        )
+                        self.storage.write_state_dict(sd, path, safe_save_file)
+                    else:
+                        self.storage.write_state_dict(sd, path, torch.save)
+                else:
+                    self.storage.write_state_dict(sd, path, torch.save)
 
 
 class TempDirCheckpointSaver(AsyncCheckpointSaver):
@@ -1063,13 +1136,15 @@ def __init__(
         local_shard_num=1,
         global_shard_num=1,
         save_timeout=CheckpointConstant.SAVE_TIMEOUT,
+        rank=0,
     ) -> None:
         super().__init__(
             checkpoint_dir,
             storage_meta,
             local_shard_num,
             global_shard_num,
             save_timeout,
+            rank=rank,
         )
 
         if self._node_rank == 0:
@@ -1267,9 +1342,6 @@ class DdpCheckpointSaver(CommonDirCheckpointSaver):
     """Persist the checkpoint from CPU memory buffer into the storage."""
 
     def persist_to_storage(self, local_shard_id: int, ckpt_config):
-        if self._node_rank != 0:
-            logger.info("Skip and only rank 0 saves checkpoint in a DDP job.")
-            return
         super().persist_to_storage(local_shard_id, ckpt_config)
 
 
@@ -1310,6 +1382,70 @@ def update_tracker_file(self, step):
         )
         self.storage.write(str(step), ds_tracker_filename)
 
+    def get_deepspeed_install_dir(self):
+        spec = importlib.util.find_spec("deepspeed")
+        deepspeed_dir = ""
+        if spec and spec.origin:
+            # spec.origin 指向 __init__.py 文件
+            module_path = spec.origin
+            deepspeed_dir = os.path.dirname(module_path)
+        return deepspeed_dir
+
+    def ucp(self, input_dir: str, output_dir: str, ucp_device_type: str):
+        import torch
+        from packaging import version
+        from torch.distributed.elastic.multiprocessing.api import (
+            SubprocessHandler,
+        )
+
+        def version_less_than_230():
+            current_version = version.parse(torch.__version__).base_version
+            return version.parse(current_version) <= version.parse("2.2.2")
+
+        def version_less_than_240():
+            current_version = version.parse(torch.__version__).base_version
+            return version.parse(current_version) <= version.parse("2.3.1")
+
+        import sys
+        import os
+
+        cmd = os.getenv("PYTHON_EXEC", sys.executable)
+        deepspeed_dir = self.get_deepspeed_install_dir()
+        args_list = [
+            deepspeed_dir + "/checkpoint/ds_to_universal.py",
+            "--input_folder",
+            f"{input_dir}",
+            "--output_folder",
+            f"{output_dir}",
+            "--inject_missing_state",
+        ]
+        if ucp_device_type != "cpu":
+            args_list.extend(["--device", ucp_device_type])
+        args = tuple(args_list)
+        if version_less_than_230():
+            handler = SubprocessHandler(cmd, args, {}, "", "")
+        else:
+            handler = SubprocessHandler(cmd, args, {}, "", "", 0)
+        ret = handler.proc.wait()
+        if ret != 0:
+            print(f"subprocess returned non-zero exit code{ret}")
+            return False
+        else:
+            return True
+
+    def get_latest_success_save_dir(self):
+        try:
+            tracker_file = os.path.join(
+                self.checkpoint_dir, CheckpointConstant.TRACER_FILE_NAME
+            )
+
+            with open(tracker_file, "r") as f:
+                step = int(f.read())
+        except FileNotFoundError:
+            return None, None
+
+        return self.checkpoint_dir, step
+
 
 class FsdpDcpSaver(CommonDirCheckpointSaver):
     """The saver saves the distributed checkpoint of FSDP into the storage."""
@@ -1355,3 +1491,4 @@ def persist_to_storage(
                 self.checkpoint_dir, CheckpointConstant.TRACER_FILE_NAME
             )
             self.storage.write(str(ckpt_config.step), tracer_file)
+            dcp_metadata = meta_dict.get("dcp_metadata", {})