Repository: intelligent-machine-learning/dlrover (language: python)

Issue:
The agent unlinks the shared memory with signals.



Discussion / hints:

Fix patch:
diff --git a/dlrover/python/common/multi_process.py b/dlrover/python/common/multi_process.py
index f816c6aef..c14d91b12 100644
--- a/dlrover/python/common/multi_process.py
+++ b/dlrover/python/common/multi_process.py
@@ -60,13 +60,13 @@ def _create_socket_client(path):
     """
     client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     connected = False
-    for _ in range(3):
+    for _ in range(30):
         try:
             client.connect(path)
             connected = True
             break
         except FileNotFoundError:
-            time.sleep(1)
+            time.sleep(0.1)
     if not connected:
         client.connect(path)
     return client
diff --git a/dlrover/python/elastic_agent/torch/ckpt_saver.py b/dlrover/python/elastic_agent/torch/ckpt_saver.py
index 6e2c2b13f..d0b673a0d 100644
--- a/dlrover/python/elastic_agent/torch/ckpt_saver.py
+++ b/dlrover/python/elastic_agent/torch/ckpt_saver.py
@@ -13,14 +13,18 @@
 
 import os
 import shutil
+import signal
 import sys
 import threading
+import time
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Callable, List, Mapping, Tuple
 
 import numpy as np
 import torch
+import torch.distributed as dist
 
 from dlrover.python.common.log import default_logger as logger
 from dlrover.python.common.multi_process import (
@@ -30,12 +34,24 @@
     SharedQueue,
 )
 
-CKPT_DIR_PREFIX = "checkpoint-"
+_CKPT_DIR_PREFIX = "checkpoint-"
 
-SAVE_STEP_QNAME_PREFIX = "checkpoint_lock_rank_"
-CKPT_META_NAME_PREFIX = "checkpoint_meta_local_rank_"
-TENSOR_SHM_NAME_PREFIX = "checkpoint_shm_local_rank_"
-SHM_LOCK_NAME_PREFIX = "shm_local_rank_"
+_SAVE_STEP_QNAME_PREFIX = "checkpoint_lock_rank_"
+_CKPT_META_NAME_PREFIX = "checkpoint_meta_local_rank_"
+_TENSOR_SHM_NAME_PREFIX = "checkpoint_shm_local_rank_"
+_SHM_LOCK_NAME_PREFIX = "shm_local_rank_"
+_WIRTING_SHM = "__WRITING_SHM__"
+
+
+def timer(func):
+    def wrapper(*args, **kwargs):
+        start = time.time()
+        result = func(*args, **kwargs)
+        t = round(time.time() - start, 3)
+        logger.info(f"Function {func.__name__} cost {t}s")
+        return result
+
+    return wrapper
 
 
 def _init_dir(dir):
@@ -44,7 +60,8 @@ def _init_dir(dir):
     os.makedirs(dir)
 
 
-def convert_torch_dtype_to_numpy(torch_dtype):
+def _convert_torch_dtype_to_numpy(torch_dtype):
+    """Conver the torch dtype to numpy dtype."""
     dtype_map = {
         torch.float32: np.float32,
         torch.float: np.float32,
@@ -64,33 +81,33 @@ def convert_torch_dtype_to_numpy(torch_dtype):
     return dtype_map[torch_dtype]
 
 
-def traverse_state_dict(value: object, visitor: Callable[[object], None]):
+def _traverse_state_dict(value: object, visitor: Callable[[object], None]):
     """
     Invoke ``visitor`` for each value recursively in ``state_dict``.
     """
     if isinstance(value, Mapping):
         temp_dict = {}
         for k, v in value.items():
-            temp_dict[k] = traverse_state_dict(v, visitor)
+            temp_dict[k] = _traverse_state_dict(v, visitor)
         return temp_dict
     elif isinstance(value, List):
         temp_list = []
         for _, v in enumerate(value):
-            temp_list.append(traverse_state_dict(v, visitor))
+            temp_list.append(_traverse_state_dict(v, visitor))
         return temp_list
     else:
         return visitor(value)
 
 
-def read_state_dict_from_shm(checkpoint_meta, tensor_shm):
-    state_dict = traverse_state_dict(
+def _read_state_dict_from_shm(checkpoint_meta, tensor_shm):
+    state_dict = _traverse_state_dict(
         checkpoint_meta,
-        lambda x: read_tensor_from_buf(x, tensor_shm),
+        lambda x: _read_tensor_from_buf(x, tensor_shm),
     )
     return state_dict
 
 
-def read_tensor_from_buf(value, shm_tensor_buffer):
+def _read_tensor_from_buf(value, shm_tensor_buffer):
     """
     Read a tensor from the buffer of shared memory.
     """
@@ -107,6 +124,43 @@ def read_tensor_from_buf(value, shm_tensor_buffer):
         return value
 
 
+def _create_shared_memory(name, create, size=0):
+    """
+    Create a shared memory.
+    """
+    if not create:
+        try:
+            return SharedMemory(name=name)
+        except FileNotFoundError:
+            return None
+    try:
+        shm = SharedMemory(
+            name=name,
+            create=create,
+            size=size,
+        )
+    except FileExistsError:
+        shm = SharedMemory(name=name)
+    return shm
+
+
+def _get_latest_checkpoint(checkpoint_dir):
+    """Get the checkpoint directory with the maximum step."""
+    if not os.path.exists(checkpoint_dir):
+        return ""
+    max_step = 0
+    for fn in os.listdir(checkpoint_dir):
+        if not fn.startswith(_CKPT_DIR_PREFIX):
+            continue
+        step = int(fn.split("-")[-1])
+        max_step = step if step > max_step else max_step
+    if max_step > 0:
+        path = os.path.join(checkpoint_dir, f"{_CKPT_DIR_PREFIX}{max_step}")
+    else:
+        path = ""
+    return path
+
+
 @dataclass
 class TensorMeta(object):
     shape: Tuple[int] = None  # type: ignore
@@ -126,29 +180,63 @@ class SaverFactory(object):
 
 
 class CheckpointSaver(metaclass=ABCMeta):
+    """
+    CheckpointSaver saves the state dict from the shared memory into
+    the storage.
+
+    Attributes:
+        checkpoint_dir (str): the directory to save the checkpointing state
+            dict to the storage if the training process fails.
+        num_proc (int): the number of training process, i.e. local world size.
+    """
+
+    _saver_instance = None
+
+    def __init__(self, checkpoint_dir, num_proc=1):
+        self.checkpoint_dir = checkpoint_dir
+        self.num_proc = num_proc
+
     @abstractmethod
-    def _save_shm_to_storage(self):
+    def _sync_shm_to_storage(self):
         pass
 
     @classmethod
-    def start_async_saving_ckpt(cls):
+    def start_async_saving_ckpt(cls, num_proc=1):
         """
         Start a thread to asynchronously save the checkpoint state dict
         from the shared memory into the storage. Firstly, it waits that
         the training process notify the saver class to create a saver.
