Repository: sql-machine-learning/elasticdl

Issue:
Model export 



Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/common/model_helper.py b/elasticdl/python/elasticdl/common/model_helper.py
--- a/elasticdl/python/elasticdl/common/model_helper.py
+++ b/elasticdl/python/elasticdl/common/model_helper.py
@@ -1,5 +1,6 @@
 # common helper methods for model manipulation.
 import importlib.util
+from elasticdl.proto import elasticdl_pb2
 
 
 def load_user_model(model_file):
@@ -16,3 +17,15 @@ def build_model(model, feature_columns):
     else:
         input_shapes = [(1,) + tuple(f_col.shape) for f_col in feature_columns]
         model.build(input_shapes)
+
+def save_checkpoint_to_file(pb_model, file_name):
+    encoded_model = pb_model.SerializeToString()
+    with open(file_name, "wb") as f:
+        f.write(encoded_model)
+
+def load_from_checkpoint_file(file_name):
+    pb_model = elasticdl_pb2.Model()
+    with open(file_name, "rb") as f:
+        pb_model.ParseFromString(f.read())
+    return pb_model
+    
diff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py
--- a/elasticdl/python/elasticdl/master/main.py
+++ b/elasticdl/python/elasticdl/master/main.py
@@ -59,6 +59,23 @@ def _parse_args():
         help="the number of workers used in training",
         default=0,
     )
+    parser.add_argument(
+        "--checkpoint_dir",
+        help="The directory to store the checkpoint files",
+        default="",
+    )
+    parser.add_argument(
+        "--checkpoint_steps",
+        type=int,
+        help="Save checkpoint every this many steps. If 0, no checkpoints to save.",
+        default=0,
+    )
+    parser.add_argument(
+        "--keep_checkpoint_max",
+        type=int,
+        help="The maximum number of recent checkpoint files to keep. If 0, keep all.",
+        default=3,
+    )
     parser.add_argument(
         "--worker_cpu_request",
         help="the minimal cpu required by worker in training",
@@ -139,6 +156,9 @@ def main():
             optimizer,
             task_q,
             init_var=model_inst.trainable_variables,
+            checkpoint_dir=args.checkpoint_dir,
+            checkpoint_steps=args.checkpoint_steps,
+            keep_checkpoint_max=args.keep_checkpoint_max,
         ),
         server,
     )
diff --git a/elasticdl/python/elasticdl/master/servicer.py b/elasticdl/python/elasticdl/master/servicer.py
--- a/elasticdl/python/elasticdl/master/servicer.py
+++ b/elasticdl/python/elasticdl/master/servicer.py
@@ -1,4 +1,5 @@
 import logging
+import os
 import threading
 import numpy as np
 
@@ -11,6 +12,7 @@
 from elasticdl.proto import elasticdl_pb2
 from elasticdl.proto import elasticdl_pb2_grpc
 from elasticdl.python.elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray
+from elasticdl.python.elasticdl.common.model_helper import save_checkpoint_to_file, load_from_checkpoint_file
 
 
 class MasterServicer(elasticdl_pb2_grpc.MasterServicer):
@@ -23,7 +25,10 @@ def __init__(
         optimizer,
         task_q,
         *,
-        init_var=[]
+        init_var,
+        checkpoint_dir,
+        checkpoint_steps,
+        keep_checkpoint_max
     ):
         # TODO: group params together into a single object.
         self._logger = logging.getLogger(__name__)
