mindspore 2.7.0__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/train/callback/_train_fault_tolerance.py

@@ -28,15 +28,15 @@ from mindspore._c_expression import _repair_device, _stop_device, _tft_sem_post,
 from mindspore._c_expression import _rebuild_world_group, _rebuild_sub_group, _finalize_comm, _clean_rootinfo
 from mindspore._c_expression import clean_tdt_channel
 from mindspore._c_expression import _pre_launch_send_recv
-from mindspore._c_expression import send_recv, reset_params
+from mindspore._c_expression import send_recv, reset_params, direct_copy_to_host
+from mindspore._c_expression import _reg_snapshot_params, _reset_snapshot_state, _clear_snapshot_saving_flag
 from mindspore._c_expression import CollectiveManager
 from mindspore._c_expression import _get_uce_process_strategy, _get_uce_mem_info
-from mindspore._c_expression import TensorPy as Tensor_
 from mindspore.ops.operations.manually_defined._inner import TensorReport
 import mindspore
 import mindspore.common.dtype as mstype
-from mindspore.parallel._recovery_context import _set_recovery_context
 from mindspore import runtime
+from mindspore._c_expression import set_is_arf
 
 
 def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
@@ -157,6 +157,7 @@ def _tft_clean_callback(is_uce_error, args, ctx):
     CollectiveManager.get_instance().resume_hccl_comm()
     logger.warning("Finish _tft_clean_callback, ret: {}".format(ret))
     if ctx.tft.tft_get_repair_type() == "recover":
+        _reset_snapshot_state()
         logger.warning(f"Destroy hcom")
         _finalize_comm()
         logger.warning(f"Destroy hcom end")
@@ -166,10 +167,10 @@ def _tft_clean_callback(is_uce_error, args, ctx):
 def _tft_stop_callback(args, cb_ctx):
     """ Callback used for TFT stop function."""
     logger.warning(f"Enter _tft_stop_callback device_id: {cb_ctx.device_id}")
-    _stop_device(cb_ctx.device_id)
     if (not cb_ctx.is_uce_rank) and (not cb_ctx._is_params_consistent()):  # pylint: disable=W0212
         raise RuntimeError("Can't stop device, because training parameters are left in inconsistent state!")
     cb_ctx.is_uce_rank = False
+    _stop_device(cb_ctx.device_id)
     if cb_ctx.tft.tft_get_repair_type() == "recover":
         logger.warning(f"Reset limit step")
         cb_ctx.tft.tft_reset_limit_step()
@@ -181,7 +182,7 @@ def _tft_rebuild_sub_groups(fault_ranks, args, ctx):
     logger.warning(f"Enter _tft_rebuild_sub_groups, device id: {ctx.device_id}")
     _rebuild_world_group()
     _rebuild_sub_group()
-    _set_recovery_context(is_arf=True)
+    set_is_arf(True)
     logger.warning(f"try to pre launch send recv before real launch")
     _pre_launch_send_recv(context.get_context('device_id'))
     logger.warning(f"Pre launch send recv before real launch end")
@@ -201,7 +202,10 @@ class TrainFaultTolerance(Callback):
         ckpt_save_path (str): Checkpoint save directory when failure occurs. When saved,
             a new directory named 'ttp_saved_checkpoints-step_{cur_step_num}'
             is created in that directory. Default: ``None``.
-        kwargs (dict): Other dictionary type parameters.
+        kwargs (dict): Other dictionary type parameters. When argument `ckpt_save_path` is ``None``, `kwargs` must
+            provide a parameter named `ckpt_save_fn`, which points to a function used to save checkpoint. The
+            prototype of `ckpt_save_fn` is ``def save_ckpt(cb_params, append_dict)``. When both `ckpt_save_path`
+            and `ckpt_save_fn` are provided, `ckpt_save_fn` is used in priority.
 
     Raises:
         Exception: TFT init failed.
@@ -328,7 +332,7 @@ class TrainFaultTolerance(Callback):
         # `def load_checkpoint() -> tuple(dict, bool)`, the return value is a tuple containing 2 values,
         # i.e. (param_dict, remove_redundancy)
         self.ckpt_load_func = kwargs.get("ckpt_load_fn", None)
-        if self._only_enable_tre():
+        if self._only_enable_tre() or self._only_enable_ckpt_d2h_async():
             return
         self.tft = _tft_handler.get_tft()
         self._check_init()
