Repository: sql-machine-learning/elasticdl

Issue:
When `get_model` returns with `model_init_status==False`,  worker reports local variables to ps for ps initialization.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -165,6 +165,8 @@ def set_model(self, model_inst):
                 self._model, self._embedding_layers
             ):
                 self._non_embed_vars[var.name] = var
+            if self._use_multi_ps:
+                self.init_ps_var_partition()
 
     def _init_embedding_layer(self):
         """
@@ -226,10 +228,22 @@ def get_model_from_master(self, version, method):
 
     def get_model_from_ps(self, version, method):
         model_version = -1
-        for stub in self._ps_stubs:
+        for ps_id, stub in enumerate(self._ps_stubs):
+            if ps_id not in self._ps_vars:
+                continue
             req = empty_pb2.Empty()
             res = stub.pull_variable(req)
 
+            if not res.model_init_status:
+                # push variable to ps for initialization
+                self.report_variable_to_ps(ps_id)
+                res = stub.pull_variable(req)
+                if not res.model_init_status:
+                    # TODO: support PS fault-tolerance
+                    raise RuntimeError(
+                        "PS pod %d cannot be initialized" % ps_id
+                    )
+
             for tensor_pb in res.model.param:
                 tensor = Tensor.from_tensor_pb(tensor_pb)
                 self._non_embed_vars[tensor.name].assign(tensor.to_ndarray())
@@ -285,7 +299,7 @@ def report_variable_to_master(self):
             )
         self._stub.ReportVariable(req)
 
-    def report_variable_to_ps(self):
+    def init_ps_var_partition(self):
         ps_vars = {}
         for v in self._non_embed_vars.values():
             if v.name not in self._var_to_ps:
@@ -297,29 +311,34 @@ def report_variable_to_ps(self):
                 ps_vars[ps_id] = [v]
             else:
                 ps_vars[ps_id].append(v)
+        self._ps_vars = ps_vars
+
+    def report_variable_to_ps(self, ps_id):
+        model = elasticdl_pb2.Model()
+        if ps_id in self._ps_vars:
+            vars = self._ps_vars[ps_id]
+            for var in vars:
+                emplace_tensor_pb_from_ndarray(
+                    model.param, var.numpy(), name=var.name
+                )
+        if self._embedding_layers:
+            embedding_infos = elasticdl_pb2.EmbeddingTableInfo
+            for layer in self._embedding_layers:
+                embedding_info = embedding_infos.add()
+                embedding_info.name = layer.name
+                embedding_info.dim = layer.output_dim
+                embedding_info.initializer = layer.embeddings_initializer
+
+        self._ps_stubs[ps_id].push_model(model)
 
+    def report_variable_to_all_ps(self):
         # TODO: call `push_model` in parallel
         for ps_id in range(len(self._ps_stubs)):
-            model = elasticdl_pb2.Model()
-            if ps_id in ps_vars:
-                vars = ps_vars[ps_id]
-                for var in vars:
-                    emplace_tensor_pb_from_ndarray(
-                        model.param, var.numpy(), name=var.name
-                    )
-            if self._embedding_layers:
-                embedding_infos = elasticdl_pb2.EmbeddingTableInfo
-                for layer in self._embedding_layers:
-                    embedding_info = embedding_infos.add()
-                    embedding_info.name = layer.name
-                    embedding_info.dim = layer.output_dim
-                    embedding_info.initializer = layer.embeddings_initializer
-
-            self._ps_stubs[ps_id].push_model(model)
+            self.report_variable_to_ps(ps_id)
 
     def report_variable(self):
         if self._use_multi_ps:
-            self.report_variable_to_ps()
+            self.report_variable_to_all_ps()
         else:
             self.report_variable_to_master()
 
@@ -502,7 +521,10 @@ def _run_model_call_before_training(self, features):
             self._non_embed_vars[var.name] = var
 
         if not self._var_created:
-            self.report_variable()
+            if self._use_multi_ps:
+                self.init_ps_var_partition()
+            else:
+                self.report_variable()
             self._var_created = True
 
         if self._need_embedding_layer_check:


Test patch:
diff --git a/elasticdl/python/tests/worker_ps_interaction_test.py b/elasticdl/python/tests/worker_ps_interaction_test.py
--- a/elasticdl/python/tests/worker_ps_interaction_test.py
+++ b/elasticdl/python/tests/worker_ps_interaction_test.py
@@ -39,8 +39,10 @@ def setUp(self):
             "mnist_functional_api.mnist_functional_api.custom_model"
         )
         self._batch_size = 16
-        ports = [12345, 12346]
-        self._pserver, self._channel = self._create_pserver_and_channel(ports)
+        self._ports = [12345, 12346]
+        self._pserver, self._channel = self._create_pserver_and_channel(
+            self._ports
+        )
 
     def tearDown(self):
         for pserver in self._pserver:
@@ -78,6 +80,14 @@ def _create_pserver_and_channel(self, ports):
             channels.append(channel)
         return pservers, channels
 
+    def _restart_pserver(self):
+        # Stop first
+        self.tearDown()
+        # Start again
+        self._pserver, self._channel = self._create_pserver_and_channel(
+            self._ports
+        )
+
     def test_worker_pull_embedding(self):
         worker = Worker(
             worker_id=0,
@@ -133,7 +143,6 @@ def test_compare_onebatch_train(self):
             ps_channels=self._channel,
         )
         worker._run_model_call_before_training(images)
-        worker.report_variable()
         worker.get_model(0, elasticdl_pb2.MINIMUM)
         w_loss, w_grads = worker.training_process_eagerly(images, labels)
         worker.report_gradient(w_grads)
@@ -198,7 +207,6 @@ def _worker_train(self, train_db, test_db, dataset, stop_step):
         for step, (x, y) in enumerate(train_db):
             if step == 0:
                 worker._run_model_call_before_training(x)
-                worker.report_variable()
 
             worker.get_model(step, elasticdl_pb2.MINIMUM)
             w_loss, w_grads = worker.training_process_eagerly(x, y)
@@ -247,8 +255,9 @@ def test_compare_mnist_train(self):
 
         tf.keras.backend.clear_session()
         tf.random.set_seed(22)
+        stop_step = 20
         worker_results = self._worker_train(
-            train_db=db, test_db=test_db, dataset="mnist", stop_step=50
+            train_db=db, test_db=test_db, dataset="mnist", stop_step=stop_step
         )
 
         tf.keras.backend.clear_session()
@@ -294,7 +303,7 @@ def test_compare_mnist_train(self):
                 )
                 acc_meter.reset_states()
 
-            if step > 50:
+            if step > stop_step:
                 break
 
         for w, l in zip(worker_results, local_results):
@@ -323,10 +332,52 @@ def __init__(self, data):
         test_db = test_db.batch(self._batch_size)
 
         worker_results = self._worker_train(
-            train_db=db, test_db=test_db, dataset="frappe", stop_step=200
+            train_db=db, test_db=test_db, dataset="frappe", stop_step=100
         )
         acc = max([r[1] for r in worker_results])
-        self.assertLess(0.7, acc)
+        self.assertLess(0.6, acc)
+
+    def test_restart_ps(self):
+        num_data = 8
+        training_data = [
+            random_batch(self._batch_size) for _ in range(num_data)
+        ]
+        workers = []
+        for w in range(2):
+            self._restart_pserver()
+            tf.keras.backend.clear_session()
+            tf.random.set_seed(22)
+            worker = Worker(
+                worker_id=0,
+                job_type=elasticdl_pb2.TRAINING,
+                minibatch_size=self._batch_size,
+                model_zoo=self._model_zoo_path,
+                model_def=self._model_def,
+                ps_channels=self._channel,
+            )
+            workers.append(worker)
+            worker._run_model_call_before_training(training_data[0][0])
+            for i in range(num_data):
+                worker.get_model(0, elasticdl_pb2.MINIMUM)
+                w_loss, w_grads = worker.training_process_eagerly(
+                    training_data[i][0], training_data[i][1]
+                )
+                worker.report_gradient(w_grads)
+                if w == 1 and i == 3:
+                    # Restart ps for the 2nd worker at i==3
+                    self._restart_pserver()
+                    # `report_variable` will be called in `get_model` to
+                    # initialize variables on ps with worker variables
+                    worker.get_model(0, elasticdl_pb2.MINIMUM)
+                    # send the grads again as these grads are not applied
+                    # on worker variables
+                    worker.report_gradient(w_grads)
+
+        for var_name in workers[0]._non_embed_vars:
+            np.testing.assert_array_equal(
+                workers[0]._non_embed_vars[var_name].numpy(),
+                workers[1]._non_embed_vars[var_name].numpy(),
+            )
 
 
 if __name__ == "__main__":