+
+        Args:
+            num_proc: the number of training process, i.e. local world size.
         """
         sq = SharedQueue(name="factory", create=True)
 
-        def _save(sq: SharedQueue):
-            class_name = sq.get()
+        def _save():
+            class_name, args = sq.get()
             class_def = getattr(sys.modules[__name__], class_name)
-            saver: CheckpointSaver = class_def()
-            saver._save_shm_to_storage()
+            if cls._saver_instance is None:
+                args["num_proc"] = num_proc
+                saver: CheckpointSaver = class_def(**args)
+                cls._saver_instance = saver
+            cls._saver_instance._sync_shm_to_storage()
 
         threading.Thread(
-            target=_save, args=(sq,), name="checkpoint-saver", daemon=True
+            target=_save, name="checkpoint-saver", daemon=True
         ).start()
 
+    @classmethod
+    def get_ckpt_saver(cls):
+        return cls._saver_instance
+
+    @abstractmethod
+    def close(self):
+        pass
+
+    @abstractmethod
+    def save_shm_to_storage(self):
+        pass
+
 
 class NoShardingSaver(CheckpointSaver):
     """
@@ -156,17 +244,17 @@ class NoShardingSaver(CheckpointSaver):
     from the shared memory created by local rank 0 to the storage.
     """
 
-    def __init__(self) -> None:
-        self._checkpoint_dir = ""
+    def __init__(self, checkpoint_dir, num_proc=1) -> None:
+        super().__init__(checkpoint_dir, num_proc)
         self._tensor_shm = None
         # Only local rank 0 save the state dict to memory in DDP.
-        qname = SAVE_STEP_QNAME_PREFIX + str(0)
+        qname = _SAVE_STEP_QNAME_PREFIX + str(0)
         self._to_save_queue = SharedQueue(name=qname, create=True)
-        meta_name = CKPT_META_NAME_PREFIX + str(0)
+        meta_name = _CKPT_META_NAME_PREFIX + str(0)
         self._shared_ckpt_meta = SharedDict(name=meta_name, create=True)
-        lock_name = SHM_LOCK_NAME_PREFIX + str(0)
+        lock_name = _SHM_LOCK_NAME_PREFIX + str(0)
         self._shm_lock = SharedLock(name=lock_name, create=True)
-        self._shm_name = TENSOR_SHM_NAME_PREFIX + str(0)
+        self._shm_name = _TENSOR_SHM_NAME_PREFIX + str(0)
 
     def __del__(self):
         self.close()
@@ -179,7 +267,7 @@ def close(self):
         self._shared_ckpt_meta.close()
         self._shm_lock.close()
 
-    def _save_shm_to_storage(self):
+    def _sync_shm_to_storage(self):
         """
         The loop to persist the state dict from the memory
         buffer into the storage.
@@ -195,7 +283,7 @@ def _save_shm_to_storage(self):
                 f"into the storage {path}."
             )
             meta_dict = self._shared_ckpt_meta.get()
-            state_dict = read_state_dict_from_shm(meta_dict, self._tensor_shm)
+            state_dict = _read_state_dict_from_shm(meta_dict, self._tensor_shm)
             self._persist_to_storage(state_dict, path)
             self._shm_lock.release()
 
@@ -204,3 +292,439 @@ def _persist_to_storage(self, state_dict, path):
         checkpoint_dir = os.path.dirname(path)
         _init_dir(checkpoint_dir)
         torch.save(state_dict, path)
