Repository: sql-machine-learning/elasticdl

Issue:
PS needs to init wrap initializer when embedding exists
For a PS that does not have non-embedding parameters, it needs to initialize wrap optimizer  if embedding layer exists.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/ps/servicer.py b/elasticdl/python/ps/servicer.py
--- a/elasticdl/python/ps/servicer.py
+++ b/elasticdl/python/ps/servicer.py
@@ -37,6 +37,7 @@ def __init__(
         self._eval_steps = evaluation_steps
         self._version_lock = threading.Lock()
         self._lock = threading.Lock()
+        self._use_wrap_opt = False
 
         self._grads_n = 0
         self._grads_buffer = {}
@@ -79,11 +80,7 @@ def push_model(self, request, _):
         with self._lock:
             accepted = self._parameters.init_from_model_pb(request)
         if accepted and self._parameters.has_embedding_params():
-            self.wrap_optimizer()
-            self._parameters.create_slot_params(
-                self._optimizer.allowed_slot_names,
-                self._optimizer.slot_initial_value,
-            )
+            self.wrap_optimizer_and_set_slot()
         return empty_pb2.Empty()
 
     def push_embedding_info(self, request, _):
@@ -91,6 +88,7 @@ def push_embedding_info(self, request, _):
             self._parameters.init_embedding_params(
                 request.embedding_table_info
             )
+            self.wrap_optimizer_and_set_slot()
         return empty_pb2.Empty()
 
     def push_gradient(self, request, _):
@@ -209,3 +207,12 @@ def _report_version(self, version):
         req = elasticdl_pb2.ReportVersionRequest()
         req.model_version = version
         self._master_stub.ReportVersion(req)
+
+    def wrap_optimizer_and_set_slot(self):
+        if not self._use_wrap_opt:
+            self.wrap_optimizer()
+            self._parameters.create_slot_params(
+                self._optimizer.allowed_slot_names,
+                self._optimizer.slot_initial_value,
+            )
+            self._use_wrap_opt = True


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
@@ -167,6 +167,22 @@ def _worker_train(self, worker_id, train_db, test_db, stop_step):
                 break
         return worker_results
 
+    def _test_deepfm_train(self, num_ps, num_worker, stop_step):
+        model_def = "deepfm_functional_api.deepfm_functional_api.custom_model"
+        self._create_pserver(model_def, num_ps)
+        db, test_db = get_frappe_dataset(self._batch_size)
+
+        self._create_worker(num_worker)
+        threads = []
+        for w in range(num_worker):
+            t = Thread(
+                target=self._worker_train, args=(w, db, test_db, stop_step)
+            )
+            t.start()
+            threads.append(t)
+        for t in threads:
+            t.join()
+
     def test_worker_pull_embedding(self):
         model_def = "mnist_functional_api.mnist_functional_api.custom_model"
         self._create_pserver(model_def, 2)
@@ -352,19 +368,16 @@ def test_deepfm_train(self):
         self.assertLess(0.6, acc)
 
     def test_deepfm_two_worker_train(self):
-        model_def = "deepfm_functional_api.deepfm_functional_api.custom_model"
-        self._create_pserver(model_def, 2)
-        db, test_db = get_frappe_dataset(self._batch_size)
-
-        self._create_worker(2)
-        t1 = Thread(target=self._worker_train, args=(0, db, test_db, 50))
-        t1.start()
-
-        t2 = Thread(target=self._worker_train, args=(1, db, test_db, 50))
-        t2.start()
-
-        t1.join()
-        t2.join()
+        num_ps = 2
+        num_worker = 2
+        stop_step = 10
+        self._test_deepfm_train(num_ps, num_worker, stop_step)
+
+    def test_deepfm_four_worker_train(self):
+        num_ps = 4
+        num_worker = 1
+        stop_step = 10
+        self._test_deepfm_train(num_ps, num_worker, stop_step)
 
     def test_restart_ps(self):
         model_def = "mnist_functional_api.mnist_functional_api.custom_model"