Repository: sql-machine-learning/elasticdl

Issue:
generate test data automatically
Now, some of our unit tests rely on some handcraft [testdata](https://github.com/sql-machine-learning/elasticdl/tree/develop/elasticdl/python/tests/testdata).

For example, there is an evaluation test in travis ci:

https://github.com/sql-machine-learning/elasticdl/blob/1f4ee6599942c897f1494673c0db63175e7f4708/scripts/client_test.sh#L35-L57

If some of the proto messages in elasticdl.proto are changed,  the testdata will be invalid. So the unit test will fail. The handcraft test data is hard to maintain.

We have to figure out a way to generate these test data automatically.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/task_dispatcher.py b/elasticdl/python/master/task_dispatcher.py
--- a/elasticdl/python/master/task_dispatcher.py
+++ b/elasticdl/python/master/task_dispatcher.py
@@ -190,6 +190,8 @@ def _create_train_end_callback_task(self):
         SavedModelExporter to execute on_train_end,we include
         a shard of data in this task.
         """
+        if not self._training_shards:
+            return
 
         self.reset_job_counters(elasticdl_pb2.TRAIN_END_CALLBACK)
         shards = self._training_shards
diff --git a/scripts/gen_mnist_checkpoint.py b/scripts/gen_mnist_checkpoint.py
new file mode 100644
--- /dev/null
+++ b/scripts/gen_mnist_checkpoint.py
@@ -0,0 +1,37 @@
+import argparse
+
+import tensorflow as tf
+
+from elasticdl.python.tests.test_utils import save_checkpoint_without_embedding
+
+
+def mnist_custom_model():
+    inputs = tf.keras.Input(shape=(28, 28), name="image")
+    x = tf.keras.layers.Reshape((28, 28, 1))(inputs)
+    x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu")(x)
+    x = tf.keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu")(x)
+    x = tf.keras.layers.BatchNormalization()(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Dropout(0.25)(x)
+    x = tf.keras.layers.Flatten()(x)
+    outputs = tf.keras.layers.Dense(10)(x)
+
+    return tf.keras.Model(inputs=inputs, outputs=outputs, name="mnist_model")
+
+
+def add_params(parser):
+    parser.add_argument(
+        "--checkpoint_dir",
+        help="The directory to store the mnist checkpoint",
+        default="",
+        type=str,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    add_params(parser)
+    args, _ = parser.parse_known_args()
+    print(args)
+    model = mnist_custom_model()
+    save_checkpoint_without_embedding(model, args.checkpoint_dir)


Test patch:
diff --git a/elasticdl/python/tests/callbacks_test.py b/elasticdl/python/tests/callbacks_test.py
--- a/elasticdl/python/tests/callbacks_test.py
+++ b/elasticdl/python/tests/callbacks_test.py
@@ -12,6 +12,7 @@
 from elasticdl.python.elasticdl.callbacks.saved_model_exporter import (
     SavedModelExporter,
 )
+from elasticdl.python.tests.test_utils import save_checkpoint_without_embedding
 from elasticdl.python.worker.task_data_service import TaskDataService
 
 
@@ -50,17 +51,20 @@ def test_on_train_end(self):
             elasticdl_pb2.TRAIN_END_CALLBACK,
         )
         task_data_service.get_dataset_by_task = mock.Mock(return_value=dataset)
-        model_handler = ModelHandler.get_model_handler(
-            distribution_strategy=DistributionStrategy.PARAMETER_SERVER,
-            checkpoint_dir="elasticdl/python/tests/testdata/functional_ckpt/",
-        )
-        saved_model_exporter = SavedModelExporter(
-            task_data_service, dataset_fn, model_handler
-        )
+
         with tempfile.TemporaryDirectory() as temp_dir_name:
+            checkpoint_dir = os.path.join(temp_dir_name, "checkpoint")
+            model = custom_model_with_embedding_layer()
+            save_checkpoint_without_embedding(model, checkpoint_dir)
+            model_handler = ModelHandler.get_model_handler(
+                distribution_strategy=DistributionStrategy.PARAMETER_SERVER,
+                checkpoint_dir=checkpoint_dir,
+            )
+            saved_model_exporter = SavedModelExporter(
+                task_data_service, dataset_fn, model_handler
+            )
             saved_model_path = os.path.join(temp_dir_name, "test_exporter")
             params = {"batch_size": 10, "saved_model_path": saved_model_path}
-            model = custom_model_with_embedding_layer()
             saved_model_exporter.set_params(params)
             saved_model_exporter.set_model(model)
             saved_model_exporter.on_train_end()
diff --git a/elasticdl/python/tests/model_handler_test.py b/elasticdl/python/tests/model_handler_test.py
--- a/elasticdl/python/tests/model_handler_test.py
+++ b/elasticdl/python/tests/model_handler_test.py
@@ -1,3 +1,5 @@
+import os
+import tempfile
 import unittest
 
 import numpy as np
@@ -147,7 +149,7 @@ def setUp(self):
         tf.keras.backend.clear_session()
         self.model_handler = ModelHandler.get_model_handler(
             distribution_strategy=DistributionStrategy.PARAMETER_SERVER,
-            checkpoint_dir="elasticdl/python/tests/testdata/functional_ckpt/",
+            checkpoint_dir="",
         )
 
     def _save_model(self, is_subclass):
@@ -185,43 +187,48 @@ def test_get_model_with_embedding_column_to_train(self):
         )
 
     def test_get_model_to_export(self):
-        self._save_model(False)
-        model_inst = custom_model_with_embedding_layer()
-        train_model = self.model_handler.get_model_to_train(model_inst)
-        export_model = self.model_handler.get_model_to_export(
-            train_model, dataset=None
-        )
+        with tempfile.TemporaryDirectory() as temp_dir:
+            self.model_handler._checkpoint_dir = os.path.join(
+                temp_dir, "test_export"
+            )
+            self._save_model(False)
+            model_inst = custom_model_with_embedding_layer()
+            train_model = self.model_handler.get_model_to_train(model_inst)
+            export_model = self.model_handler.get_model_to_export(
+                train_model, dataset=None
+            )
 
-        test_data = tf.constant([0])
-        result = export_model.call(test_data).numpy()
-        self.assertEqual(result[0][0], 3.0)
+            test_data = tf.constant([0])
+            result = export_model.call(test_data).numpy()
+            self.assertEqual(result[0][0], 3.0)
 
     def test_get_subclass_model_to_export(self):
-        self.model_handler._checkpoint_dir = (
-            "elasticdl/python/tests/testdata/subclass_ckpt/"
-        )
-        self._save_model(True)
-
-        def _get_dataset():
-            dataset = tf.data.Dataset.from_tensor_slices(
-                np.random.randint(0, 10, (10, 4))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            self.model_handler._checkpoint_dir = os.path.join(
+                temp_dir, "test_export"
             )
-            dataset = dataset.batch(2)
-            return dataset
+            self._save_model(True)
 
-        model_inst = CustomModel()
-        dataset = _get_dataset()
+            def _get_dataset():
+                dataset = tf.data.Dataset.from_tensor_slices(
+                    np.random.randint(0, 10, (10, 4))
+                )
+                dataset = dataset.batch(2)
+                return dataset
 
-        train_model = self.model_handler.get_model_to_train(model_inst)
-        self.assertEqual(type(train_model.embedding), Embedding)
+            model_inst = CustomModel()
+            dataset = _get_dataset()
 
-        export_model = self.model_handler.get_model_to_export(
-            train_model, dataset=dataset
-        )
+            train_model = self.model_handler.get_model_to_train(model_inst)
+            self.assertEqual(type(train_model.embedding), Embedding)
 
-        test_data = tf.constant([0])
-        result = export_model.call(test_data).numpy()
-        self.assertEqual(result[0][0], 3.0)
+            export_model = self.model_handler.get_model_to_export(
+                train_model, dataset=dataset
+            )
+
+            test_data = tf.constant([0])
+            result = export_model.call(test_data).numpy()
+            self.assertEqual(result[0][0], 3.0)
 
     def test_get_model_with_sparse_to_train(self):
         model_inst = custom_model_with_sparse_embedding()
@@ -229,24 +236,29 @@ def test_get_model_with_sparse_to_train(self):
         self.assertEqual(type(model_inst.layers[1]), Embedding)
 
     def test_get_model_with_sparse_to_export(self):
-        self._save_model(False)
-        model_inst = custom_model_with_sparse_embedding()
-        train_model = self.model_handler.get_model_to_train(model_inst)
-
-        # Model handler will restore model parameters from the checkpoint
-        # directory and assign parameters to train_model.
-        export_model = self.model_handler.get_model_to_export(
-            train_model, dataset=None
-        )
-        test_data = tf.SparseTensor(
-            indices=[[0, 0]], values=[0], dense_shape=(1, 1)
-        )
-        result = export_model.call(test_data).numpy()
+        with tempfile.TemporaryDirectory() as temp_dir:
+            self.model_handler._checkpoint_dir = os.path.join(
+                temp_dir, "test_export"
+            )
+            self._save_model(False)
+            model_inst = custom_model_with_sparse_embedding()
+            train_model = self.model_handler.get_model_to_train(model_inst)
+
+            # Model handler will restore model parameters from the checkpoint
+            # directory and assign parameters to train_model.
+            export_model = self.model_handler.get_model_to_export(
+                train_model, dataset=None
+            )
+            test_data = tf.SparseTensor(
+                indices=[[0, 0]], values=[0], dense_shape=(1, 1)
+            )
+            result = export_model.call(test_data).numpy()
 
-        # The embedding table in checkpoint file is
-        # [[1.0, 1.0], [1.0, 1.0], [1.0,1.0], [1.0, 1.0]], weights in the dense
-        # layer is [[1.0],[1.0]], bias is [1.0]. So the result is 3.0.
-        self.assertEqual(result[0][0], 3.0)
+            # The embedding table in checkpoint file is
+            # [[1.0, 1.0], [1.0, 1.0], [1.0,1.0], [1.0, 1.0]], weights in the
+            # dense layer is [[1.0],[1.0]], bias is [1.0]. So the result
+            # is 3.0.
+            self.assertEqual(result[0][0], 3.0)
 
 
 if __name__ == "__main__":
diff --git a/elasticdl/python/tests/test_utils.py b/elasticdl/python/tests/test_utils.py
--- a/elasticdl/python/tests/test_utils.py
+++ b/elasticdl/python/tests/test_utils.py
@@ -23,13 +23,14 @@
     get_module_file_path,
     load_module,
 )
+from elasticdl.python.common.save_utils import CheckpointSaver
 from elasticdl.python.data.recordio_gen.frappe_recordio_gen import (
     load_raw_data,
 )
 from elasticdl.python.master.evaluation_service import EvaluationService
 from elasticdl.python.master.servicer import MasterServicer
 from elasticdl.python.master.task_dispatcher import _TaskDispatcher
-from elasticdl.python.ps.parameter_server import ParameterServer
+from elasticdl.python.ps.parameter_server import Parameters, ParameterServer
 from elasticdl.python.tests.in_process_master import InProcessMaster
 from elasticdl.python.worker.worker import Worker
 
@@ -631,3 +632,13 @@ def get_frappe_dataset(batch_size):
     test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
     test_db = test_db.batch(batch_size)
     return db, test_db
+
+
+def save_checkpoint_without_embedding(model, checkpoint_dir, version=100):
+    checkpoint_saver = CheckpointSaver(checkpoint_dir, 0, 0, False)
+    params = Parameters()
+    for var in model.trainable_variables:
+        params.non_embedding_params[var.name] = var
+    params.version = version
+    model_pb = params.to_model_pb()
+    checkpoint_saver.save(version, model_pb, False)
diff --git a/elasticdl/python/tests/testdata/functional_ckpt/version-100/variables-0-of-1.ckpt b/elasticdl/python/tests/testdata/functional_ckpt/version-100/variables-0-of-1.ckpt
deleted file mode 100644
Binary files a/elasticdl/python/tests/testdata/functional_ckpt/version-100/variables-0-of-1.ckpt and /dev/null differ
diff --git a/elasticdl/python/tests/testdata/mnist_functional_api_model/version-110/variables-0-of-1.ckpt b/elasticdl/python/tests/testdata/mnist_functional_api_model/version-110/variables-0-of-1.ckpt
deleted file mode 100644
Binary files a/elasticdl/python/tests/testdata/mnist_functional_api_model/version-110/variables-0-of-1.ckpt and /dev/null differ
diff --git a/elasticdl/python/tests/testdata/subclass_ckpt/version-100/variables-0-of-1.ckpt b/elasticdl/python/tests/testdata/subclass_ckpt/version-100/variables-0-of-1.ckpt
deleted file mode 100644
Binary files a/elasticdl/python/tests/testdata/subclass_ckpt/version-100/variables-0-of-1.ckpt and /dev/null differ
diff --git a/scripts/client_test.sh b/scripts/client_test.sh
--- a/scripts/client_test.sh
+++ b/scripts/client_test.sh
@@ -4,6 +4,11 @@ JOB_TYPE=$1
 PS_NUM=$2
 WORKER_NUM=$3
 
+# Generate checkpoint for mnist to test evaluation and prediction
+MNIST_CKPT_DIR=model_zoo/test_data/mnist_ckpt/
+python -m scripts.gen_mnist_checkpoint --checkpoint_dir=${MNIST_CKPT_DIR}
+
+
 if [[ "$JOB_TYPE" == "train" ]]; then
     elasticdl train \
       --image_base=elasticdl:ci \
@@ -37,7 +42,7 @@ elif [[ "$JOB_TYPE" == "evaluate" ]]; then
       --image_base=elasticdl:ci \
       --model_zoo=model_zoo \
       --model_def=mnist_functional_api.mnist_functional_api.custom_model \
-      --checkpoint_dir_for_init=elasticdl/python/tests/testdata/mnist_functional_api_model/version-110  \
+      --checkpoint_dir_for_init=model_zoo/${MNIST_CKPT_DIR}/version-100  \
       --validation_data=/data/mnist/test \
       --num_epochs=1 \
       --master_resource_request="cpu=0.3,memory=1024Mi" \
@@ -60,7 +65,7 @@ elif [[ "$JOB_TYPE" == "predict" ]]; then
       --image_base=elasticdl:ci \
       --model_zoo=model_zoo \
       --model_def=mnist_functional_api.mnist_functional_api.custom_model \
-      --checkpoint_dir_for_init=elasticdl/python/tests/testdata/mnist_functional_api_model/version-110 \
+      --checkpoint_dir_for_init=model_zoo/${MNIST_CKPT_DIR}/version-100 \
       --prediction_data=/data/mnist/test \
       --master_resource_request="cpu=0.2,memory=1024Mi" \
       --master_resource_limit="cpu=1,memory=2048Mi" \