@@ -341,8 +345,7 @@ class TrainFaultTolerance(Callback):
         self.is_uce_rank = False
 
         self.assign = mindspore.ops.Assign()
-        self.g_one = Parameter(Tensor([1], dtype=mstype.int32))
-        self.s1 = mindspore.hal.Stream()
+        self.g_one = Tensor([1], dtype=mstype.int32)
         _tft_sem_enable()
         self._tft_register()
 
@@ -352,7 +355,21 @@ class TrainFaultTolerance(Callback):
         non_tre_flags = ["TTP:1", "UCE:1", "ARF:1"]
         if any(flag in env_enable for flag in non_tre_flags):
             return False
-        return "TRE:1" in env_enable
+        return "TRE:1" in env_enable or "TRE:2" in env_enable
+
+    @staticmethod
+    def _only_enable_ckpt_d2h_async():
+        """Check whether only set MS_ENABLE_CKPT_D2H_ASYNC=1 without setting MS_ENABLE_TFT"""
+        if os.getenv("MS_ENABLE_TFT", "") != "":
+            return False
+        return os.getenv("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+
+    @staticmethod
+    def _enable_snapshot():
+        """Check whether parameter snapshot enabled"""
+        enable_step_tre = "TRE:2" in os.getenv("MS_ENABLE_TFT", "")
+        enable_ckpt_d2h_async = os.getenv("MS_ENABLE_CKPT_D2H_ASYNC") == "1"
+        return enable_step_tre or enable_ckpt_d2h_async
 
     def _only_enable_tsp(self):
         """Check if only configured MS_ENABLE_TFT='{TSP:1}'"""
@@ -387,9 +404,7 @@ class TrainFaultTolerance(Callback):
     def _is_params_consistent(self):
         for key, param in self.cb_params.train_network.parameters_and_names():
             if "tft_g_one_flag" in key:
-                with mindspore.hal.StreamCtx(self.s1):
-                    tft_g_one_flag = Tensor(Tensor_.move_to(param, "CPU", False))
-                self.s1.synchronize()
+                tft_g_one_flag = direct_copy_to_host(param)
                 return int(tft_g_one_flag) == 1
         return False
 
@@ -434,7 +449,7 @@ class TrainFaultTolerance(Callback):
                 super(TFTOptSubCls, self).__init__(*args, **kwargs)
                 self.report = TensorReport()
                 self.report_end = TensorReport()
-                self.report_end.add_prim_attr("side_effect_mem", True).add_prim_attr("optimizer_end", True)
+                self.report_end.add_prim_attr("optimizer_end", True)
                 self.depend = ops.Depend()
                 self.allreduce_sum = ops.AllReduce()
                 self.allreduce_sum.add_prim_attr("tft_report_before", True)
@@ -448,7 +463,27 @@ class TrainFaultTolerance(Callback):
                 self.report_end("tft_report", self.tft_g_one_flag)
                 return opt_ret
 
-        return TFTOptSubCls
+        class TFTOptSnapShotCls(origin_opt_cls):
+            """
+            Optimizer wrapper class when using tft.
+            """
+
+            def __init__(self, *args, **kwargs):
+                super(TFTOptSnapShotCls, self).__init__(*args, **kwargs)
+                self.report = TensorReport()
+                self.report.add_prim_attr("side_effect_mem", True).add_prim_attr("snapshot", True)
+                self.dummy_input = Tensor([1], dtype=mstype.int32)
+
+            def construct(self, gradients, **kwargs):
+                """Add fake op TensorReport to insert wait event for copying parameters"""
+                self.report("tft_report", self.dummy_input)
+                opt_ret = super(TFTOptSnapShotCls, self).construct(gradients, **kwargs)
+                return opt_ret
+
+        env_tft = os.getenv('MS_ENABLE_TFT', '')
+        features = ['TTP:1', 'UCE:1', 'ARF:1']
+        need_redundancy = any([env_tft.find(feat) >= 0 for feat in features])
+        return TFTOptSubCls if need_redundancy else TFTOptSnapShotCls
 
     def _tft_register(self):
         """Register callback functions."""
@@ -476,6 +511,17 @@ class TrainFaultTolerance(Callback):
             _clean_rootinfo()
             self.clean_unique_id = True
 
+    def on_train_step_begin(self, run_context):
+        """
+        Clear saving snapshot state at each step begin.
+
+        Args:
+            run_context (RunContext): Context of the train running. Refer to
+                                      :class:`mindspore.train.RunContext` for detail.
+        """
+        if self._enable_snapshot():
+            _clear_snapshot_saving_flag()
+
     def on_train_step_end(self, run_context):
         """
         Report status to MindIO TFT after every step finished.
@@ -484,7 +530,7 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
-        if self._only_enable_tre():
+        if self._only_enable_tre() or self._only_enable_ckpt_d2h_async():
             return
 
         cb_params = run_context.original_args()
@@ -524,10 +570,15 @@ class TrainFaultTolerance(Callback):
             run_context (RunContext): Context of the train running. Refer to
                                       :class:`mindspore.train.RunContext` for detail.
         """
+        cb_params = run_context.original_args()
+        if self._enable_snapshot():
+            param_dict = {}
+            for param in cb_params.train_network.trainable_params():
+                param_dict[param.name] = param
+            _reg_snapshot_params(param_dict)
         if self._only_enable_tsp():
             return
-        cb_params = run_context.original_args()
-        if self._only_enable_tre():
+        if self._only_enable_tre() or self._only_enable_ckpt_d2h_async():
             self.cb_params = cb_params
             return
         sink_size = cb_params.get("sink_size", 0)

mindspore/train/data_sink.py

@@ -238,12 +238,8 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
 
         real_sink_fun = _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config)
 
