Repository: sql-machine-learning/elasticdl

Issue:
Add arguments in worker.main to pass redis endpoint
According [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#master-start-distributed-storage-service-when-needed), master should tell workers how to access redis cluster. 


Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/main.py b/elasticdl/python/master/main.py
--- a/elasticdl/python/master/main.py
+++ b/elasticdl/python/master/main.py
@@ -167,6 +167,29 @@ def main():
         evaluation_service.start()
         task_d.set_evaluation_service(evaluation_service)
 
+    embedding_service_endpoint = None
+    # Search for embedding layers in the model,
+    # if found, initialize embedding service
+    layers = find_layer(model_inst, Embedding)
+    if layers:
+        embedding_service = EmbeddingService()
+        embedding_service_endpoint = embedding_service.start_embedding_service(
+            job_name=args.job_name,
+            image_name=args.worker_image,
+            namespace=args.namespace,
+            resource_request=args.master_resource_request,
+            resource_limit=args.master_resource_limit,
+            pod_priority=args.worker_pod_priority,
+            volume=args.volume,
+            image_pull_policy=args.image_pull_policy,
+            restart_policy=args.restart_policy,
+            cluster_spec=args.cluster_spec,
+        )
+        logger.info(
+            "Embedding service start succeeded. The endpoint is %s."
+            % str(embedding_service_endpoint)
+        )
+
     # The master service
     logger.info("Starting master service")
     server = grpc.server(
@@ -188,33 +211,13 @@ def main():
         checkpoint_filename_for_init=args.checkpoint_filename_for_init,
         checkpoint_service=checkpoint_service,
         evaluation_service=evaluation_service,
+        embedding_service_endpoint=embedding_service_endpoint,
     )
     elasticdl_pb2_grpc.add_MasterServicer_to_server(master_servicer, server)
     server.add_insecure_port("[::]:{}".format(args.port))
     server.start()
     logger.info("Server started at port: %d", args.port)
 
-    # Search for embedding layers in the model,
-    # if found, initialize embedding service
-    layers = find_layer(model_inst, Embedding)
-    if layers:
-        embedding_service = EmbeddingService()
-        redis_address_map = embedding_service.start_embedding_service(
-            job_name=args.job_name,
-            image_name=args.worker_image,
-            namespace=args.namespace,
-            resource_request=args.master_resource_request,
-            resource_limit=args.master_resource_limit,
-            pod_priority=args.worker_pod_priority,
-            volume=args.volume,
-            image_pull_policy=args.image_pull_policy,
-            restart_policy=args.restart_policy,
-            cluster_spec=args.cluster_spec,
-        )
-        logger.info(
-            "Embedding service start succeeded: %s" % str(redis_address_map)
-        )
-
     worker_manager = None
     if args.num_workers:
         assert args.worker_image, "Worker image cannot be empty"
@@ -243,6 +246,8 @@ def main():
             job_type,
             "--minibatch_size",
             str(args.minibatch_size),
+            "--embedding_service_endpoint",
+            str(embedding_service_endpoint),
         ]
 
         worker_manager = WorkerManager(
diff --git a/elasticdl/python/master/servicer.py b/elasticdl/python/master/servicer.py
--- a/elasticdl/python/master/servicer.py
+++ b/elasticdl/python/master/servicer.py
@@ -32,7 +32,7 @@ def __init__(
         checkpoint_filename_for_init,
         checkpoint_service,
         evaluation_service,
-        embedding_service_address=None
+        embedding_service_endpoint=None
     ):
         # TODO: group params together into a single object.
         self._opt = optimizer
@@ -50,7 +50,7 @@ def __init__(
         # optimizer's apply_gradients() function.
         self._model = {}
         self._version = 0
-        self.embedding_service_address = embedding_service_address
+        self._embedding_service_endpoint = embedding_service_endpoint
         self._init_model(checkpoint_filename_for_init, init_var)
 
         self._checkpoint_service = checkpoint_service
@@ -172,7 +172,7 @@ def _update_edl_embedding_table(self, name_var_list):
             EmbeddingService.update_embedding(
                 keys=keys,
                 embeddings=embeddings,
-                embedding_service_address=self.embedding_service_address,
+                embedding_service_endpoint=self._embedding_service_endpoint,
             )
 
     def _update_model(self):
@@ -198,7 +198,9 @@ def _update_model(self):
                 unique_ids_list.append(unique_ids)
                 grads_idx_transformed = tf.IndexedSlices(grads.values, idx)
                 embeddings = EmbeddingService.lookup_embedding(
-                    embedding_service_address=self.embedding_service_address,
+                    embedding_service_endpoint=(
+                        self._embedding_service_endpoint
+                    ),
                     keys=[
                         Embedding.get_key([layer_name, i])
                         for i in unique_ids.numpy()
diff --git a/elasticdl/python/worker/main.py b/elasticdl/python/worker/main.py
--- a/elasticdl/python/worker/main.py
+++ b/elasticdl/python/worker/main.py
@@ -73,6 +73,13 @@ def _parse_args():
         help="The dictionary of model parameters in a string that will be "
         'used to instantiate the model, e.g. "param1=1,param2=2"',
     )
+    parser.add_argument(
+        "--embedding_service_endpoint",
+        type=str,
+        default="{}",
+        help="The endpoint of embedding service, "
+        "e.g. \"{'ip_0': [port_0,port_1]}\"",
+    )
 
     return parser.parse_args()
 
@@ -99,6 +106,7 @@ def main():
         args.minibatch_size,
         get_model_file(args.model_zoo, args.model_def),
         channel=channel,
+        embedding_service_endpoint=eval(args.embedding_service_endpoint),
         dataset_fn=args.dataset_fn,
         loss=args.loss,
         optimizer=args.optimizer,


Test patch: