internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: sql-machine-learning/elasticdl

Issue:
The status of the worker pod should be Error when there is any error happened.
The status of the worker pod is still Running and will never stop even if there are some errors in function `handle_task` like:
```bash
Traceback (most recent call last):
  File "/elasticdl/python/elasticdl/worker/worker.py", line 236, in run
    self._handle_task(task)
  File "/elasticdl/python/elasticdl/worker/worker.py", line 190, in _handle_task
    task, record_buf, min_model_version
  File "/elasticdl/python/elasticdl/worker/worker.py", line 211, in _process_minibatch
    features, labels
  File "/elasticdl/python/elasticdl/worker/worker.py", line 168, in _run_training_task
    accepted, min_model_version = self.report_gradient(grads)
  File "/elasticdl/python/elasticdl/worker/worker.py", line 115, in report_gradient
    res = self._stub.ReportGradient(req)
  File "/usr/local/lib/python3.6/dist-packages/grpc/_channel.py", line 565, in __call__
    return _end_unary_response_blocking(state, call, False, None)
  File "/usr/local/lib/python3.6/dist-packages/grpc/_channel.py", line 467, in _end_unary_response_blocking
    raise _Rendezvous(state, None, None, deadline)


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/worker/worker.py b/elasticdl/python/elasticdl/worker/worker.py
--- a/elasticdl/python/elasticdl/worker/worker.py
+++ b/elasticdl/python/elasticdl/worker/worker.py
@@ -234,7 +234,11 @@ def run(self):
             err_msg = ""
             try:
                 self._handle_task(task)
+            except RuntimeError as err:
+                err_msg = str(err)
+                traceback.print_exc()
             except Exception as ex:
                 err_msg = str(ex)
                 traceback.print_exc()
+                raise ex
             self.report_task_result(task.task_id, err_msg)


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Hide TensorBoard service REST call logs
Currently these logs appear on master pod's log which is not necessary at user level:
```
I0624 15:50:54.834580 140272641951488 _internal.py:122] ::ffff:10.138.0.35 - - [24/Jun/2019 15:50:54] "GET /data/plugin/scalars/tags HTTP/1.1" 200 -
W0624 15:50:55.150964 140556554397440 servicer.py:195] Task result for outdated version 30 dropped
I0624 15:50:57.245738 140272633558784 _internal.py:122] ::ffff:10.36.1.1 - - [24/Jun/2019 15:50:57] "GET /data/environment HTTP/1.1" 200 -
I0624 15:50:57.250612 140271752836864 _internal.py:122] ::ffff:10.138.0.35 - - [24/Jun/2019 15:50:57] "GET /data/experiments HTTP/1.1" 200 -
I0624 15:50:57.252254 140272105903872 _internal.py:122] ::ffff:10.36.1.1 - - [24/Jun/2019 15:50:57] "GET /data/runs HTTP/1.1" 200 -
I0624 15:50:57.255696 140272641951488 _internal.py:122] ::ffff:10.138.0.35 - - [24/Jun/2019 15:50:57] "GET /data/plugins_listing HTTP/1.1" 200 -
I0624 15:50:57.430979 140272641951488 _internal.py:122] ::ffff:10.138.0.35 - - [24/Jun/2019 15:50:57] "GET /data/plugin/scalars/tags HTTP/1.1" 200 -
```


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/master/tensorboard_service.py b/elasticdl/python/elasticdl/master/tensorboard_service.py
--- a/elasticdl/python/elasticdl/master/tensorboard_service.py
+++ b/elasticdl/python/elasticdl/master/tensorboard_service.py
@@ -34,7 +34,10 @@ def write_dict_to_summary(self, dictionary, version):
 
     def start(self):
         self.tb_process = subprocess.Popen(
-            ["tensorboard --logdir " + self._tensorboard_log_dir], shell=True
+            ["tensorboard --logdir " + self._tensorboard_log_dir],
+            shell=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.STDOUT,
         )
 
     def keep_running(self):


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Worker pod should exit only after all tasks in task_q are done
Currently, a worker pod exits when GetTask does not get any task.
GetTask returns null task if there is no task in `to_do` queue.
But there may exist tasks in `doing` queue. New tasks may be added into `to_do` queue later:
1. New evaluation tasks
2. Tasks in `doing` queue may fail, and put back into `to_do` queue.

Thus, the worker pod should wait until all tasks are done.


Discussion:

Resolving patch:
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
@@ -103,6 +103,9 @@ def GetTask(self, request, _):
             # For evaluation task, it will use the fixed version model
             if task.type == elasticdl_pb2.EVALUATION:
                 res.model_version = task.model_version
+        elif not self._task_q.finished():
+            # Not all tasks are finished, wait in case of new tasks later.
+            res.type = elasticdl_pb2.WAIT
         return res
 
     def GetModel(self, request, _):
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
@@ -1,4 +1,5 @@
 import logging
+import time
 import traceback
 from contextlib import closing
 
@@ -224,11 +225,16 @@ def run(self):
         """
         while True:
             task = self.get_task()
-            self._logger.info("Receive a new task: %d", task.task_id)
             if not task.shard_file_name:
-                # No more task
-                self._logger.info("No more task, stopping")
-                break
+                if task.type == elasticdl_pb2.WAIT:
+                    # Wait a few seconds then try to get_task again
+                    time.sleep(5)
+                    continue
+                else:
+                    # No more task
+                    self._logger.info("No more task, stopping")
+                    break
+            self._logger.info("Receive a new task: %d", task.task_id)
             err_msg = ""
             try:
                 self._handle_task(task)


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Add label to pod so that users know the job has finished
TensorBoard needs to keep running even when tasks are finished. We should add labels to master pod so that users know the job has finished.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/constants.py b/elasticdl/python/common/constants.py
--- a/elasticdl/python/common/constants.py
+++ b/elasticdl/python/common/constants.py
@@ -3,3 +3,9 @@ class GRPC(object):
     # It's too small to send model parameters.
     MAX_SEND_MESSAGE_LENGTH = 256 * 1024 * 1024
     MAX_RECEIVE_MESSAGE_LENGTH = 256 * 1024 * 1024
+
+
+class WorkerManagerStatus(object):
+    PENDING = "Pending"
+    RUNNING = "Running"
+    FINISHED = "Finished"
diff --git a/elasticdl/python/master/k8s_worker_manager.py b/elasticdl/python/master/k8s_worker_manager.py
--- a/elasticdl/python/master/k8s_worker_manager.py
+++ b/elasticdl/python/master/k8s_worker_manager.py
@@ -76,6 +76,12 @@ def _start_worker(self, worker_id):
             self._pod_name_to_id[name] = worker_id
             self._pods_phase[worker_id] = (name, None)
 
+    def update_status(self, status):
+        master_name = self._k8s_client.get_master_pod_name()
+        self._k8s_client.patch_labels_to_pod(
+            master_name, labels_dict={"status": status}
+        )
+
     def start_workers(self):
         for i in range(self._num_workers):
             self._start_worker(self._next_worker_id())
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
@@ -8,7 +8,7 @@
 import recordio
 
 from elasticdl.proto import elasticdl_pb2_grpc
-from elasticdl.python.common.constants import GRPC
+from elasticdl.python.common.constants import GRPC, WorkerManagerStatus
 from elasticdl.python.common.model_helper import get_model_file, load_module
 from elasticdl.python.master.args import parse_args
 from elasticdl.python.master.checkpoint_service import CheckpointService
@@ -159,6 +159,7 @@ def main():
     server.start()
     logger.info("Server started at port: %d", args.port)
 
+    worker_manager = None
     if args.num_workers:
         assert args.worker_image, "Worker image cannot be empty"
 
@@ -199,8 +200,10 @@ def main():
             image_pull_policy=args.image_pull_policy,
             restart_policy=args.restart_policy,
         )
+        worker_manager.update_status(WorkerManagerStatus.PENDING)
         logger.info("Launching %d workers", args.num_workers)
         worker_manager.start_workers()
+        worker_manager.update_status(WorkerManagerStatus.RUNNING)
 
         if tb_service:
             worker_manager.start_tensorboard_service()
@@ -208,6 +211,8 @@ def main():
     try:
         while True:
             if task_d.finished():
+                if worker_manager:
+                    worker_manager.update_status(WorkerManagerStatus.FINISHED)
                 break
             time.sleep(30)
     except KeyboardInterrupt:


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Print out name of the created master pod after submitting a job via client
Currently, when submitting a job via `elasticdl train --job_name=xxx`, no master pod information will be printed out and users have to guess the name to master pod from the job name they provided. 

We should print out the name to master pod when job has been submitted successfully.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/api.py b/elasticdl/python/elasticdl/api.py
--- a/elasticdl/python/elasticdl/api.py
+++ b/elasticdl/python/elasticdl/api.py
@@ -85,21 +85,7 @@ def train(args):
         else args.master_resource_request
     )
 
-    k8s.Client(
-        image_name=image_name,
-        namespace=args.namespace,
-        job_name=args.job_name,
-        event_callback=None,
-        cluster_spec=args.cluster_spec,
-    ).create_master(
-        resource_requests=args.master_resource_request,
-        resource_limits=args.master_resource_limit,
-        args=container_args,
-        pod_priority=args.master_pod_priority,
-        image_pull_policy=args.image_pull_policy,
-        restart_policy=args.restart_policy,
-        volume=args.volume,
-    )
+    _submit_job(image_name, args, container_args)
     # TODO: print dashboard url after launching the master pod
 
 
@@ -157,20 +143,30 @@ def evaluate(args):
         else args.master_resource_request
     )
 
-    k8s.Client(
+    _submit_job(image_name, args, container_args)
+
+
+def _submit_job(image_name, client_args, container_args):
+    client = k8s.Client(
         image_name=image_name,
-        namespace=args.namespace,
-        job_name=args.job_name,
+        namespace=client_args.namespace,
+        job_name=client_args.job_name,
         event_callback=None,
-        cluster_spec=args.cluster_spec,
-    ).create_master(
-        resource_requests=args.master_resource_request,
-        resource_limits=args.master_resource_limit,
+        cluster_spec=client_args.cluster_spec,
+    )
+
+    client.create_master(
+        resource_requests=client_args.master_resource_request,
+        resource_limits=client_args.master_resource_limit,
         args=container_args,
-        pod_priority=args.master_pod_priority,
-        image_pull_policy=args.image_pull_policy,
-        restart_policy=args.restart_policy,
-        volume=args.volume,
+        pod_priority=client_args.master_pod_priority,
+        image_pull_policy=client_args.image_pull_policy,
+        restart_policy=client_args.restart_policy,
+        volume=client_args.volume,
+    )
+    print(
+        "ElasticDL job %s was successfully submitted. The master pod is: %s."
+        % (client_args.job_name, client.get_master_pod_name())
     )
 
 


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Use tf.function-decorated functions to accelerate model compute



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
@@ -155,7 +155,8 @@ def _create_variable_and_report(self, features):
         self.report_variable()
         self._var_created = True
 
-    def _run_training_task(self, features, labels):
+    @tf.function
+    def training_process(self, features, labels):
         with tf.GradientTape() as tape:
             outputs = self._model.call(features, training=True)
             loss = self._loss(outputs, labels)
@@ -163,16 +164,30 @@ def _run_training_task(self, features, labels):
             if self._model.losses:
                 loss += math_ops.add_n(self._model.losses)
         grads = tape.gradient(loss, self._model.trainable_variables)
+        return loss, grads
+
+    @tf.function
+    def evaluation_process(self, features, labels):
+        outputs = self._model.call(features, training=False)
+        evaluation_metrics = self._eval_metrics_fn(outputs, labels)
+        return evaluation_metrics
+
+    @tf.function
+    def predict_process(self, features):
+        outputs = self._model.call(features, training=False)
+        return outputs
+
+    def _run_training_task(self, features, labels):
+        loss, grads = self.training_process(features, labels)
         accepted, min_model_version = self.report_gradient(grads)
         return accepted, min_model_version, loss
 
     def _run_evaluation_task(self, features, labels):
-        outputs = self._model.call(features, training=False)
-        evaluation_metrics = self._eval_metrics_fn(outputs, labels)
+        evaluation_metrics = self.evaluation_process(features, labels)
         return self.report_evaluation_metrics(evaluation_metrics)
 
     def _run_prediction_task(self, features):
-        predictions = self._model.call(features, training=False)
+        predictions = self.predict_process(features)
         return self.report_prediction_outputs(predictions)
 
     def _handle_task(self, task):


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Move _pos_int and _non_neg_int to python/common/args.py
`_pos_int` and `_non_neg_int` should be common functions


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/args.py b/elasticdl/python/common/args.py
--- a/elasticdl/python/common/args.py
+++ b/elasticdl/python/common/args.py
@@ -1,3 +1,19 @@
+def pos_int(arg):
+    res = int(arg)
+    if res <= 0:
+        raise ValueError("Positive integer argument required. Got %s" % res)
+    return res
+
+
+def non_neg_int(arg):
+    res = int(arg)
+    if res < 0:
+        raise ValueError(
+            "Non-negative integer argument required. Get %s" % res
+        )
+    return res
+
+
 def add_common_params(parser):
     parser.add_argument(
         "--model_zoo",
diff --git a/elasticdl/python/master/args.py b/elasticdl/python/master/args.py
--- a/elasticdl/python/master/args.py
+++ b/elasticdl/python/master/args.py
@@ -1,22 +1,10 @@
 import argparse
 
-from elasticdl.python.common.args import add_common_params, add_train_params
-
-
-def _pos_int(arg):
-    res = int(arg)
-    if res <= 0:
-        raise ValueError("Positive integer argument required. Got %s" % res)
-    return res
-
-
-def _non_neg_int(arg):
-    res = int(arg)
-    if res < 0:
-        raise ValueError(
-            "Non-negative integer argument required. Get %s" % res
-        )
-    return res
+from elasticdl.python.common.args import (
+    add_common_params,
+    add_train_params,
+    pos_int,
+)
 
 
 def parse_args():
@@ -24,7 +12,7 @@ def parse_args():
     parser.add_argument(
         "--port",
         default=50001,
-        type=_pos_int,
+        type=pos_int,
         help="The listening port of master",
     )
     parser.add_argument(


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Master init Redis when needed
When a model contains embedding layer, master should initialize embedding service. Refer to [design doc](https://github.com/wangkuiyi/elasticdl/blob/develop/elasticdl/doc/distributed_embedding_layer_design.md#master-start-distributed-storage-service-when-needed)


Discussion:
we can create Redis service like tensorboard service and start a worker with specified command.

Resolving patch:
diff --git a/elasticdl/python/elasticdl/api.py b/elasticdl/python/elasticdl/api.py
--- a/elasticdl/python/elasticdl/api.py
+++ b/elasticdl/python/elasticdl/api.py
@@ -28,6 +28,10 @@ def train(args):
         _cluster_spec_def_in_docker(args.cluster_spec),
         "--num_workers",
         str(args.num_workers),
+        "--master_resource_request",
+        args.master_resource_request,
+        "--master_resource_limit",
+        args.master_resource_limit,
         "--worker_resource_request",
         args.worker_resource_request,
         "--worker_resource_limit",
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
@@ -12,12 +12,15 @@
     JobType,
     WorkerManagerStatus,
 )
+from elasticdl.python.common.embedding_service import EmbeddingService
 from elasticdl.python.common.log_util import get_logger
 from elasticdl.python.common.model_helper import (
+    find_layer,
     get_model_file,
     load_model_from_module,
     load_module,
 )
+from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.master.args import parse_args
 from elasticdl.python.master.checkpoint_service import CheckpointService
 from elasticdl.python.master.evaluation_service import EvaluationService
@@ -191,6 +194,27 @@ def main():
     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"


Test patch:
Repository: sql-machine-learning/elasticdl

Issue:
Fix lookup_embedding call in embedding layer
elasticdl.layers.Embedding.lookup_embedding still uses old worker.lookup_embedding method.
We need to update it to the revised version.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -87,8 +87,8 @@ def get_key(name_list):
         return "-".join(map(str, name_list))
 
     def lookup_embedding(self, unique_ids):
-        batch_embedding = self.worker.embedding_lookup(
-            unique_ids, self._name, self.embedding_initializer
+        batch_embedding = self.worker.lookup_embedding(
+            unique_ids, self._name, self.embedding_initializer, self.output_dim
         )
         return batch_embedding
 


Test patch:
diff --git a/elasticdl/python/tests/layer_test.py b/elasticdl/python/tests/layer_test.py
--- a/elasticdl/python/tests/layer_test.py
+++ b/elasticdl/python/tests/layer_test.py
@@ -65,7 +65,7 @@ def __init__(self, embedding_size, output_dim):
         for i in range(embedding_size):
             self.embedding[i].fill(i)
 
-    def embedding_lookup(self, ids, name, embedding_initializer):
+    def lookup_embedding(self, ids, name, embedding_initializer, output_dim):
         values = np.take(self.embedding, ids, axis=0)
         return values
Repository: sql-machine-learning/elasticdl

Issue:
tf.reshape in Embedding layer call supports for partially-defined shape
Need to convert None to -1 in TensorShape for 2D inputs (batched input).


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -112,9 +112,12 @@ def call(self, input):
             self.tape.watch(batch_embedding_tensor)
             self.bet_ids_pair.append((batch_embedding_tensor, unique_ids))
         outputs = tf.gather(batch_embedding_tensor, idx)
-        outputs = tf.reshape(
-            outputs, ids.get_shape().concatenate(self.output_dim)
-        )
+        # tf.reshape does not support shape with None. Replace None with -1.
+        if ids.get_shape().rank == 2:
+            output_shape = (-1, ids.get_shape()[1], self.output_dim)
+        else:
+            output_shape = ids.get_shape().concatenate(self.output_dim)
+        outputs = tf.reshape(outputs, output_shape)
         # TODO: support combiner for dense input
         return outputs
 


Test patch:
diff --git a/elasticdl/python/tests/layer_test.py b/elasticdl/python/tests/layer_test.py
--- a/elasticdl/python/tests/layer_test.py
+++ b/elasticdl/python/tests/layer_test.py
@@ -136,12 +136,20 @@ def test_embedding_layer(self):
         for index, idx in enumerate(ids):
             self.assertTrue((values[index] == results[index]).all())
 
-        model = tf.keras.models.Sequential([layer])
-        outputs = model.call(np.array([ids, ids]))
-        values = outputs.numpy()[1]
-        for index, idx in enumerate(ids):
-            correct_value = np.array([idx] * output_dim, dtype=np.float32)
-            self.assertTrue((values[index] == correct_value).all())
+        # Keras model without/with input_layer
+        model_without_input_layer = tf.keras.models.Sequential([layer])
+        inputs = tf.keras.Input(shape=(7,))
+        embeddings = layer(inputs)
+        model_with_input_layer = tf.keras.Model(
+            inputs=inputs, outputs=embeddings
+        )
+        models = [model_without_input_layer, model_with_input_layer]
+        for model in models:
+            outputs = model.call(tf.constant([ids, ids]))
+            values = outputs.numpy()[1]
+            for index, idx in enumerate(ids):
+                correct_value = np.array([idx] * output_dim, dtype=np.float32)
+                self.assertTrue((values[index] == correct_value).all())
 
     def test_embedding_layer_with_input_length(self):
         output_dim = 8