+
+    def save_shm_to_storage(self):
+        """
+        Save the state dict in the shared memory into the storage. The agent
+        can call the method to save the state dict into the storage if the
+        training process fails or the agent wants to restart training
+        processes.
+        """
+        if self._tensor_shm is None:
+            return
+        acquired = self._shm_lock.acquire()
+        if not acquired:
+            # The training process does not release the lock because it fails
+            # when writing the state dict into the shared memory. The shared
+            # memory may be dirty and the saver cannot save it to the storage.
+            return
+        meta_dict = self._shared_ckpt_meta.get()
+        step = meta_dict["step"]
+        path = os.path.join(
+            self.checkpoint_dir, f"checkpoint-{step}/checkpoint.pt"
+        )
+        state_dict = _read_state_dict_from_shm(meta_dict, self._tensor_shm)
+        self._persist_to_storage(state_dict, path)
+        self._shm_lock.release()
+        logger.info(
+            "Save the checkpointing state dict from the shared "
+            f"memory to {path}."
+        )
+
+
+class CheckpointEngine(metaclass=ABCMeta):
+    """
+    The checkpoint engine synchronously writes the state dict into
+    the shared memory and notify the agent in main process to
+    asynchronously save the state dict from the shared memory into
+    the storage. Writing to memory is significantly quicker
+    than writing to storage. The engine only blocks the training
+    with a little time. Users can frequently call `save_to_memory` in
+    the training loop and call `save_to_storage`.
+
+    If the training process fail, the agent in main process can continuely
+    saves the the state dict from the shared memory into the storage.
+
+    Attributes:
+        checkpoint_dir (str):  the directory to save the temp checkpoint
+            if the training process fails.
+
+    Examples::
+        >>> engine = NoShardingCheckpointEngine(
+        >>>     checkpoint_dir="/tmp/checkpoint/"
+        >>> )
+        >>> for step, data in enumerate(dataloader):
+        >>>     ...
+        >>>     state_dict = model.state_dict()
+        >>>     if step % 5 == 0:
+        >>>         engine.save_to_memory(state_dict, step)
+        >>>     elif step % 100 == 0:
+        >>>         path = f"/tmp/checkpoint/ckpt-{step}.pt"
+        >>>         engine.save_to_storage(state_dict, path, step)
+        >>> sate_dict = engine.load()
+    """
+
+    def __init__(self, checkpoint_dir):
+        self.checkpoint_dir = checkpoint_dir
+        if dist.is_initialized():
+            self._rank = dist.get_rank()
+            self._local_rank = int(os.environ["LOCAL_RANK"])
+            self._saver_group = dist.new_group(
+                backend="gloo", timeout=timedelta(seconds=30)
+            )
+        else:
+            self._rank = 0
+            self._local_rank = int(os.getenv("LOCAL_RANK", 0))
+            self._saver_group = None
+
+        self._buffer_size = 0
+        self._cached_step = 0
+        self._meta_dict = dict()
+        self._shm_name = ""
+        self._tensor_shm: SharedMemory = None
+        self._shared_ckpt_meta: SharedDict = None
+        self._shm_lock: SharedLock = None
+        self._to_save_queue: SharedQueue = None
+        self._notify_agent_to_create_saver()
+        self._init_shared_objs()
+
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        if self._tensor_shm:
+            self._tensor_shm.close()
+
+    @abstractmethod
+    def _init_shared_objs(self):
+        """
+        Initialize the shared queue, lock and memory to communiate
+        with the agent in the main process.
+        """
+        pass
+
+    @abstractmethod
+    def _notify_agent_to_create_saver(self):
+        """
+        Notify the agent in the main process to create a checkpointing
+        saver to save the state dict from the shared memory into the storage.
+        """
+        pass
+
+    @timer
+    def save_to_memory(self, state_dict, step):
+        """
+        Synchonously Saves the state dict into the shared memory with the main
+        process. If the agent in the main process is saving the shared memory
+        into the storage, the method will skip to write the shared memory.
+
+        Args:
+            state_dict (dict): the state dict of model and optimizer to save.
+            step (int): the iteration step.
+        """
+        if "step" not in state_dict:
+            state_dict["step"] = step
+        if _WIRTING_SHM in state_dict:
+            raise ValueError(f"state_dict cannot have the key {_WIRTING_SHM}.")
+
+        if self._tensor_shm is None:
+            self._make_state_dict_buffer(state_dict)
+        acquired = self._shm_lock.acquire(blocking=False)
+        all_rank_ready = self._check_all_rank_ready(acquired)
+        if not all_rank_ready:
+            logger.info(
+                f"Rank {self._rank} skips the save the checkpoint "
+                f"in CPU memory since it is saving the latest "
+                "checkpoint from the CPU memory into the storage."
+            )
+            if acquired:
+                self._shm_lock.release()
+            return
+        self._copy_state_dict_to_shm(state_dict)
+
+        if acquired:
+            self._shm_lock.release()
+        self._cached_step = step
+
+    def _create_tensor_meta(self, value: torch.Tensor):
+        """
+        Create a tensor meta of a tensor and compute the total
+        size of the state dict.
+        """
+        if not torch.is_tensor(value):
+            return value
+        dtype = _convert_torch_dtype_to_numpy(value.dtype)
+        meta = TensorMeta(
+            shape=tuple(value.shape),  # type: ignore
+            dtype=dtype,
+            element_size=value.element_size(),
+            numel=value.numel(),
+            offset=self._buffer_size,
+        )
+        self._buffer_size += value.numel() * value.element_size()
+        return meta
+
+    def _make_state_dict_buffer(self, state_dict):
+        """
+        Make the shared memory to store the state dict.
+        """
+        self._meta_dict = _traverse_state_dict(
+            state_dict, self._create_tensor_meta
+        )
+
+        # Update the meta dict in the main process.
+        self._shared_ckpt_meta.update(self._meta_dict)
+        self._tensor_shm = _create_shared_memory(
+            name=self._shm_name,
+            create=True,
+            size=self._buffer_size,
+        )
+
+    def _copy_state_dict_to_shm(self, state_dict):
+        """
+        Copy the state dict from CPU memory buffer into the shared memory.
+        """
+
+        def _tarverse_copy(value, meta):
+            if isinstance(value, Mapping):
+                for k, v in value.items():
+                    if isinstance(v, (Mapping, List)):
+                        m = meta[k]
+                        _tarverse_copy(v, m)
+                    elif torch.is_tensor(v):
+                        m = meta[k]
+                        self._write_shared_memory(v, m)
+                    else:
+                        meta[k] = v
+            elif isinstance(value, List):
+                for i, v in enumerate(value):
+                    if isinstance(v, (Mapping, List)):
+                        m = meta[i]
+                        _tarverse_copy(v, m)
+                    elif torch.is_tensor(v):
+                        m = meta[i]
+                        self._write_shared_memory(v, m)
+                    else:
+                        meta[i] = v
+
+        self._meta_dict[_WIRTING_SHM] = True
+        self._shared_ckpt_meta.update(self._meta_dict)
+        _tarverse_copy(state_dict, self._meta_dict)
+        # Update the meta dict in the main process.
+        self._meta_dict[_WIRTING_SHM] = False
+        self._shared_ckpt_meta.update(self._meta_dict)
+
+    def _write_shared_memory(self, value, meta: TensorMeta):
+        """
+        Write a CPU tensor into the shared memory.
+        """
+        data_array = value.cpu().numpy()
+        write_array = np.ndarray(
+            data_array.shape,
+            dtype=data_array.dtype,
+            buffer=self._tensor_shm.buf,
+            offset=meta.offset,
+        )
+        if data_array.shape == ():
+            write_array.fill(data_array)
+        else:
+            write_array[:] = data_array[:]
+
+    def _check_all_rank_ready(self, ready):
+        """
+        Check wether all ranks are ready.
+        """
+        if not self._saver_group:
+            return ready
+        value = 0 if ready else 1
+        t = torch.tensor([value], dtype=torch.int64)
+        dist.all_reduce(t, group=self._saver_group)
+        return t == 0
+
+    @timer
+    def save_to_storage(self, state_dict, path, step):
+        """
+        Asynchonously saves the state dict into the storage. It synchonously
+        saves the state dict into the shared memory and put the path
+        into a shared queue. The agent in the main process waits for the queue
+        for save the state dict in the shared memory into the storage.
+
+        Args:
+            state_dict (dict): the state dict of model and optimizer to save.
+            path (str): optional, the file path to save the checkpoint. If the
+                path is not defined, the engine will save the state dict into
+                the shared memory not the storage.
+            step (int): the iteration step.
+        """
+        if step > self._cached_step:
+            self.save_to_memory(state_dict, step)
+        if path:
+            self._to_save_queue.put(path)
+
+    def load(self, resume_path=""):
+        """
+        The method firstly try to load the state dict from the shared memory.
+        If there is no state dict in the shared memory, the method will
+        load the state dict from the storage.
+
+        Returns:
+            A dict.
+        """
+        state_dict = self._load_from_shared_memory()
+        if state_dict:
+            return state_dict
+        state_dict = self._load_from_storage(resume_path)
+        return state_dict
+
+    def _load_from_shared_memory(self):
+        """
+        Load the state dict from the shared memory.
+
+        Returns:
+            A dict.
+        """
+        if self._tensor_shm is None:
+            self._tensor_shm = _create_shared_memory(
+                self._shm_name,
+                create=False,
+            )
+        if not self._tensor_shm:
+            return None
+        meta_dict = self._shared_ckpt_meta.get()
+        if meta_dict.get(_WIRTING_SHM, False):
+            return None
+        state_dict = _read_state_dict_from_shm(meta_dict, self._tensor_shm)
+        return state_dict
+
+    def _load_from_storage(self, resume_path=""):
+        """
+        Load the state dict from the CPU memory if the state dict is complete
+        in CPU memory. Otherwise, the function will load the state dict from
+        the storage.
+
+        Args:
+            resume_path (str, optional): , If the resume_path is an empty
+                string, the function will load the latest checkpoint file in
+                the checkpoint directory.
+
+        Returns:
+            A dict:
+                a dictionary containing a whole state of the modules in the
+                checkpointing file.
+        """
+        if resume_path:
+            state_dict = torch.load(resume_path)
+        else:
+            state_dict = self._load_from_historic_checkpoint()
+        return state_dict
+
+    def _load_from_historic_checkpoint(self):
+        """Locd checkpoint from the lastest complete checkpoint."""
+        while True:
+            latest_ckpt_dir = _get_latest_checkpoint(self.checkpoint_dir)
+            if not latest_ckpt_dir:
+                return {}
+
+            resume_path = os.path.join(latest_ckpt_dir, "checkpoint.pt")
+            if not os.path.exists(resume_path):
+                shutil.rmtree(latest_ckpt_dir)
+                continue
+            try:
+                state_dict = torch.load(resume_path)
+                logger.info(f"Load checkpoint from {resume_path}")
+                return state_dict
+            except Exception:
+                logger.warning(
+                    f"Fail to load checkpoint from {resume_path}."
+                    " Roll back to the last checkpoint file."
+                )
+                shutil.rmtree(latest_ckpt_dir)
+
+
+class ShardingCheckpointEngine(CheckpointEngine):
+    """
+    The engine to save the sharding model and optimizer state dict
+    into the memory and storage. We can use it to save the model and optimizer
+    using FSDP, Zero-3 or Megatron-LM.
+    """
+
+    def __init__(self, checkpoint_dir):
+        super().__init__(checkpoint_dir)
+
+    def _notify_agent_to_create_saver(self):
+        # TODO: implement the saver in the agent to support saving
+        # sharding state dict.
+        pass
+
+    def _init_shared_objs(self):
+        meta_name = _CKPT_META_NAME_PREFIX + str(self._local_rank)
+        self._shared_ckpt_meta = SharedDict(name=meta_name, create=False)
+        lock_name = _SHM_LOCK_NAME_PREFIX + str(self._local_rank)
+        self._shm_lock = SharedLock(name=lock_name, create=False)
+        qname = _SAVE_STEP_QNAME_PREFIX + str(self._local_rank)
+        self._to_save_queue = SharedQueue(name=qname, create=False)
+        self._shm_name = _TENSOR_SHM_NAME_PREFIX + str(self._local_rank)
+
+
+class NoShardingCheckpointEngine(CheckpointEngine):
+    """
+    The engine saves the model and optimizer state dict without sharding
+    in a local or DDP job.
+    """
+
+    def __init__(self, checkpoint_dir):
+        super().__init__(checkpoint_dir)
+
+    def _notify_agent_to_create_saver(self):
+        queue = SharedQueue(name="factory")
+        args = {"checkpoint_dir": self.checkpoint_dir}
+        queue.put(("NoShardingSaver", args))
+        queue.close()
+
+    def _init_shared_objs(self):
+        """
+        Initialize the shared object with the main process.
+        Without model sharding, all ranks share the same shared memory
+        created by the local rank 0 on a node.
+        """
+        meta_name = _CKPT_META_NAME_PREFIX + str(0)
+        self._shared_ckpt_meta = SharedDict(name=meta_name, create=False)
+        lock_name = _SHM_LOCK_NAME_PREFIX + str(0)
+        self._shm_lock = SharedLock(name=lock_name, create=False)
+        self._shm_lock.release()
+        qname = _SAVE_STEP_QNAME_PREFIX + str(0)
+        self._to_save_queue = SharedQueue(name=qname, create=False)
+        self._shm_name = _TENSOR_SHM_NAME_PREFIX + str(0)
+
+    @timer
+    def save_to_memory(self, state_dict, step):
+        """
+        Synchonously Saves the state dict into the shared memory with the main
+        process. If the agent in the main process is saving the shared memory
+        into the storage, the method will skip to write the shared memory.
+
+        Args:
+            state_dict (dict): the state dict of model and optimizer to save.
+            step (int): the iteration step.
+        """
+        if self._local_rank == 0:
+            super().save_to_memory(state_dict, step)
+
+    @timer
+    def save_to_storage(self, state_dict, path, step):
+        """
+        Asynchonously saves the state dict into the storage. It synchonously
+        saves the state dict into the shared memory and put the path
+        into a shared queue. The agent in the main process waits for the queue
+        for save the state dict in the shared memory into the storage.
+
+        Args:
+            state_dict (dict): the state dict of model and optimizer to save.
+            step (int): the iteration step.
+            path (str): optional, the file path to save the checkpoint. If the
+                path is not defined, the engine will save the state dict into
+                the shared memory not the storage.
+        """
+        if self._rank == 0:
+            super().save_to_storage(state_dict, path, step)
+
+
+def _clean_shm_handler(signum, frame):
+    """Clean the shared memory from ^C and "killall python" etc."""
+    saver: CheckpointSaver = CheckpointSaver.get_ckpt_saver()
+    if saver:
+        saver.close()
+
+
+signal.signal(signal.SIGINT, _clean_shm_handler)
+signal.signal(signal.SIGTERM, _clean_shm_handler)
diff --git a/dlrover/python/elastic_agent/torch/training.py b/dlrover/python/elastic_agent/torch/training.py
index 8723b8a21..66136bdac 100644
--- a/dlrover/python/elastic_agent/torch/training.py
+++ b/dlrover/python/elastic_agent/torch/training.py
@@ -499,7 +499,10 @@ def _initialize_workers(self, worker_group):
                 break
 
     def _invoke_run(self, role: str = DEFAULT_ROLE) -> RunResult:
-        # NOTE: currently only works for a single role
+        # Start a thread to save the checkpointing state dict from
+        # the shared memory to the storage.
+        proc_num = self._worker_group.spec.local_world_size
+        CheckpointSaver.start_async_saving_ckpt(proc_num)
 
         spec = self._worker_group.spec
         role = spec.role
@@ -588,8 +591,18 @@ def _report_failure_to_master(self, failures: Dict[int, ProcessFailure]):
     def _restart_workers(self, worker_group: WorkerGroup):
         self._restart_count += 1
         self._remaining_restarts -= 1
+        self._save_ckpt_to_storage()
         super()._restart_workers(worker_group)
 
+    def _save_ckpt_to_storage(self):
+        """
+        The agent can save the checkpointing state dict in the shared
+        memory into the storage before restarting training processes.
+        """
+        saver: CheckpointSaver = CheckpointSaver.get_ckpt_saver()
+        if saver:
+            saver.save_shm_to_storage()
+
     def _membership_changed(self, role, rdzv_handler: RendezvousHandler):
         # Timeout may happen when to query TCPStore.
         if self._config.network_check:
@@ -682,8 +695,6 @@ def launch_agent(
         log_dir=config.log_dir,
     )
 
-    CheckpointSaver.start_async_saving_ckpt()
-
     shutdown_rdzv = True
     try:
         metrics.initialize_metrics(metrics.MetricsConfig(config.metrics_cfg))
diff --git a/dlrover/trainer/torch/elastic/checkpoint.py b/dlrover/trainer/torch/elastic/checkpoint.py
index b803c5c0a..1dea49eb2 100644
--- a/dlrover/trainer/torch/elastic/checkpoint.py
+++ b/dlrover/trainer/torch/elastic/checkpoint.py
@@ -13,13 +13,8 @@
 
 import os
 import shutil
-import time
 from abc import ABCMeta, abstractmethod
-from datetime import timedelta
-from typing import List, Mapping
 
-import numpy as np
-import torch
 import torch.distributed as dist
 from torch.distributed.fsdp import FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
@@ -28,80 +23,20 @@
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from dlrover.python.common.log import default_logger as logger
-from dlrover.python.common.multi_process import (
-    SharedDict,
-    SharedLock,
-    SharedMemory,
-    SharedQueue,
-)
 from dlrover.python.elastic_agent.torch.ckpt_saver import (
-    CKPT_META_NAME_PREFIX,
-    SAVE_STEP_QNAME_PREFIX,
-    SHM_LOCK_NAME_PREFIX,
-    TENSOR_SHM_NAME_PREFIX,
-    TensorMeta,
-    convert_torch_dtype_to_numpy,
-    read_state_dict_from_shm,
-    traverse_state_dict,
+    CheckpointEngine,
+    NoShardingCheckpointEngine,
 )
 from dlrover.trainer.torch.elastic.sampler import ElasticDistributedSampler
 
 CKPT_DIR_PREFIX = "checkpoint-"
 
 
-def timer(func):
-    def wrapper(*args, **kwargs):
-        start = time.time()
-        result = func(*args, **kwargs)
-        t = round(time.time() - start, 3)
-        logger.info(f"Function {func.__name__} cost {t}s")
-        return result
-
-    return wrapper
-
-
 def _sync():
     if dist.is_initialized():
         dist.barrier()
 
 