-        loop = sink_size
-        if jit_config is not None and context.get_context('mode') == context.GRAPH_MODE:
-            loop = 1
-
         out = None
-        for _ in range(loop):
+        for _ in range(sink_size):
             out = real_sink_fun()
 
         return out

mindspore/train/model.py

@@ -28,7 +28,7 @@ import numpy as np
 
 import mindspore
 from mindspore import log as logger
-from mindspore.train.serialization import save_checkpoint, load_checkpoint
+from mindspore.train.serialization import save_checkpoint
 from mindspore.train.callback._checkpoint import ModelCheckpoint, _chg_ckpt_file_name_if_same_exist
 from mindspore.common.tensor import Tensor
 from mindspore.train.metrics import get_metrics, get_metric_fn
@@ -40,16 +40,12 @@ from mindspore.train.callback import __all__ as internal_cb_names
 from mindspore.train.callback._cluster_monitor import ClusterMonitor
 from mindspore import context
 from mindspore.parallel._utils import _get_parallel_mode, _get_device_num, _get_parameter_broadcast, \
-    _device_number_check, _parameter_broadcast_check, _parallel_predict_check, \
-    _reset_op_id_with_offset
-from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _is_ps_mode, \
-    _cache_enable, _enable_distributed_mindrt
+    _device_number_check, _parameter_broadcast_check, _parallel_predict_check
 from mindspore.train.metrics import Loss
 from mindspore.log import vlog_print
 from mindspore import nn
 from mindspore.boost import AutoBoost
 from mindspore.context import ParallelMode
-from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
 from mindspore.train.dataset_helper import DatasetHelper, connect_network_with_dataset
 from mindspore.common.api import _pynative_executor, ARG_SPECIFIED, TOTAL_ARG_LEN
 from mindspore.dataset.core.config import get_debug_mode
@@ -57,7 +53,8 @@ from mindspore.dataset.engine.datasets import _set_training_dataset, _reset_trai
 from mindspore.train import amp
 from mindspore._c_expression import _framework_profiler_step_start, _framework_profiler_step_end
 from mindspore._c_expression import _get_optimzer_timestamps
-from mindspore._c_expression import clean_tdt_channel, _clean_rootinfo
+from mindspore._c_expression import clean_tdt_channel, _clean_rootinfo, check_is_arf, set_is_arf
+from mindspore._c_expression import _get_snapshot_params, _is_snapshot_valid
 
 from mindspore.parallel._utils import _init_auto_parallel_context, _clear_auto_parallel_context
 from .serialization import load_param_into_net
@@ -163,7 +160,7 @@ def _handle_exception_info(obj, uce_env, tft, e):
         tft.tft_report_error(force_stop_err)
     elif "ARF FINISH" in e_str:
         logger.warning(f"ARF FINISH")
