chore: import upstream snapshot with attribution

python/paddle/base/incubate/__init__.py

@@ -0,0 +1,17 @@
+#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+# incubate directory is mainly for internal use
+# after we have tested incubate APIs in industrial application for a period
+# we will move stable functions into base
+__version__ = '0.1.0'

python/paddle/base/incubate/checkpoint/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

python/paddle/base/incubate/checkpoint/auto_checkpoint.py

@@ -0,0 +1,710 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import os
+import sys
+import time
+from threading import current_thread
+
+from paddle.base import compiler, unique_name
+from paddle.base.framework import Program, in_dygraph_mode
+
+from .checkpoint_saver import CheckpointSaver, PaddleModel, SerializableBase
+
+g_train_epoch_range = None
+g_checker = None
+
+logger = None
+
+generator = unique_name.UniqueNameGenerator()
+
+CONST_CHECKPOINT = "checkpoint"
+CONST_MEMORYINIT = "memory_init"
+
+# auto checkpoint by dataloader event.
+CONST_DACP_TYPE = "dacp"
+# auto checkpoint by loop range.
+CONST_ACP_TYPE = "acp"
+g_acp_type = None
+g_program_attr = {}  # program_name->can_be_auto_checkpoint
+
+
+def _get_logger(log_level, name="auto_checkpoint"):
+    global logger
+    if logger is not None:
+        return logger
+
+    logger = logging.getLogger(name)
+    logger.setLevel(log_level)
+    logger.propagate = False
+
+    log_handler = logging.StreamHandler()
+    log_format = logging.Formatter(
+        '%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
+    )
+    log_handler.setFormatter(log_format)
+    logger.addHandler(log_handler)
+
+    return logger
+
+
+def _thread_checker():
+    assert current_thread().name == "MainThread", (
+        "auto checkpoint must run under main thread"
+    )
+
+
+class AutoCheckpointChecker:
+    def __init__(self):
+        self._run_env = None
+        self._platform = None
+        self._job_id = None
+        self._hdfs_home = None
+        self._hdfs_name = None
+        self._hdfs_ugi = None
+        self._hdfs_checkpoint_path = None
+        self._trainer_id = None
+        self._ce_test = None
+
+        self._run_env = os.getenv("PADDLE_RUNNING_ENV")
+        if self._run_env != "PADDLE_EDL_AUTO_CHECKPOINT":
+            return
+
+        try:
+            self._platform = os.environ["PADDLE_RUNNING_PLATFORM"]
+            self._job_id = os.environ["PADDLE_JOB_ID"]
+            self._hdfs_home = os.environ["PADDLE_EDL_HDFS_HOME"]
+            self._hdfs_name = os.environ["PADDLE_EDL_HDFS_NAME"]
+            self._hdfs_ugi = os.environ["PADDLE_EDL_HDFS_UGI"]
+            self._hdfs_checkpoint_path = os.environ[
+                "PADDLE_EDL_HDFS_CHECKPOINT_PATH"
+            ]
+            self._trainer_id = int(os.environ["PADDLE_TRAINER_ID"])
+
+            self._ce_test = int(os.getenv("PADDLE_EDL_ONLY_FOR_CE_TEST", "0"))
+            self._fs_cache = os.getenv("PADDLE_EDL_FS_CACHE", ".cache")
+
+            self._save_checkpoint_inter = int(
+                os.getenv("PADDLE_EDL_SAVE_CHECKPOINT_INTER", "900")
+            )  # s
+
+            if not self._ce_test:
+                assert (
+                    len(self._hdfs_home) > 3
+                    and len(self._hdfs_name) > 6
+                    and len(self._hdfs_ugi) > 3
+                    and len(self._hdfs_checkpoint_path) > 0
+                ), "hdfs environ must set"
+            else:
+                assert (
+                    len(self._hdfs_home) > 3
+                    and len(self._hdfs_checkpoint_path) > 0
+                ), "hdfs environ must set"
+
+        except Exception as e:
+            logger.fatal(f"exception:{e}")
+            sys.exit(1)
+
+    def get_range_checkpoint_path(self, name):
+        return f"{self.hdfs_checkpoint_path}/{self.job_id}/range/{name}"
+
+    def get_exe_checkpoint_path(self, name):
+        return f"{self.hdfs_checkpoint_path}/{self.job_id}/exe/{name}"
+
+    def get_job_path(self):
+        return f"{self.hdfs_checkpoint_path}/{self.job_id}"
+
+    @property
+    def save_checkpoint_inter(self):
+        return self._save_checkpoint_inter
+
+    def valid(self):
+        if in_dygraph_mode():
+            return False
+
+        return (
+            self._run_env is not None
+            and self._platform is not None
+            and self._job_id is not None
+            and self._hdfs_home is not None
+            and self._hdfs_name is not None
+            and self._hdfs_ugi is not None
+            and self._hdfs_checkpoint_path is not None
+            and self._trainer_id is not None
+        )
+
+    def __str__(self):
+        return f"run_env:{self._run_env} platform:{self._platform} job_id:{self._hdfs_home} \
+            hdfs_home:{self._hdfs_name} hdfs_name:{self._hdfs_ugi} hdfs_ugi:{self._hdfs_checkpoint_path} \
+            hdfs_checkpoint_path:{self._trainer_id} trainer_id:{self._ce_test} ce_test"
+
+    @property
+    def trainer_id(self):
+        return self._trainer_id
+
+    @property
+    def run_env(self):
+        return self._run_env
+
+    @property
+    def platform(self):
+        return self._platform
+
+    @property
+    def job_id(self):
+        return self._job_id
+
+    @property
+    def hdfs_home(self):
+        return self._hdfs_home
+
+    @property
+    def hdfs_name(self):
+        return self._hdfs_name
+
+    @property
+    def ce_test(self):
+        return self._ce_test
+
+    @property
+    def hdfs_ugi(self):
+        return self._hdfs_ugi
+
+    @property
+    def hdfs_checkpoint_path(self):
+        return self._hdfs_checkpoint_path
+
+    @staticmethod
+    def generate_range_name():
+        return generator("_range_")
+
+
+class ExeTrainStatus(SerializableBase):
+    def __init__(self):
+        self._epoch_no = -1  # start epoch_no
+        self._hash_key = None
+        self._key = None
+        self._checkpoint_path = None
+        self._checkpoint_no = None
+        self._restored_from = None
+        self._exe = None
+        self._program = None
+        self._exe_name = None
+        self._program_name = None
+
+        self._file_name = "exe_train_status"
+
+    def __eq__(self, t):
+        return (
+            self._epoch_no == t._epoch_no
+            and self._hash_key == t._hash_key
+            and self._key == t._key
+            and self._checkpoint_path == t._checkpoint_path
+            and self._checkpoint_no == t._checkpoint_no
+            and self._exe_name == t._exe_name
+            and self._program_name == t._program_name
+        )
+
+    def __ne__(self, t):
+        return not self == t
+
+    def serialize(self, path):
+        file_name = f"{path}/{self._file_name}"
+        with open(file_name, 'w') as f:
+            s = self._serialize()
+            f.write(s)
+
+    def _serialize(self, pop_keys=["restored_from"]):
+        d = self._to_dict()
+        for k in pop_keys:
+            d.pop(k, None)
+        return json.dumps(d)
+
+    def deserialize(self, path):
+        d = None
+        file_name = f"{path}/{self._file_name}"
+        with open(file_name, 'r') as f:
+            s = f.read()
+            self._deserialize(s)
+
+    def _deserialize(self, s):
+        d = json.loads(s)
+        self._epoch_no = d["epoch_no"]
+        self._key = d["key"]
+        self._hash_key = d["hash_key"]
+        self._checkpoint_path = d["checkpoint_path"]
+        self._checkpoint_no = d["checkpoint_no"]
+        self._exe_name = d["exe_name"]
+        self._program_name = d["program_name"]
+
+    def _to_dict(self):
+        return {
+            "epoch_no": self._epoch_no,
+            "key": self._key,
+            "hash_key": self._hash_key,
+            "checkpoint_path": self._checkpoint_path,
+            "restored_from": self._restored_from,
+            "exe_name": self._exe_name,
+            "program_name": self._program_name,
+            "checkpoint_no": self._checkpoint_no,
+        }
+
+    def __str__(self):
+        return self._serialize([])
+
+
+class TrainEpochRange(SerializableBase):
+    def __init__(
+        self, max_epoch_num, name, checkpoint_inter=None, restored=True
+    ):
+        self._max_epoch_num = max_epoch_num
+        self._epoch_no = -1  # current epoch_no
+        self._name = name
+        self._restored_from = None
+        self._exe_status = {}
+        self._flag_generated = False
+
+        self._checker = g_checker
+        if checkpoint_inter is not None:
+            self._save_checkpoint_inter = checkpoint_inter
+        else:
+            self._save_checkpoint_inter = self._checker.save_checkpoint_inter
+        assert self._save_checkpoint_inter >= 0, (
+            f"checkpoint inter:{self._save_checkpoint_inter} must >=0"
+        )
+        self._last_checkpoint_time = time.time()
+
+        self._load_cp_nos = None
+        self._checkpoint_epoch_no = None
+
+        if not self._checker.valid():
+            return
+
+        self._file_name = "range_train_status"
+
+        if not restored:
+            return
+
+        self._checkpoint_path = self._checker.get_range_checkpoint_path(name)
+
+        config = {
+            "fs.default.name": self._checker.hdfs_name,
+            "hadoop.job.ugi": self._checker.hdfs_ugi,
+        }
+
+        if self._checker.ce_test:
+            config = None
+
+        from paddle.distributed.fleet.utils.fs import HDFSClient
+
+        self._hdfs = HDFSClient(self._checker.hdfs_home, config)
+
+        self._cper = CheckpointSaver(self._hdfs)
+
+        _thread_checker()
+
+        self._get_last_valid_checkpoint()
+
+    def _look_for_valid(self, cp_nos):
+        cps = []
+        epoch_no = -1
+        for i in cp_nos[::-1]:
+            t = TrainEpochRange(self._max_epoch_num, self.name, restored=False)
+            self._cper.load_checkpoint(
+                self._checkpoint_path,
+                [t],
+                self._checker.trainer_id,
+                checkpoint_no=i,
+                local_cache_path=self._checker._fs_cache,
+            )
+            cps.append(t)
+            logger.debug(f"look for valid:{i} t:{t._serialize()}")
+            if epoch_no < 0:
+                epoch_no = t._epoch_no
+            else:
+                if epoch_no - t._epoch_no >= 1:
+                    return t, i
+        return None, None
+
+    def _get_last_valid_checkpoint(self):
+        self._load_cp_nos = self._cper.get_checkpoint_no(self._checkpoint_path)
+        logger.info(f"find checkpoint nos:{self._load_cp_nos}")
+
+        if len(self._load_cp_nos) < 1:
+            self._restored_from = CONST_MEMORYINIT
+            return
+
+        if g_acp_type == CONST_ACP_TYPE:
+            # get the last one
+            self._cper.load_checkpoint(
+                self._checkpoint_path,
+                [self],
+                self._checker.trainer_id,
+                local_cache_path=self._checker._fs_cache,
+            )
+            self._restored_from = CONST_CHECKPOINT
+            self._checkpoint_epoch_no = self._epoch_no
+
+            logger.info(f"load tain_epoch_range checkpoint:{self._serialize()}")
+
+        elif g_acp_type == CONST_DACP_TYPE:
+            t, i = self._look_for_valid(self._load_cp_nos)
+            if t is None:
+                self._restored_from = CONST_MEMORYINIT
+                return
+
+            self._cper.load_checkpoint(
+                self._checkpoint_path,
+                [self],
+                self._checker.trainer_id,
+                checkpoint_no=i,
+                local_cache_path=self._checker._fs_cache,
+            )
+
+            self._restored_from = CONST_CHECKPOINT
+            self._checkpoint_epoch_no = self._epoch_no
+            logger.info(f"load tain_epoch_range checkpoint:{self._serialize()}")
+        else:
+            raise AssertionError(f"not supported acp_type:{g_acp_type}")
+
+    def _to_dict(self):
+        d = {
+            "max_epoch_num": self._max_epoch_num,
+            "epoch_no": self._epoch_no,
+            "name": self._name,
+            "checkpoint_path": self._checkpoint_path,
+            "restored_from": self._restored_from,
+            "checkpoint_epoch_no": self._checkpoint_epoch_no,
+        }
+        return d
+
+    def __str__(self):
+        return self._serialize([])
+
+    @property
+    def name(self):
+        return self._name
+
+    def serialize(self, path):
+        file_name = f"{path}/{self._file_name}"
+        with open(file_name, 'w') as f:
+            s = self._serialize()
+            f.write(s)
+
+    def _serialize(self, pop_keys=["restored_from", "checkpoint_epoch_no"]):
+        # self
+        d = self._to_dict()
+        for k in pop_keys:
+            d.pop(k, None)
+
+        # registered exes
+        d["exe_status"] = {}
+        e = d["exe_status"]
+        for k, t in self._exe_status.items():
+            e[t._key] = t._serialize()
+        return json.dumps(d)
+
+    @property
+    def restored_from(self):
+        return self._restored_from
+
+    def deserialize(self, path):
+        d = None
+        file_name = f"{path}/{self._file_name}"
+        with open(file_name, 'r') as f:
+            d = json.load(f)
+
+        # self
+        self._max_epoch_num = d["max_epoch_num"]
+        self._epoch_no = d["epoch_no"]
+        self._name = d["name"]
+        self._checkpoint_path = d["checkpoint_path"]
+
+        # exes status
+        e = d["exe_status"]
+        for k, v in e.items():
+            t = ExeTrainStatus()
+            t._deserialize(v)
+            self._exe_status[k] = t
+
+    def next(self):
+        _thread_checker()
+
+        if self._max_epoch_num < 0:
+            self._max_epoch_num = sys.maxsize
+
+        assert self._epoch_no >= -1, (
+            f"self._epoch_no:{self._epoch_no} must >=-1"
+        )
+
+        self._last_checkpoint_time = time.time()
+        start = self._epoch_no + 1
+        logger.info(
+            f"started epoch_no:{start} max_epoch_num:{self._max_epoch_num}"
+        )
+
+        for i in range(start, self._max_epoch_num):
+            self._epoch_no = i
+            yield i
+
+            self.save_checkpoint()
+
+    def get(self):
+        return self._epoch_no
+
+    def save_checkpoint(self):
+        # not save last one because exe and program can't be restored.
+        if self._checker.trainer_id == 0:
+            if (
+                time.time() - self._last_checkpoint_time
+                >= self._save_checkpoint_inter
+            ):
+                if g_acp_type == CONST_ACP_TYPE:
+                    # not save the last one
+                    if (
+                        self._max_epoch_num > 0
+                        and self._epoch_no != self._max_epoch_num - 1
+                    ):
+                        self._save_checkpoint()
+                elif g_acp_type == CONST_DACP_TYPE:
+                    self._save_checkpoint()
+                else:
+                    raise AssertionError("not supported acp_type:{g_acp_type}")
+            self._last_checkpoint_time = time.time()
+
+    def _save_checkpoint(self):
+        """
+        status => /jobid/xxx_range_xx/range/
+        model =>                       /exe/
+        """
+        if not self._checker.valid():
+            return
+
+        e = self._exe_status
+        for k, t in self._exe_status.items():
+            m = PaddleModel(t._exe, t._program)
+            p = self._checker.get_exe_checkpoint_path(t._hash_key)
+            t._epoch_no = self.get()
+            path, checkpoint_no = self._cper.save_checkpoint(
+                p,
+                [m],
+                self._checker.trainer_id,
+                local_cache_path=self._checker._fs_cache,
+            )
+            # index info
+            t._checkpoint_path = path
+            t._checkpoint_no = checkpoint_no
+
+            e[t._key] = t
+
+            logger.debug(f"save executor checkpoint:{t._serialize()}")
+
+        if len(self._exe_status) > 0:
+            self._cper.save_checkpoint(
+                self._checkpoint_path,
+                [self],
+                local_cache_path=self._checker._fs_cache,
+            )
+            logger.info(
+                f"save train_epoch_range checkpoint:{self._serialize()}"
+            )
+
+            self._generate_flag()
+
+    def _generate_flag(self):
+        if self._flag_generated:
+            return
+
+        name = "can_be_auto_checkpoint.flag"
+        path = self._checker.get_job_path() + "/" + name
+        logger.info("this job can_be_auto_checkpoint")
+        self._hdfs.mkdirs(self._checker.get_job_path())
+        self._hdfs.touch(path, exist_ok=True)
+
+        self._flag_generated = True
+
+
+def _get_train_epoch_range():
+    return g_train_epoch_range
+
+
+def _check_program_oprole(program):
+    global_block = program.global_block()
+    has_backward = False
+    has_opt = False
+    for idx, op in enumerate(global_block.ops):
+        if op._is_backward_op():
+            has_backward = True
+
+        if op._is_optimize_op():
+            has_opt = True
+
+        if has_backward and has_opt:
+            return True
+
+    return False
+
+
+def _can_auto_checkpoint(prog):
+    if not isinstance(prog, compiler.CompiledProgram) and not isinstance(
+        prog, Program
+    ):
+        return False
+
+    if isinstance(prog, compiler.CompiledProgram):
+        if prog._program is None or prog._program._is_distributed:
+            return False
+    else:
+        if prog._is_distributed:
+            return False
+
+    program = _get_valid_program(prog)
+
+    if program._auto_checkpoint_name in g_program_attr:
+        if not g_program_attr[program._auto_checkpoint_name]:
+            return False
+    else:
+        ret = False
+        if isinstance(program, compiler.CompiledProgram):
+            ret = _check_program_oprole(program._program)
+        else:
+            ret = _check_program_oprole(program)
+
+        g_program_attr[program._auto_checkpoint_name] = ret
+        if not ret:
+            logger.debug(
+                f"program {program._auto_checkpoint_name} need't to auto checkpoint"
+            )
+            return False
+
+    return g_checker.valid() and g_train_epoch_range is not None
+
+
+def _get_running_key(exe_name, program_name):
+    return f"{exe_name}_{program_name}"
+
+
+def _get_checker():
+    _get_logger(20)
+    global g_checker
+    if g_checker is None:
+        g_checker = AutoCheckpointChecker()
+
+    return g_checker
+
+
+def _normal_yield(max_epoch_num):
+    if max_epoch_num < 0:
+        max_epoch_num = sys.maxsize
+    yield from range(0, max_epoch_num)
+
+
+def train_epoch_range(max_epoch_num, save_checkpoint_inter=None):
+    global g_acp_type
+    if not _get_checker().valid():
+        logger.warning(
+            "auto checkpoint will take effect automatically on PaddleCloud"
+        )
+        for i in _normal_yield(max_epoch_num):
+            yield i
+
+        return
+
+    if g_acp_type == CONST_DACP_TYPE:
+        for i in _normal_yield(max_epoch_num):
+            yield i
+
+        return
+
+    g_acp_type = CONST_ACP_TYPE
+    logger.info(f"acp_type:{g_acp_type}")
+
+    global g_train_epoch_range
+    try:
+        g_train_epoch_range = TrainEpochRange(
+            max_epoch_num,
+            g_checker.generate_range_name(),
+            checkpoint_inter=save_checkpoint_inter,
+        )
+
+        for i in g_train_epoch_range.next():
+            yield i
+    finally:
+        g_train_epoch_range = None
+
+
+def _get_valid_program(prog):
+    if isinstance(prog, compiler.CompiledProgram):
+        return prog._program
+
+    return prog
+
+
+def _auto_checkpoint(exe, prog):
+    _get_checker()
+
+    assert exe._auto_checkpoint_name is not None
+    if not _can_auto_checkpoint(prog):
+        return
+
+    program = _get_valid_program(prog)
+    assert program._auto_checkpoint_name is not None
+
+    exe_status = g_train_epoch_range._exe_status
+    key = _get_running_key(
+        exe._auto_checkpoint_name, program._auto_checkpoint_name
+    )
+
+    if g_train_epoch_range.restored_from == CONST_CHECKPOINT:
+        assert key in exe_status, (
+            f"when restored key:{key} must be in train_epoch_range:{g_train_epoch_range}"
+        )
+
+    t = None
+    if key in exe_status:
+        t = exe_status[key]
+        if t._restored_from is None:
+            a = CheckpointSaver(g_train_epoch_range._hdfs)
+            m = PaddleModel(exe, program)
+            a.load_checkpoint(
+                g_checker.get_exe_checkpoint_path(key),
+                [m],
+                trainer_id=g_checker.trainer_id,
+                checkpoint_no=t._checkpoint_no,
+                local_cache_path=g_checker._fs_cache,
+            )
+            t._restored_from = CONST_CHECKPOINT
+            logger.info(f"load executor checkpoint {t}")
+        t._exe = exe
+        t._program = program
+        t._epoch_no = g_train_epoch_range.get()
+    else:
+        t = ExeTrainStatus()
+        t._epoch_no = g_train_epoch_range.get()
+        t._hash_key = key
+        t._key = key
+        t._restored_from = CONST_MEMORYINIT
+        t._exe = exe
+        t._program = program
+        t._exe_name = exe._auto_checkpoint_name
+        t._program_name = program._auto_checkpoint_name
+
+        # register this <exe,program,io>
+        exe_status[key] = t
+
+        logger.info("not found checkpoint, so train from epoch 0")
+
+    _thread_checker()