-def _get_latest_checkpoint(checkpoint_dir):
-    """Get the checkpoint directory with the maximum step."""
-    if not os.path.exists(checkpoint_dir):
-        return ""
-    max_step = 0
-    for fn in os.listdir(checkpoint_dir):
-        if not fn.startswith(CKPT_DIR_PREFIX):
-            continue
-        step = int(fn.split("-")[-1])
-        max_step = step if step > max_step else max_step
-    if max_step > 0:
-        path = os.path.join(checkpoint_dir, f"{CKPT_DIR_PREFIX}{max_step}")
-    else:
-        path = ""
-    return path
-
-
-def _create_shared_memory(name, create, size=0):
-    """
-    Create a shared memory.
-    """
-    if not create:
-        try:
-            return SharedMemory(name=name)
-        except FileNotFoundError:
-            return None
-    try:
-        shm = SharedMemory(
-            name=name,
-            create=create,
-            size=size,
-        )
-    except FileExistsError:
-        shm = SharedMemory(name=name)
-    return shm
-
-
 def _keep_topk_checkpoint(checkpoint_dir, max_to_keep):
     """Keep top k checkpoints and remove other checkpoints.
 
@@ -125,398 +60,6 @@ def _keep_topk_checkpoint(checkpoint_dir, max_to_keep):
         shutil.rmtree(dir_name)
 
 
-class CheckpointEngine(metaclass=ABCMeta):
-    """
-    The checkpoint engine synchronously writes the state dict into
-    the shared memory and notify the agent in main process to
-    asynchronously save the state dict from the shared memory into
-    the storage. Writing to memory is significantly quicker
-    than writing to storage. The engine only blocks the training
-    with a little time. Users can frequently call `save_to_memory` in
-    the training loop and call `save_to_storage`.
-
-    If the training process fail, the agent in main process can continuely
-    saves the the state dict from the shared memory into the storage.
-
-    Attributes:
-        checkpoint_dir (str):  the directory to save the temp checkpoint
-            if the training process fails.
-
-    Examples::
-        >>> engine = NoShardingCheckpointEngine(
-        >>>     checkpoint_dir="/tmp/checkpoint/"
-        >>> )
-        >>> for step, data in enumerate(dataloader):
-        >>>     ...
-        >>>     state_dict = model.state_dict()
-        >>>     if step % 5 == 0:
-        >>>         engine.save_to_memory(state_dict, step)
-        >>>     elif step % 100 == 0:
-        >>>         path = f"/tmp/checkpoint/ckpt-{step}.pt"
-        >>>         engine.save_to_storage(state_dict, path, step)
-        >>> sate_dict = engine.load()
-    """
-
-    def __init__(self, checkpoint_dir):
-        self.checkpoint_dir = checkpoint_dir
-        if dist.is_initialized():
-            self._rank = dist.get_rank()
-            self._local_rank = int(os.environ["LOCAL_RANK"])
-            self._saver_group = dist.new_group(
-                backend="gloo", timeout=timedelta(seconds=30)
-            )
-        else:
-            self._rank = 0
-            self._local_rank = int(os.getenv("LOCAL_RANK", 0))
-            self._saver_group = None
-
-        self._buffer_size = 0
-        self._cached_step = 0
-        self._meta_dict = dict()
-        self._shm_name = ""
-        self._tensor_shm: SharedMemory = None
-        self._shared_ckpt_meta: SharedDict = None
-        self._shm_buffer_lock: SharedLock = None
-        self._to_save_queue: SharedQueue = None
-        self._notify_agent_to_create_saver()
-        self._init_shared_objs()
-
-    def __del__(self):
-        self.close()
-
-    def close(self):
-        if self._shared_ckpt_meta:
-            self._shared_ckpt_meta.close()
-        if self._shm_buffer_lock:
-            self._shm_buffer_lock.close()
-        if self._to_save_queue:
-            self._to_save_queue.close()
-        if self._tensor_shm:
-            self._tensor_shm.close()
-
-    @abstractmethod
-    def _init_shared_objs(self):
-        """
-        Initialize the shared queue, lock and memory to communiate
-        with the agent in the main process.
-        """
-        pass
-
-    @abstractmethod
-    def _notify_agent_to_create_saver(self):
-        """
-        Notify the agent in the main process to create a checkpointing
-        saver to save the state dict from the shared memory into the storage.
-        """
-        pass
-
-    def _create_tensor_meta(self, value: torch.Tensor):
-        """
-        Create a tensor meta of a tensor and compute the total
-        size of the state dict.
-        """
-        if not torch.is_tensor(value):
-            return value
-        dtype = convert_torch_dtype_to_numpy(value.dtype)
-        meta = TensorMeta(
-            shape=tuple(value.shape),  # type: ignore
-            dtype=dtype,
-            element_size=value.element_size(),
-            numel=value.numel(),
-            offset=self._buffer_size,
-        )
-        self._buffer_size += value.numel() * value.element_size()
-        return meta
-
-    def _make_state_dict_buffer(self, state_dict):
-        """
-        Make the shared memory to store the state dict.
-        """
-        self._meta_dict = traverse_state_dict(
-            state_dict, self._create_tensor_meta
-        )
-
-        # Update the meta dict in the main process.
-        self._shared_ckpt_meta.update(self._meta_dict)
-        self._tensor_shm = _create_shared_memory(
-            name=self._shm_name,
-            create=True,
-            size=self._buffer_size,
-        )
-
-    def _copy_state_dict_to_shm(self, state_dict):
-        """
-        Copy the state dict from CPU memory buffer into the shared memory.
-        """
-
-        def _tarverse_copy(value, meta):
-            if isinstance(value, Mapping):
-                for k, v in value.items():
-                    if isinstance(v, (Mapping, List)):
-                        m = meta[k]
-                        _tarverse_copy(v, m)
-                    elif torch.is_tensor(v):
-                        m = meta[k]
-                        self._write_shared_memory(v, m)
-                    else:
-                        meta[k] = v
-            elif isinstance(value, List):
-                for i, v in enumerate(value):
-                    if isinstance(v, (Mapping, List)):
-                        m = meta[i]
-                        _tarverse_copy(v, m)
-                    elif torch.is_tensor(v):
-                        m = meta[i]
-                        self._write_shared_memory(v, m)
-                    else:
-                        meta[i] = v
-
-        _tarverse_copy(state_dict, self._meta_dict)
-        # Update the meta dict in the main process.
-        self._shared_ckpt_meta.update(self._meta_dict)
-
-    def _write_shared_memory(self, value, meta: TensorMeta):
-        """
-        Write a CPU tensor into the shared memory.
-        """
-        data_array = value.cpu().numpy()
-        write_array = np.ndarray(
-            data_array.shape,
-            dtype=data_array.dtype,
-            buffer=self._tensor_shm.buf,
-            offset=meta.offset,
-        )
-        if data_array.shape == ():
-            write_array.fill(data_array)
-        else:
-            write_array[:] = data_array[:]
-
-    @timer
-    def save_to_memory(self, state_dict, step):
-        """
-        Synchonously Saves the state dict into the shared memory with the main
-        process. If the agent in the main process is saving the shared memory
-        into the storage, the method will skip to write the shared memory.
-
-        Args:
-            state_dict (dict): the state dict of model and optimizer to save.
-            step (int): the iteration step.
-        """
-        state_dict["step"] = step
-        if self._tensor_shm is None:
-            self._make_state_dict_buffer(state_dict)
-        acquired = self._shm_buffer_lock.acquire(blocking=False)
-        all_rank_ready = self._check_all_rank_ready(acquired)
-        if not all_rank_ready:
-            logger.info(
-                f"Rank {self._rank} skips the save the checkpoint "
-                f"in CPU memory since it is saving the latest "
-                "checkpoint from the CPU memory into the storage."
-            )
-            if acquired:
-                self._shm_buffer_lock.release()
-            return
-        self._copy_state_dict_to_shm(state_dict)
-
-        if acquired:
-            self._shm_buffer_lock.release()
-        self._cached_step = step
-
-    def _check_all_rank_ready(self, ready):
-        """
-        Check wether all ranks are ready.
-        """
-        if not self._saver_group:
-            return ready
-        value = 0 if ready else 1
-        t = torch.tensor([value], dtype=torch.int64)
-        dist.all_reduce(t, group=self._saver_group)
-        return t == 0
-
-    @timer
-    def save_to_storage(self, state_dict, path, step):
-        """
-        Asynchonously saves the state dict into the storage. It synchonously
-        saves the state dict into the shared memory and put the path
-        into a shared queue. The agent in the main process waits for the queue
-        for save the state dict in the shared memory into the storage.
-
-        Args:
-            state_dict (dict): the state dict of model and optimizer to save.
-            path (str): optional, the file path to save the checkpoint. If the
-                path is not defined, the engine will save the state dict into
-                the shared memory not the storage.
-            step (int): the iteration step.
-        """
-        if step > self._cached_step:
-            self.save_to_memory(state_dict, step)
-        if path:
-            self._to_save_queue.put(path)
-
-    def load(self, resume_path=""):
-        """
-        The method firstly try to load the state dict from the shared memory.
-        If there is no state dict in the shared memory, the method will
-        load the state dict from the storage.
-
-        Returns:
-            A dict.
-        """
-        state_dict = self._load_from_shared_memory()
-        if state_dict:
-            return state_dict
-        state_dict = self._load_from_storage(resume_path)
-        return state_dict
-
-    def _load_from_shared_memory(self):
-        """
-        Load the state dict from the shared memory.
-
-        Returns:
-            A dict.
-        """
-        if self._tensor_shm is None:
-            self._tensor_shm = _create_shared_memory(
-                self._shm_name,
-                create=False,
-            )
-        if not self._tensor_shm:
-            return None
-        meta_dict = self._shared_ckpt_meta.get()
-        state_dict = read_state_dict_from_shm(meta_dict, self._tensor_shm)
-        return state_dict
-
-    def _load_from_storage(self, resume_path=""):
-        """
-        Load the state dict from the CPU memory if the state dict is complete
-        in CPU memory. Otherwise, the function will load the state dict from
-        the storage.
-
-        Args:
-            resume_path (str, optional): , If the resume_path is an empty
-                string, the function will load the latest checkpoint file in
-                the checkpoint directory.
-
-        Returns:
-            A dict:
-                a dictionary containing a whole state of the modules in the
-                checkpointing file.
-        """
-        if resume_path:
-            state_dict = torch.load(resume_path)
-        else:
-            state_dict = self._load_from_historic_checkpoint()
-        return state_dict
-
-    def _load_from_historic_checkpoint(self):
-        """Locd checkpoint from the lastest complete checkpoint."""
-        while True:
-            latest_ckpt_dir = _get_latest_checkpoint(self.checkpoint_dir)
-            if not latest_ckpt_dir:
-                return {}
-
-            resume_path = os.path.join(latest_ckpt_dir, "checkpoint.pt")
-            if not os.path.exists(resume_path):
-                shutil.rmtree(latest_ckpt_dir)
-                continue
-            try:
-                state_dict = torch.load(resume_path)
-                logger.info(f"Load checkpoint from {resume_path}")
-                return state_dict
-            except Exception:
-                logger.warning(
-                    f"Fail to load checkpoint from {resume_path}."
-                    " Roll back to the last checkpoint file."
-                )
-                shutil.rmtree(latest_ckpt_dir)
-
-
-class ShardingCheckpointEngine(CheckpointEngine):
-    """
-    The engine to save the sharding model and optimizer state dict
-    into the memory and storage. We can use it to save the model and optimizer
-    using FSDP, Zero-3 or Megatron-LM.
-    """
-
-    def __init__(self, checkpoint_dir):
-        super().__init__(checkpoint_dir)
-
-    def _notify_agent_to_create_saver(self):
-        # TODO: implement the saver in the agent to support saving
-        # sharding state dict.
-        pass
-
-    def _init_shared_objs(self):
-        meta_name = CKPT_META_NAME_PREFIX + str(self._local_rank)
-        self._shared_ckpt_meta = SharedDict(name=meta_name, create=False)
-        lock_name = SHM_LOCK_NAME_PREFIX + str(self._local_rank)
-        self._shm_buffer_lock = SharedLock(name=lock_name, create=False)
-        qname = SAVE_STEP_QNAME_PREFIX + str(self._local_rank)
-        self._to_save_queue = SharedQueue(name=qname, create=False)
-        self._shm_name = TENSOR_SHM_NAME_PREFIX + str(self._local_rank)
-
-
-class NoShardingCheckpointEngine(CheckpointEngine):
-    """
-    The engine saves the model and optimizer state dict without sharding
-    in a local or DDP job.
-    """
-
-    def __init__(self, checkpoint_dir):
-        super().__init__(checkpoint_dir)
-
-    def _notify_agent_to_create_saver(self):
-        queue = SharedQueue(name="factory")
-        queue.put("NoShardingSaver")
-        queue.close()
-
-    def _init_shared_objs(self):
-        """
-        Initialize the shared object with the main process.
-        Without model sharding, all ranks share the same shared memory
-        created by the local rank 0 on a node.
-        """
-        meta_name = CKPT_META_NAME_PREFIX + str(0)
-        self._shared_ckpt_meta = SharedDict(name=meta_name, create=False)
-        lock_name = SHM_LOCK_NAME_PREFIX + str(0)
-        self._shm_buffer_lock = SharedLock(name=lock_name, create=False)
-        qname = SAVE_STEP_QNAME_PREFIX + str(0)
-        self._to_save_queue = SharedQueue(name=qname, create=False)
-        self._shm_name = TENSOR_SHM_NAME_PREFIX + str(0)
-
-    @timer
-    def save_to_memory(self, state_dict, step):
-        """
-        Synchonously Saves the state dict into the shared memory with the main
-        process. If the agent in the main process is saving the shared memory
-        into the storage, the method will skip to write the shared memory.
-
-        Args:
-            state_dict (dict): the state dict of model and optimizer to save.
-            step (int): the iteration step.
-        """
-        if self._local_rank == 0:
-            super().save_to_memory(state_dict, step)
-
-    @timer
-    def save_to_storage(self, state_dict, path, step):
-        """
-        Asynchonously saves the state dict into the storage. It synchonously
-        saves the state dict into the shared memory and put the path
-        into a shared queue. The agent in the main process waits for the queue
-        for save the state dict in the shared memory into the storage.
-
-        Args:
-            state_dict (dict): the state dict of model and optimizer to save.
-            step (int): the iteration step.
-            path (str): optional, the file path to save the checkpoint. If the
-                path is not defined, the engine will save the state dict into
-                the shared memory not the storage.
-        """
-        if self._rank == 0:
-            super().save_to_storage(state_dict, path, step)
-
-
 class CheckpointManger(metaclass=ABCMeta):
     """CheckpontManager can save and load checkpoint states.
 