@@ -42,6 +47,15 @@ def __init__(
         self._evaluation_metrics = {}
         for var in init_var:
             self.set_model_var(var.name, var.numpy())
+        self._checkpoint_dir = checkpoint_dir
+        self._checkpoint_steps = checkpoint_steps
+        self._keep_checkpoint_max = keep_checkpoint_max
+        if not self._checkpoint_dir:
+            self._checkpoint_dir = os.getcwd() + "/checkpoint_dir"
+        if self._checkpoint_steps:
+            os.makedirs(self._checkpoint_dir, exist_ok=True)
+            if self._keep_checkpoint_max > 0:
+                self._checkpoint_list = []
 
     def set_model_var(self, name, value):
         """Add or set model variable. Value should be a float32 ndarray"""
@@ -70,17 +84,33 @@ def GetTask(self, request, _):
     def GetModel(self, request, _):
         _ = self._validate_model_version(request.min_version)
 
-        res = elasticdl_pb2.Model()
         with self._lock:
-            res.version = self._version
-            for k, v in self._model.items():
-                res.param[k].CopyFrom(ndarray_to_tensor(v.numpy()))
+            res = self._get_model_no_lock()
         return res
 
     def _update_model_version(self):
         assert self._lock.locked()
         self._version += 1
 
+    def save_checkpoint(self):
+        file_name = "{}/model_v{}.chkpt".format(self._checkpoint_dir, self._version)
+        pb_model = self._get_model_no_lock()
+        save_checkpoint_to_file(pb_model, file_name)
+        if self._keep_checkpoint_max:
+            self._checkpoint_list.append(file_name)
+            while len(self._checkpoint_list) > self._keep_checkpoint_max:
+                file_to_delete = self._checkpoint_list.pop(0)
+                os.remove(file_to_delete)
+
+    def load_checkpoint_file(self, file_name):
+        pb_model = load_from_checkpoint_file(file_name)
+
+        for k, v in self._model.items():
+            # Assumes all variables exist in pb_model.param.
+            v.assign(
+                tensor_to_ndarray(pb_model.param[k]))
+        self._model_version = pb_model.version
+
     def _update_model(self):
         assert self._lock.locked()
         grad_var = []
@@ -92,6 +122,13 @@ def _update_model(self):
         self._gradient_sum.clear()
         self._grad_n = 0
 
+    def _get_model_no_lock(self):
+        pb_model = elasticdl_pb2.Model()
+        pb_model.version = self._version
+        for k, v in self._model.items():
+            pb_model.param[k].CopyFrom(ndarray_to_tensor(v.numpy()))
+        return pb_model
+
     def _validate_model_version(self, request_model_version):
         if request_model_version > self._version:
             err_msg = "Model version %d not available yet, current version: %d" % (
@@ -143,6 +180,15 @@ def ReportGradient(self, request, _):
             self._grad_n += 1
             if self._grad_n >= self._grad_to_wait:
                 self._update_model()
+                if self._checkpoint_steps and self._version % self._checkpoint_steps == 0:
+                    try:
+                        self.save_checkpoint()
+                    except:
+                        self._logger.warning(
+                            "Failed to save checkpoint file for model version {}".format(self._version)
+                        )
+
+
         res.accepted = True
         res.model_version = self._version
         return res


Test patch:
diff --git a/elasticdl/python/tests/checkpoint_test.py b/elasticdl/python/tests/checkpoint_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/checkpoint_test.py
@@ -0,0 +1,146 @@
+import random
+import os
+import mock
+import grpc
+import tempfile
+import unittest
+import numpy as np
+import recordio
+
+import tensorflow as tf
+
+tf.enable_eager_execution()
+
+from contextlib import closing
+from elasticdl.python.elasticdl.master.servicer import MasterServicer
+from elasticdl.python.elasticdl.worker.worker import Worker
+from elasticdl.python.elasticdl.common.model_helper import load_user_model
+from elasticdl.python.elasticdl.master.task_queue import _TaskQueue
+from elasticdl.python.elasticdl.common.model_helper import save_checkpoint_to_file, load_from_checkpoint_file
+from elasticdl.proto import elasticdl_pb2
+from elasticdl.python.data.codec import BytesCodec
+
+
+_module_file = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "test_module.py"
+)
+
+m = load_user_model(_module_file)
+columns = m.feature_columns() + m.label_columns()
+
+
+def create_recordio_file(size, codec_type):
+    codec = None
+    if codec_type == "bytes":
+        codec = BytesCodec(columns)
+    elif codec_type == "tf_example":
+        codec = TFExampleCodec(columns)
+
+    temp_file = tempfile.NamedTemporaryFile(delete=False)
+    with closing(recordio.Writer(temp_file.name)) as f:
+        for _ in range(size):
+            x = np.random.rand((1)).astype(np.float32)
+            y = 2 * x + 1
+            f.write(codec.encode({"x": x, "y": y}))
+    return temp_file.name
+
+
+class CheckpointTest(unittest.TestCase):
+    def testSaveLoadCheckpoint(self):
+        master = MasterServicer(
+            2, 3, None, None,
+            init_var=[],
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0
+        )
+
+        model_inst = m.model
+        for variable in model_inst.trainable_variables:
+            master.set_model_var(variable.name, variable.numpy())
+
+        req = elasticdl_pb2.GetModelRequest()
+        req.min_version = 0
+        model = master.GetModel(req, None)
+
+        tmp_file = tempfile.NamedTemporaryFile()
+        save_checkpoint_to_file(model, tmp_file.name)
+
+        pb_model = load_from_checkpoint_file(tmp_file.name)
+
+        self.assertEqual(model.version, pb_model.version)
+        for k in model.param:
+            self.assertEqual(model.param[k], pb_model.param[k])
+
+    def testCheckpintArguments(self):
+        """
+        Run Worker.distributed_train with a local master.
+        grpc calls are mocked by local master call.
+        """
+
+        def mock_GetTask(req):
+            return master.GetTask(req, None)
+
+        def mock_GetModel(req):
+            return master.GetModel(req, None)
+
+        def mock_ReportGradient(req):
+            return master.ReportGradient(req, None)
+
+        def mock_ReportTaskResult(req):
+            return master.ReportTaskResult(req, None)
+
+        codec_type = "bytes"
+        channel = grpc.insecure_channel("localhost:9999")
+        worker = Worker(
+            1,
+            _module_file,
+            channel=channel,
+            codec_type=codec_type,
+        )
+
+        # save checkpoint file every 2 steps
+        # keep at most 5 recent checkpoint files
+        checkpoint_dir = tempfile.mkdtemp()
+        checkpoint_steps = 2
+        keep_checkpoint_max = 5
+
+        filename = create_recordio_file(128, codec_type)
+        task_q = _TaskQueue({filename: 128}, record_per_task=64, num_epoch=1)
+        master = MasterServicer(2,
+                                2,
+                                worker._opt_fn(),
+                                task_q,
+                                init_var=[],
+                                checkpoint_dir=checkpoint_dir,
+                                checkpoint_steps=checkpoint_steps,
+                                keep_checkpoint_max=keep_checkpoint_max
+                                )
+
+        for var in worker._model.trainable_variables:
+            master.set_model_var(var.name, var.numpy())
+
+        with mock.patch.object(
+            worker._stub, "GetTask", mock_GetTask
+        ), mock.patch.object(
+            worker._stub, "GetModel", mock_GetModel
+        ), mock.patch.object(
+            worker._stub, "ReportGradient", mock_ReportGradient
+        ), mock.patch.object(
+            worker._stub, "ReportTaskResult", mock_ReportTaskResult
+        ):
+            try:
+                worker.distributed_train()
+                res = True
+            except Exception as ex:
+                print(ex)
+                res = False
+
+        self.assertTrue(res)
+        checkpoint_files = sorted(os.listdir(checkpoint_dir))
+        self.assertEqual(checkpoint_files,
+                         ['model_v24.chkpt', 'model_v26.chkpt', 'model_v28.chkpt', 'model_v30.chkpt', 'model_v32.chkpt'])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/elasticdl/python/tests/servicer_test.py b/elasticdl/python/tests/servicer_test.py
--- a/elasticdl/python/tests/servicer_test.py
+++ b/elasticdl/python/tests/servicer_test.py
@@ -45,6 +45,10 @@ def testGetEmptyTask(self):
             3,
             None,
             _TaskQueue({}, record_per_task=3, num_epoch=2),
+            init_var=[],
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0
         )
 
         req = elasticdl_pb2.GetTaskRequest()
