add checkpoint util class and implement

paddle/fluid/operators/listen_and_serv_op.cc

@@ -187,6 +187,7 @@ void ListenAndServOp::RunSyncLoop(framework::Executor *executor,
     for (auto &var : sparse_vars) {
       var->GetMutable<framework::SelectedRows>()->mutable_rows()->clear();
     }
+
     rpc_service_->SetCond(1);
     // FIXME(typhoonzero): use another condition to sync wait clients get.
     rpc_service_->WaitClientGet(fan_in);

python/paddle/fluid/io.py

@@ -13,21 +13,17 @@
 # limitations under the License.
 
 import os
+import time
+import shutil
 
 from paddle.fluid.evaluator import Evaluator
 from paddle.fluid.framework import Program, Parameter, default_main_program, Variable
 from . import core
 
 __all__ = [
-    'save_vars',
-    'save_params',
-    'save_persistables',
-    'load_vars',
-    'load_params',
-    'load_persistables',
-    'save_inference_model',
-    'load_inference_model',
-    'get_inference_program',
+    'save_vars', 'save_params', 'save_persistables', 'load_vars', 'load_params',
+    'load_persistables', 'save_inference_model', 'load_inference_model',
+    'get_inference_program', 'save_checkpoint', 'restore_checkpoint'
 ]
 
 
@@ -195,6 +191,8 @@ def load_vars(executor,
         load_var_map = {}
         for each_var in vars:
             assert isinstance(each_var, Variable)
+            if each_var.type == core.VarDesc.VarType.RAW:
+                continue
             new_var = _clone_var_in_block_(load_block, each_var)
             if filename is None:
                 load_block.append_op(
@@ -454,3 +452,149 @@ def get_parameter_value_by_name(name, executor, program=None):
         program = default_main_program()
     var = program.global_block().var(name)
     return get_parameter_value(var, executor)
+
+
+SUCCESS = "_SUCCESS"
+BEGIN_SECS = None
+
+
+def save_checkpoint(executor,
+                    dirname,
+                    keep_max=3,
+                    save_secs=600,
+                    main_program=None):
+    """
+    Save Variables to Checkpint Dir
+
+    :param dirname
+    :param keep_max
+    :param save_secs
+    :param main_program
+    """
+    if dirname is None:
+        raise Exception("save checkpoint dir can not be none")
+
+    if not os.path.isdir(dirname):
+        os.makedirs(dirname)
+
+    global BEGIN_SECS
+    if BEGIN_SECS is not None:
+        if time.time() - BEGIN_SECS < save_secs:
+            return
+    BEGIN_SECS = time.time()
+
+    serial = _get_lastest_checkpoint_dir(dirname) + 1
+    cur_dir = os.path.join(dirname, str(serial))
+    # save_persistables(executor, cur_dir, main_program)
+    save_vars(
+        executor,
+        dirname=cur_dir,
+        main_program=main_program,
+        vars=None,
+        predicate=is_checkpoint_var,
+        filename=None)
+    _write_success(cur_dir)
+    _lru_delete(dirname, keep_max)
+
+
+def restore_checkpoint(dirname, executor, main_program=None):
+    """
+    Load Variables from Checkpint Dir
+
+    :param dirname
+    :param executor
+    :param main_program
+    """
+    if dirname is None and os.path.isdir(dirname):
+        raise Exception("restore checkpoint can not load variables from %s" %
+                        dirname)
+    serial = _get_lastest_checkpoint_dir(dirname)
+
+    if serial < 0:
+        return
+    cur_dir = os.path.join(dirname, str(serial))
+    # load_persistables(executor, cur_dir, main_program)
+    load_vars(
+        executor,
+        dirname=cur_dir,
+        main_program=main_program,
+        predicate=is_checkpoint_var,
+        filename=None)
+
+
+def is_checkpoint_var(var):
+    if var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH or \
+            var.desc.type() == core.VarDesc.VarType.FETCH_LIST or \
+            var.desc.type() == core.VarDesc.VarType.RAW:
+        return False
+
+    if var.name.endswith("@GRAD"):
+        return False
+
+    return var.persistable
+
+
+def _lru_delete(dirname, keep_max=3):
+    """
+    retain checkpoint nums with keep_max
+    """
+    dirs = os.listdir(dirname)
+    serials = []
+    for serial in dirs:
+        try:
+            serials.append(int(serial))
+        except ValueError:
+            continue
+
+    if len(serials) <= keep_max:
+        return
+
+    serials.sort(reverse=True)
+    serials = serials[keep_max:]
+    for serial in serials:
+        cur_dir = os.path.join(dirname, str(serial))
+        shutil.rmtree(cur_dir)
+
+
+def _write_success(dirname):
+    """
+    write _SUCCESS to checkpoint dir
+    """
+    success_file = os.path.join(dirname, SUCCESS)
+    with open(success_file, 'a'):
+        pass
+
+
+def _get_lastest_checkpoint_dir(checkpoint_dir):
+    """
+    get the biggest number in checkpoint_dir, which has _SUCCESS
+    """
+    if not checkpoint_dir.strip():
+        return -1
+
+    def has_success(checkpoint_dir, cur_dir):
+        """
+        is _SUCCESS in this dir
+        """
+        if not os.path.isdir(os.path.join(checkpoint_dir, cur_dir)):
+            return -1
+
+        try:
+            int(cur_dir)
+        except ValueError:
+            return -1
+
+        success_path = os.path.join(checkpoint_dir, cur_dir, SUCCESS)
+        if os.path.isfile(success_path):
+            return int(cur_dir)
+
+    if not os.path.isdir(checkpoint_dir):
+        return -1
+
+    current_dir = -1
+    dirs = os.listdir(checkpoint_dir)
+    for cur_dir in dirs:
+        success_num = has_success(checkpoint_dir, cur_dir)
+        if success_num > current_dir:
+            current_dir = success_num
+    return current_dir