python/paddle/base/incubate/checkpoint/checkpoint_saver.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...compiler import CompiledProgram
+
+
+class SerializableBase:
+    def serialize(self, path):
+        raise NotImplementedError
+
+    def deserialize(self, path):
+        raise NotImplementedError
+
+
+class PaddleModel(SerializableBase):
+    def __init__(self, exe, program):
+        self._exe = exe
+        self._origin_program = program
+        self._program = program
+        if isinstance(program, CompiledProgram):
+            self._program = program._program
+
+        self._file_name = "_paddle_fleet_param__"
+
+    def serialize(self, path):
+        from paddle.distributed.io import save_persistables
+
+        save_persistables(
+            executor=self._exe,
+            dirname=path,
+            main_program=self._program,
+            filename=self._file_name,
+        )
+
+    def deserialize(self, path):
+        from paddle.distributed.io import load_persistables
+
+        load_persistables(
+            executor=self._exe,
+            dirname=path,
+            main_program=self._program,
+            filename=self._file_name,
+        )
+
+
+class CheckpointSaver:
+    def __init__(self, fs):
+        self._fs = fs
+        self._checkpoint_prefix = "__paddle_checkpoint__"
+
+    def save_checkpoint(
+        self, path, slists, trainer_id=None, local_cache_path=".cache"
+    ):
+        """
+        Serialize objects in slists to path
+        Return really saved path and checkpoint_no
+        """
+        if not self._fs.is_exist(path):
+            self._fs.mkdirs(path)
+        else:
+            assert self._fs.is_dir(path), f"path:{path} must be a directory"
+
+        max_no = self._get_last_checkpoint_no(path)
+        if max_no < 0:
+            max_no = -1
+        max_no += 1
+
+        real_path = f"{path}/{self._checkpoint_prefix}.{max_no}"
+        tmp_path = f"{real_path}.tmp"
+        saved_path = tmp_path
+
+        from paddle.distributed.fleet.utils.fs import LocalFS
+
+        local_fs = LocalFS()
+
+        cache_path = None
+        if self._fs.need_upload_download():
+            cache_path = f"{local_cache_path}/{self._checkpoint_prefix}.{max_no}.saved_cache"
+
+            if trainer_id is not None:
+                cache_path = f"{cache_path}.{trainer_id}"
+
+            if not local_fs.is_exist(cache_path):
+                local_fs.mkdirs(cache_path)
+            else:
+                assert local_fs.is_dir(cache_path), (
+                    f"cache path:{cache_path} must be a directory"
+                )
+
+            saved_path = cache_path
+
+        for s in slists:
+            s.serialize(saved_path)
+
+        if self._fs.need_upload_download():
+            self._fs.delete(tmp_path)
+            self._fs.upload(cache_path, tmp_path)
+            local_fs.delete(cache_path)
+        self._fs.mv(tmp_path, real_path)
+
+        return real_path, max_no
+
+    def load_checkpoint(
+        self,
+        path,
+        slists,
+        trainer_id,
+        local_cache_path=".cache",
+        checkpoint_no=None,
+        ignore_empty=True,
+    ):
+        """
+        Deserialize objects in slists from path
+        Return really load path
+        """
+        if checkpoint_no is None:
+            max_no = self._get_last_checkpoint_no(path)
+
+            if not ignore_empty:
+                assert max_no >= 0, "Can't find checkpoint"
+
+            if max_no < 0:
+                return None
+
+            checkpoint_no = max_no
+        else:
+            assert isinstance(checkpoint_no, int)
+            assert checkpoint_no >= 0
+
+        from paddle.distributed.fleet.utils.fs import LocalFS
+
+        local_fs = LocalFS()
+        if self._fs.need_upload_download():
+            cache_path = f"{local_cache_path}/{self._checkpoint_prefix}.{checkpoint_no}.load_cache"
+
+            if trainer_id is not None:
+                cache_path = f"{cache_path}.{trainer_id}"
+
+            if not local_fs.is_exist(local_cache_path):
+                local_fs.mkdirs(local_cache_path)
+            if local_fs.is_exist(cache_path):
+                local_fs.delete(cache_path)
+
+        real_path = f"{path}/{self._checkpoint_prefix}.{checkpoint_no}"
+        load_path = real_path
+        if self._fs.need_upload_download():
+            self._fs.download(real_path, cache_path)
+            load_path = cache_path
+
+        for s in slists:
+            s.deserialize(load_path)
+
+        if self._fs.need_upload_download() and cache_path:
+            local_fs.delete(cache_path)
+
+        return real_path
+
+    def get_checkpoint_no(self, root_path):
+        a = []
+        dirs = self._fs.list_dirs(root_path)
+        for d in dirs:
+            g = d.split(".")
+            if len(g) != 2:
+                continue
+
+            if g[0] != self._checkpoint_prefix:
+                continue
+
+            try:
+                n = int(g[1])
+                a.append(n)
+            except:
+                continue
+
+        a.sort()
+        return a
+
+    def _get_last_checkpoint_no(self, root_path):
+        """
+        only get the first depth
+        """
+        a = self.get_checkpoint_no(root_path)
+        if len(a) > 0:
+            return a[-1]
+
+        return -1