-        _set_recovery_context(is_arf=True)
+        set_is_arf(True)
         tft.tft_report_error(tft.ReportState.RS_PREREPAIR_FINISH.value)
     else:
         logger.error("uce wrapper caught other RuntimeError, enter MindIO TTP process.", exc_info=True)
@@ -175,7 +172,12 @@ def _handle_training_result_error(model, tft_obj):
     """
     Handle training result error for resuming training.
     """
-    ckpt_load_fn = tft_obj.ckpt_load_func
+    def load_snapshot_params():
+        param_dict = {}
+        for name, tensor in _get_snapshot_params().items():
+            param_dict[name] = mindspore.Parameter(tensor, name=name)
+        return (param_dict, False)
+    ckpt_load_fn = load_snapshot_params if _is_snapshot_valid() else tft_obj.ckpt_load_func
     train_network = tft_obj.cb_params.train_network
     logger.warning("Process training result error start.")
     # 1. Clear tdt channel
@@ -234,6 +236,20 @@ def _update_ckpt_callback_info(resume_train_step, **kwargs):
         ckpt_obj._append_step_num = resume_train_step
 
 
+def _get_tft_obj(**kwargs):
+    """
+    Get TrainFaultTolerance from kwargs of callback
+    """
+    obj = None
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), TrainFaultTolerance):
+        obj = kwargs.get('callbacks')
+    if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
+        for item in kwargs.get('callbacks'):
+            if isinstance(item, TrainFaultTolerance):
+                obj = item
+    return obj
+
+
 def _handle_tft(func):
     """
     Decorator function, which starts uce handle process when an exception occurs during training.
@@ -241,17 +257,11 @@ def _handle_tft(func):
 
     @wraps(func)
     def wrapper(self, *args, **kwargs):
-        obj = None
-        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), TrainFaultTolerance):
-            obj = kwargs.get('callbacks')
-        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
-            for item in kwargs.get('callbacks'):
-                if isinstance(item, TrainFaultTolerance):
-                    obj = item
-        if obj:
+        obj = _get_tft_obj(**kwargs)
+        if obj and not TrainFaultTolerance._only_enable_ckpt_d2h_async():
             tft_env = os.getenv("MS_ENABLE_TFT", "")
             uce_env = "UCE:1" in tft_env or "ARF:1" in tft_env or "HCCE:1" in tft_env
-            tre_env = "TRE:1" in tft_env
+            tre_env = "TRE:1" in tft_env or "TRE:2" in tft_env
             while True:
                 try:
                     return func(self, *args, **kwargs)
@@ -556,9 +566,7 @@ class Model:
         self._current_epoch_num = 0
         self._current_step_num = 0
         self.epoch_iter = 0
-        self.enable_recovery = False
         self._backbone_is_train = True
-        self.need_load_ckpt = False
         self._lite_full_predictor = None
         self._lite_incremental_predictor = None
         self._mindspore_lite = None
@@ -731,10 +739,7 @@ class Model:
         metrics = dict()
         # There's no need for server to execute eval, just give fake metrics.
         for key, value in self._metric_fns.items():
-            if not _is_role_pserver():
-                metrics[key] = value.eval()
-            else:
-                metrics[key] = 1
+            metrics[key] = value.eval()
         return metrics
 
     def _get_scaling_sens(self):
@@ -768,7 +773,7 @@ class Model:
             logger.info("Begin to connect network with dataset.")
             network = connect_network_with_dataset(network, dataset_helper)
 
-        if (_get_recovery_context("enable_recovery") or self._need_reset_data) and is_train:
+        if self._need_reset_data and is_train:
             _set_training_dataset(dataset_helper)
 
         network.set_train(is_train)
@@ -810,9 +815,7 @@ class Model:
         :param cb_params: callback params
         :return: none
         """