@@ -61,7 +65,12 @@ def testGetEmptyTask(self):
         self.assertEqual(1, task.model_version)
 
     def testGetModel(self):
-        master = MasterServicer(2, 3, None, None)
+        master = MasterServicer(2, 3, None, None,
+                                init_var=[],
+                                checkpoint_dir="",
+                                checkpoint_steps=0,
+                                keep_checkpoint_max=0
+        )
         req = elasticdl_pb2.GetModelRequest()
         req.min_version = 0
 
@@ -110,6 +119,10 @@ def makeGrad():
             3,
             tf.train.GradientDescentOptimizer(0.1),
             None,
+            init_var=[],
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0
         )
         master._version = 1
         master.set_model_var("x", np.array([2.0], dtype=np.float32))
@@ -196,6 +209,10 @@ def makeEvaluationMetrics():
             3,
             tf.train.GradientDescentOptimizer(0.1),
             None,
+            init_var=[],
+            checkpoint_dir="",
+            checkpoint_steps=0,
+            keep_checkpoint_max=0
         )
         master._version = 1
 
@@ -225,7 +242,12 @@ def testReportTaskResult(self):
         task_q = _TaskQueue(
             {"shard_1": 10, "shard_2": 9}, record_per_task=3, num_epoch=2
         )
-        master = MasterServicer(3, 3, None, task_q)
+        master = MasterServicer(3, 3, None, task_q,
+                                init_var=[],
+                                checkpoint_dir="",
+                                checkpoint_steps=0,
+                                keep_checkpoint_max=0
+        )
 
         # task to number of runs.
         tasks = defaultdict(int)
@@ -259,7 +281,12 @@ def testReportTaskResult(self):
         )
 
     def testUserDefinedModel(self):
-        master = MasterServicer(2, 3, None, None)
+        master = MasterServicer(2, 3, None, None,
+                                init_var=[],
+                                checkpoint_dir="",
+                                checkpoint_steps=0,
+                                keep_checkpoint_max=0
+        )
         req = elasticdl_pb2.GetModelRequest()
         req.min_version = 0
 
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -96,7 +96,11 @@ def mock_ReportTaskResult(req):
 
         filename = create_recordio_file(128, codec_type)
         task_q = _TaskQueue({filename: 128}, record_per_task=64, num_epoch=1)
-        master = MasterServicer(2, 16, worker._opt_fn(), task_q)
+        master = MasterServicer(2, 16, worker._opt_fn(), task_q,
+                                init_var=[],
+                                checkpoint_dir="",
+                                checkpoint_steps=0,
+                                keep_checkpoint_max=0)
 
         for var in worker._model.trainable_variables:
             master.set_model_var(var.name, var.numpy())