Test patch:
diff --git a/dlrover/python/tests/test_ckpt_saver.py b/dlrover/python/tests/test_ckpt_saver.py
index 76e7e5870..c8ae3e834 100644
--- a/dlrover/python/tests/test_ckpt_saver.py
+++ b/dlrover/python/tests/test_ckpt_saver.py
@@ -12,6 +12,8 @@
 # limitations under the License.
 
 import os
+import tempfile
+import time
 import unittest
 
 import numpy as np
@@ -19,11 +21,16 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from dlrover.python.common.multi_process import SharedMemory
+from dlrover.python.common.multi_process import SharedMemory, SharedQueue
 from dlrover.python.elastic_agent.torch.ckpt_saver import (
+    _WIRTING_SHM,
+    CheckpointSaver,
+    NoShardingCheckpointEngine,
     NoShardingSaver,
-    convert_torch_dtype_to_numpy,
-    traverse_state_dict,
+    ShardingCheckpointEngine,
+    _clean_shm_handler,
+    _convert_torch_dtype_to_numpy,
+    _traverse_state_dict,
 )
 
 
@@ -51,8 +58,20 @@ def forward(self, x):
 
 
 class CheckpointSaverTest(unittest.TestCase):
+    def test_create_checkpoint_saver(self):
+        CheckpointSaver.start_async_saving_ckpt(num_proc=8)
+        sq = SharedQueue(name="factory", create=False)
+        args = {"checkpoint_dir": "test_dir"}
+        sq.put(("NoShardingSaver", args))
+        for _ in range(10):
+            if CheckpointSaver._saver_instance is None:
+                time.sleep(0.5)
+            else:
+                break
+        self.assertEqual(CheckpointSaver._saver_instance.num_proc, 8)
+
     def test_close_saver(self):