-        if os.environ.get("MS_ENABLE_CKPT_D2H_ASYNC") != "1":
-            return
-        if context.get_context("device_target") == "Ascend":
+        if TrainFaultTolerance._enable_snapshot() and context.get_context("device_target") == "Ascend":
             cb_params.need_ckpt, cb_params.save_checkpoint_steps, \
             cb_params.last_triggered_step = self._check_need_ckpt(cb_params.list_callback)
             logger.info(f"need_ckpt:{cb_params.need_ckpt},"
@@ -1018,13 +1021,10 @@ class Model:
         callbacks = cb_params.list_callback
         cb_params.train_dataset_element = None
         cb_params.network = self._network
-        # Embedding cache server only run one step.
-        if _is_role_pserver() and _cache_enable():
-            epoch = 1
         cb_params.last_save_ckpt_step = None
         cb_params.latest_ckpt_file = None
         cb_params.loss_scale_mananger = self._loss_scale_manager
-        cb_params.is_arf = _get_recovery_context("is_arf")
+        cb_params.is_arf = check_is_arf()
         cb_params.initial_step = self._initial_step
 
         # build callback list
@@ -1086,12 +1086,6 @@ class Model:
             dataset_helper = train_dataset._dataset_helper
 
         self.epoch_iter = 0
-        self._check_enable_recovery()
-        # Used to check whether need perform recovery for process which is restarted.
-        self._check_need_load_ckpt(cb_params, dataset_size, sink_size)
-        # Check whether this process is embedding cache server.
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
-
         while self.epoch_iter < (epoch - initial_epoch):
             cb_params.cur_epoch_num = self.epoch_iter + 1 + initial_epoch
             self._current_epoch_num = cb_params.cur_epoch_num
@@ -1107,11 +1101,6 @@ class Model:
             cb_params.train_network = train_network
             cb_params.dataset_helper = dataset_helper
 
-            # Perform recovery for process which is restarted.
-            self._reset_training_step_for_abnormal_process(cb_params, dataset_helper)
-            # Perform recovery for process which is not restarted.
-            self._reset_training_step_for_normal_process(cb_params, dataset_helper)
-
             # For data sink dataset_helper only iter once, other wise iter epoch_size times.
             for inputs in dataset_helper:
                 if is_graph:
@@ -1126,36 +1115,17 @@ class Model:
                 outputs = train_network(*inputs)
                 cb_params.net_outputs = outputs
 
-                # In disaster recovery scenarios, need not to execute callbacks if this step executes failed.
-                need_exec_callback_step_end = not (self.enable_recovery and _get_recovery_context("need_reset"))
-                if need_exec_callback_step_end:
-                    list_callback.on_train_step_end(run_context)
+                list_callback.on_train_step_end(run_context)
+
                 if cb_params.is_arf:
                     cb_params.is_arf = False
-                    _set_recovery_context(is_arf=False)
+                    set_is_arf(False)
                 _clean_rootinfo()
 
-                # Embedding cache server only run one step.
-                if is_embedding_cache_server:
-                    break
-
             dataset_helper.continue_send()
 
-            # When it's distributed training and using MindRT,
-            # the node id should be reset to start from 0.
-            # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-            if _enable_distributed_mindrt():
-                _reset_op_id_with_offset()
-
             self._eval_during_train(valid_infos, cb_params, list_callback)
-
-            # In disaster recovery scenarios, need not to execute callbacks if this epoch executes failed.
-            # Embedding cache server need not do epoch end callback, this process only run one step.
-            need_exec_callback_epoch_end = not ((self.enable_recovery and _get_recovery_context("need_reset"))
-                                                or is_embedding_cache_server)
-
-            if need_exec_callback_epoch_end:
-                list_callback.on_train_epoch_end(run_context)
+            list_callback.on_train_epoch_end(run_context)
             if "metrics" in cb_params or "eval_results" in cb_params:
                 cb_params.pop("metrics", None)
                 cb_params.pop("eval_results", None)
@@ -1164,12 +1134,7 @@ class Model:
             if should_stop:
                 break
 
-            need_reset_to_beginning = self.enable_recovery and _get_recovery_context("need_reset") \
-                                      and not _get_recovery_context("latest_ckpt_file")
             self.epoch_iter += 1
-            if need_reset_to_beginning:
-                self.epoch_iter = 0
-                cb_params.cur_step_num = 0
 
         dataset_helper.stop_send()
         dataset_helper.release()
@@ -1203,93 +1168,6 @@ class Model:
             cb_params.dataset_sink_mode = train_dataset_sink_mode
             cb_params.net_outputs = train_net_outputs
 
-    def _check_enable_recovery(self):
-        """
-        Check whether enable recovery and execution mode consistency.
-        """
-
-        enable_recovery = _get_recovery_context("enable_recovery") and context.get_context("device_target") == "GPU"
-        if not enable_recovery:
-            self.enable_recovery = False
-        else:
-            self.enable_recovery = enable_recovery and _is_role_worker()
-
-    def _check_need_load_ckpt(self, cb_params, dataset_size, sink_size=-1):
-        """
-        Check whether need to load checkpoint after abnormal process restart.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-            dataset_size (int): The number of batches in a dataset.
-            sink_size (int): Control the amount of data in each sink. Default: -1.
-        """
-        if context.get_context("device_target") != "GPU":
-            return
-        if not self.enable_recovery:
-            self.need_load_ckpt = False
-
-        cb_params.latest_ckpt_file = _get_recovery_context("latest_ckpt_file")
-        if cb_params.latest_ckpt_file:
-            recovery_epoch_num = _get_recovery_context("latest_ckpt_epoch")
-            recovery_step_num = _get_recovery_context("latest_ckpt_step")
-            dataset_sink_size = sink_size if sink_size > 0 else dataset_size
-            cb_params.cur_step_num = (recovery_epoch_num - 1) * dataset_sink_size + recovery_step_num
-            cb_params.last_save_ckpt_step = cb_params.cur_step_num
-            self.epoch_iter = recovery_epoch_num
-            self.need_load_ckpt = True
-        else:
-            self.need_load_ckpt = False
-
-    def _reset_training_step_for_abnormal_process(self, cb_params, dataset_helper):
-        """
-        Execute recovery for abnormal exit process when restart.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-        """
-
-        if self.need_load_ckpt:
-            try:
-                load_checkpoint(cb_params.latest_ckpt_file, cb_params.train_network)
-            except BaseException as e:
-                os.remove(cb_params.latest_ckpt_file)
-                raise RuntimeError(e.__str__() + ", load ckpt failed and remove the ckpt: " \
-                                   + cb_params.latest_ckpt_file) from e
-            _reset_training_dataset(cb_params.cur_step_num, dataset_helper.iter.dataset.get_dataset_size())
-            self.need_load_ckpt = False
-
-    def _reset_training_step_for_normal_process(self, cb_params, dataset_helper):
-        """
-        Execute recovery for normal process when there is process exit abnormally.
-
-        Args:
-            cb_params (_InternalCallbackParam): Callback parameters.
-            dataset_helper (DatasetHelper): A class to process the MindData dataset,
-                it provides the type, shape and queue name of the dataset to wrap the `GetNext`.
-        """
-
-        if self.enable_recovery and _get_recovery_context("need_reset"):
-            cb_params.latest_ckpt_file = _get_recovery_context("latest_ckpt_file")
-            if cb_params.latest_ckpt_file:
-                try:
-                    load_checkpoint(cb_params.latest_ckpt_file, cb_params.train_network)
-                except BaseException as e:
-                    os.remove(cb_params.latest_ckpt_file)
-                    raise RuntimeError(e.__str__() + ", load ckpt failed and remove the ckpt: "\
-                         + cb_params.latest_ckpt_file) from e
-
-                recovery_epoch_num = _get_recovery_context("latest_ckpt_epoch")
-                recovery_step_num = _get_recovery_context("latest_ckpt_step")
-                cb_params.cur_step_num = (recovery_epoch_num - 1) * dataset_helper.sink_size() + recovery_step_num
-                self.epoch_iter = recovery_epoch_num
-                cb_params.cur_epoch_num = self.epoch_iter + 1
-                cb_params.last_save_ckpt_step = cb_params.cur_step_num
-                _reset_training_dataset(cb_params.cur_step_num, dataset_helper.iter.dataset.get_dataset_size())
-            else:
-                _reset_training_dataset(0, dataset_helper.iter.dataset.get_dataset_size())
-
-            _set_recovery_context(need_reset=False)
-
     def _train_process(self, epoch, train_dataset, list_callback=None, cb_params=None, initial_epoch=0,
                        valid_infos=None):
         """
@@ -1314,7 +1192,6 @@ class Model:
         cb_params.dataset_sink_mode = False
         run_context = RunContext(cb_params)
         list_callback.on_train_begin(run_context)
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
 
         for i in range(initial_epoch, epoch):
             cb_params.cur_epoch_num = i + 1
@@ -1345,21 +1222,12 @@ class Model:
                 list_callback.on_train_step_end(run_context)
                 if cb_params.is_arf:
                     cb_params.is_arf = False
-                    _set_recovery_context(is_arf=False)
+                    set_is_arf(False)
                 _clean_rootinfo()
-                # Embedding cache server only run one step.
-                if is_embedding_cache_server:
-                    break
                 should_stop = run_context.get_stop_requested()
                 if should_stop:
                     break
 
-            # When it's distributed training and using MindRT,
-            # the node id should be reset to start from 0.
-            # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-            if _enable_distributed_mindrt():
-                _reset_op_id_with_offset()
-
             self._eval_during_train(valid_infos, cb_params, list_callback)
 
             train_dataset.reset()
@@ -1367,9 +1235,7 @@ class Model:
             # if param is cache enable, flush data from cache to host before epoch end
             self._flush_from_cache(cb_params)
 
-            # Embedding cache server need not do epoch end callback, this process only run one step.
-            if not is_embedding_cache_server:
-                list_callback.on_train_epoch_end(run_context)
+            list_callback.on_train_epoch_end(run_context)
             if "metrics" in cb_params or "eval_results" in cb_params:
                 cb_params.pop("metrics", None)
                 cb_params.pop("eval_results", None)
@@ -1446,10 +1312,6 @@ class Model:
         """
         _init_auto_parallel_context(self._network)
         _check_tft()
-        device_target = context.get_context("device_target")
-        if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
-            logger.info("For PS mode, reset datasink mode to False when using Ascend or CPU backend.")
-            dataset_sink_mode = False
 
         Validator.check_bool(dataset_sink_mode)
         if isinstance(self._train_network, nn.GraphCell) and dataset_sink_mode:
@@ -1461,11 +1323,6 @@ class Model:
                              "the value of epoch in train {} separately."
                              .format(train_dataset._warmup_epoch, epoch))
 
-        # Parameter server and embedding cache mode check.
-        if _is_ps_mode():
-            if not dataset_sink_mode and _cache_enable():
-                raise ValueError("Embedding cache mode should run with 'dataset_sink_mode=True'.")
-
         self._check_sink_mode_for_ds_debug_mode(dataset_sink_mode)
 
         Validator.check_is_int(sink_size)
@@ -1496,12 +1353,6 @@ class Model:
                     sink_size=sink_size,
                     initial_epoch=initial_epoch)
 
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
-
         _clear_auto_parallel_context(self._network)
 
     @staticmethod
@@ -1599,10 +1450,6 @@ class Model:
             >>> model.fit(2, train_dataset, valid_dataset)
         """
         _init_auto_parallel_context(self._network)
-        device_target = context.get_context("device_target")
-        if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
-            logger.info("For PS mode, reset datasink mode to False when using Ascend or CPU backend.")
-            dataset_sink_mode = False
 
         dataset_sink_mode = Validator.check_bool(dataset_sink_mode)
         valid_dataset_sink_mode = Validator.check_bool(valid_dataset_sink_mode)
@@ -1896,13 +1743,6 @@ class Model:
 
         self._clear_metrics()
 
-        # Embedding cache server as a storage service, no need to execute eval.
-        is_embedding_cache_server = _is_role_pserver() and _cache_enable()
-        if is_embedding_cache_server:
-            metrics = self._get_metrics()
-            cb_params.metrics = metrics
-            return metrics
-
         if context.get_context("device_target") == "CPU" and dataset_sink_mode:
             dataset_sink_mode = False
             logger.info("CPU cannot support dataset sink mode currently."
@@ -1914,13 +1754,7 @@ class Model:
             else:
                 eval_result = self._eval_process(valid_dataset, list_callback, cb_params)
 
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
         _clear_auto_parallel_context(self._network)
-
         return eval_result
 
     def _predict_lite(self, *predict_data, config=None):
@@ -2171,13 +2005,6 @@ class Model:
         result = self._predict_network(*predict_data)
 
         check_output_data(result)
-
-        # When it's distributed training and using MindRT,
-        # the node id should be reset to start from 0.
-        # This is to avoid the timeout when finding the actor route tables in 'train' and 'eval' case(or 'fit').
-        if _enable_distributed_mindrt():
-            _reset_op_id_with_offset()
-
         return result
 
     def _infer_train_check(self, train_dataset, dataset_sink_mode, sink_size):