-        saver = NoShardingSaver()
+        saver = NoShardingSaver("test_ckpt")
         saver._tensor_shm = SharedMemory(name="test", create=True, size=1024)
         saver.close()
         saver.close()
@@ -67,15 +86,44 @@ def visitor(value):
             model=model.state_dict(),
             step=step,
         )
-        new_dict = traverse_state_dict(state_dict, visitor)
+        new_dict = _traverse_state_dict(state_dict, visitor)
         self.assertEqual(new_dict, state_dict)
 
     def test_convert_torch_dtype_to_numpy(self):
-        np_dtype = convert_torch_dtype_to_numpy(torch.float32)
+        np_dtype = _convert_torch_dtype_to_numpy(torch.float32)
         self.assertEqual(np_dtype, np.float32)
 
-        np_dtype = convert_torch_dtype_to_numpy(torch.float)
+        np_dtype = _convert_torch_dtype_to_numpy(torch.float)
         self.assertEqual(np_dtype, np.float32)
 
-        np_dtype = convert_torch_dtype_to_numpy(torch.int32)
+        np_dtype = _convert_torch_dtype_to_numpy(torch.int32)
         self.assertEqual(np_dtype, np.int32)
+
+    def test_save_to_storage(self):
+        model = SimpleNet()
+        step = 100
+        state_dict = dict(
+            model=model.state_dict(),
+            step=step,
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            CheckpointSaver._saver_instance = NoShardingSaver(tmpdir)
+            sq = SharedQueue(name="factory", create=True)
+            saving_engine = NoShardingCheckpointEngine(tmpdir)
+            saving_engine.save_to_memory(state_dict, step)
+            self.assertFalse(saving_engine._meta_dict[_WIRTING_SHM])
+            saver: CheckpointSaver = CheckpointSaver.get_ckpt_saver()
+            saver._tensor_shm = SharedMemory(name=saver._shm_name)
+            saver.save_shm_to_storage()
+            ckpt_files = os.listdir(tmpdir)
+            self.assertEqual(len(ckpt_files), 1)
+            sq.close()
+            _clean_shm_handler(None, None)
+
+    def test_sharding_checkpoint_engine(self):
+        os.environ["LOCAL_RANK"] = "1"
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            engine = ShardingCheckpointEngine(tmpdirname)
+            self.assertEqual(
+                engine._shared_ckpt_meta._name, "checkpoint_meta_local_rank_1"
+            )
diff --git a/dlrover/python/tests/test_elastic_training_agent.py b/dlrover/python/tests/test_elastic_training_agent.py
index b0e44dc5e..c6b977665 100644
--- a/dlrover/python/tests/test_elastic_training_agent.py
+++ b/dlrover/python/tests/test_elastic_training_agent.py
@@ -29,6 +29,10 @@
     build_master_client,
 )
 from dlrover.python.elastic_agent.monitor.training import TorchTrainingMonitor
+from dlrover.python.elastic_agent.torch.ckpt_saver import (
+    CheckpointSaver,
+    NoShardingSaver,
+)
 from dlrover.python.elastic_agent.torch.training import (
     ElasticLaunchConfig,
     ElasticTrainingAgent,
@@ -270,6 +274,9 @@ def test_restart_training(self):
             start_method=self.config.start_method,
             log_dir=self.config.log_dir,
         )
+        saver = NoShardingSaver("/tmp/test")
+        CheckpointSaver._saver_instance = saver
+        agent._save_ckpt_to_storage()
         agent._stop_workers_to_restart()
 
 
diff --git a/dlrover/trainer/tests/torch/checkpoint_test.py b/dlrover/trainer/tests/torch/checkpoint_test.py
index 3f7b02eaf..2bd056a51 100644
--- a/dlrover/trainer/tests/torch/checkpoint_test.py
+++ b/dlrover/trainer/tests/torch/checkpoint_test.py
@@ -26,14 +26,13 @@
 from torch.utils.data import DataLoader, Dataset
 
 from dlrover.python.common import grpc
-from dlrover.python.elastic_agent.torch.ckpt_saver import CheckpointSaver
-from dlrover.trainer.torch.elastic.checkpoint import (
-    CheckpointManger,
+from dlrover.python.elastic_agent.torch.ckpt_saver import (
+    CheckpointSaver,
     NoShardingCheckpointEngine,
-    ShardingCheckpointEngine,
     _create_shared_memory,
     _get_latest_checkpoint,
 )
+from dlrover.trainer.torch.elastic.checkpoint import CheckpointManger
 from dlrover.trainer.torch.elastic.sampler import ElasticDistributedSampler
 
 
@@ -103,6 +102,7 @@ def _wait_async_saving_finished(dir_name, step):
 
 class CheckpointManagerTest(unittest.TestCase):
     def setUp(self):
+        CheckpointSaver._saver_instance = None
         CheckpointSaver.start_async_saving_ckpt()
 
     def test_create_shared_memory(self):
@@ -172,11 +172,3 @@ def test_load_no_sharding(self):
                 loaded_value = loaded_state_dict["model"][key]
                 self.assertTrue(torch.equal(value, loaded_value))
             engine.close()
-
-    def test_sharding_checkpoint_engine(self):
-        os.environ["LOCAL_RANK"] = "1"
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            engine = ShardingCheckpointEngine(tmpdirname)
-            self.assertEqual(
-                engine._shared_ckpt_meta._name, "checkpoint_meta_local_rank_1"
-            )