prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  [PS-1] RPC service `push_gradient` async-SGD implementation
  Async apply gradients to non-embedding parameters. Updating embedding parameters will be supported in issue #1306 .
  
  </issue>
  <code>
  [start of elasticdl/python/ps/parameters.py]
  1 import tensorflow as tf
  2 
  3 from elasticdl.python.common.tensor import tensor_pb_to_ndarray
  4 from elasticdl.python.ps.embedding_table import create_embedding_table
  5 
  6 
  7 class Parameters(object):
  8     """
  9     There are two kinds of parameters:
  10 
  11     1. non-embedding parameters, or dense tensor parameters. We save it
  12        in a hashmap `non-embedding_params`, the key is the parameter name,
  13        the value is a tf.Variable` object.
  14     2. embedding parameters, or row-sparse parameters. We save it in a
  15        hashmap `embedding_params`, the key is the embedding layer name,
  16        the value is an `EmbeddingTable` object.
  17 
  18     """
  19 
  20     def __init__(self):
  21         self.version = 0
  22         self.init_status = False
  23         self.non_embedding_params = {}
  24         self.embedding_params = {}
  25 
  26     def get_embedding_param(self, name, indices):
  27         if name not in self.embedding_params:
  28             raise ValueError(
  29                 "Please initialize embedding param %s first!", name
  30             )
  31         return self.embedding_params[name].get(indices)
  32 
  33     def set_embedding_param(self, name, indices, values):
  34         if name not in self.embedding_params:
  35             raise ValueError(
  36                 "Please initialize embedding param %s first!", name
  37             )
  38         self.embedding_params[name].set(indices, values)
  39 
  40     def init_from_model_pb(self, model_pb):
  41         if not self.init_status:
  42             tensors_pb = model_pb.param
  43             embeddings_pb = model_pb.embedding_table_info
  44             self._init_non_embedding_params(tensors_pb)
  45             self._init_embedding_params(embeddings_pb)
  46             self.version = model_pb.version
  47             self.init_status = True
  48 
  49     def _init_non_embedding_params(self, tensors_pb):
  50         for pb in tensors_pb:
  51             name = pb.name
  52             arr = tensor_pb_to_ndarray(pb)
  53             var = tf.Variable(name=name, initial_value=arr, trainable=True)
  54             self.non_embedding_params[name] = var
  55 
  56     def _init_embedding_params(self, embeddings_pb):
  57         for pb in embeddings_pb:
  58             self.embedding_params[pb.name] = create_embedding_table(pb)
  59
  [end of elasticdl/python/ps/parameters.py]
  [start of elasticdl/python/ps/servicer.py]
  1 import threading
  2 
  3 from google.protobuf import empty_pb2
  4 
  5 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
  6 from elasticdl.python.common.dtypes import dtype_numpy_to_tensor
  7 from elasticdl.python.common.tensor import Tensor, serialize_tensor
  8 
  9 
  10 class PserverServicer(elasticdl_pb2_grpc.PserverServicer):
  11     """PS service implementation"""
  12 
  13     def __init__(
  14         self,
  15         parameters,
  16         grads_to_wait,
  17         optimizer,
  18         lr_staleness_modulation=False,
  19         use_async=False,
  20     ):
  21         self._parameters = parameters
  22         self._grads_to_wait = grads_to_wait
  23         self._optimizer = optimizer
  24         self._lr_staleness_modulation = lr_staleness_modulation
  25         self._use_async = use_async
  26         self._version = 0
  27         self._lock = threading.Lock()
  28 
  29     def pull_variable(self, request, _):
  30         """
  31         Response with all non-embedding parameters if initialized.
  32         """
  33         res = elasticdl_pb2.PullVariableResponse()
  34         if not self._parameters.init_status:
  35             res.model_init_status = False
  36             return res
  37 
  38         # Only sync-SGD needs lock
  39         # TODO: use a read-write lock to support multiple concurrent reads
  40         if not self._use_async:
  41             self._lock.acquire()
  42         res.model.version = self._parameters.version
  43         for name, var in self._parameters.non_embedding_params.items():
  44             tensor = res.model.param.add()
  45             tensor.name = name
  46             tensor.dim.extend(var.shape.as_list())
  47             var_values = var.numpy()
  48             tensor.content = var_values.tobytes()
  49             tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)
  50         if not self._use_async:
  51             self._lock.release()
  52         res.model_init_status = True
  53         return res
  54 
  55     def pull_embedding_vector(self, request, _):
  56         ret = elasticdl_pb2.Tensor()
  57         if not request.ids:
  58             return ret
  59         embedding_vectors = self._parameters.get_embedding_param(
  60             request.name, request.ids
  61         )
  62         tensor = Tensor(values=embedding_vectors)
  63         serialize_tensor(tensor, ret)
  64         return ret
  65 
  66     def push_model(self, request, _):
  67         with self._lock:
  68             self._parameters.init_from_model_pb(request)
  69         return empty_pb2.Empty()
  70 
  71     def push_gradient(self, request, _):
  72         # TODO: implement this RPC service
  73         return elasticdl_pb2.PushGradientResponse()
  74
  [end of elasticdl/python/ps/servicer.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py
  --- a/elasticdl/python/ps/parameters.py
  +++ b/elasticdl/python/ps/parameters.py
  @@ -23,6 +23,9 @@
           self.non_embedding_params = {}
           self.embedding_params = {}
   
  +    def get_non_embedding_param(self, name, default_value=None):
  +        return self.non_embedding_params.get(name, default_value)
  +
       def get_embedding_param(self, name, indices):
           if name not in self.embedding_params:
               raise ValueError(
  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
  @@ -4,6 +4,7 @@
   
   from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
   from elasticdl.python.common.dtypes import dtype_numpy_to_tensor
  +from elasticdl.python.common.log_utils import default_logger as logger
   from elasticdl.python.common.tensor import Tensor, serialize_tensor
   
   
  @@ -23,7 +24,7 @@
           self._optimizer = optimizer
           self._lr_staleness_modulation = lr_staleness_modulation
           self._use_async = use_async
  -        self._version = 0
  +        self._version_lock = threading.Lock()
           self._lock = threading.Lock()
   
       def pull_variable(self, request, _):
  @@ -69,5 +70,29 @@
           return empty_pb2.Empty()
   
       def push_gradient(self, request, _):
  -        # TODO: implement this RPC service
  +        if self._use_async:
  +            grad_vars = []
  +            for pb in request.gradients:
  +                tensor = Tensor.from_tensor_pb(pb)
  +                var = self._parameters.get_non_embedding_param(tensor.name)
  +                if var is None:
  +                    logger.warning(
  +                        "Gradients with invalid name %s" % tensor.name
  +                    )
  +                    continue
  +                grad = tensor.to_tf_tensor()
  +                grad_vars.append((grad, var))
  +
  +            self._optimizer.apply_gradients(grad_vars)
  +            with self._version_lock:
  +                self._parameters.version += 1
  +
  +            res = elasticdl_pb2.PushGradientResponse()
  +            res.accepted = True
  +            res.model_version = self._parameters.version
  +            return res
  +
  +        raise NotImplementedError(
  +            "Updating parameters synchronously is not implemented."
  +        )
           return elasticdl_pb2.PushGradientResponse()
  
verification_info: {"golden_diff": "diff --git a/elasticdl/python/ps/parameters.py b/elasticdl/python/ps/parameters.py\n--- a/elasticdl/python/ps/parameters.py\n+++ b/elasticdl/python/ps/parameters.py\n@@ -23,6 +23,9 @@\n         self.non_embedding_params = {}\n         self.embedding_params = {}\n \n+    def get_non_embedding_param(self, name, default_value=None):\n+        return self.non_embedding_params.get(name, default_value)\n+\n     def get_embedding_param(self, name, indices):\n         if name not in self.embedding_params:\n             raise ValueError(\ndiff --git a/elasticdl/python/ps/servicer.py b/elasticdl/python/ps/servicer.py\n--- a/elasticdl/python/ps/servicer.py\n+++ b/elasticdl/python/ps/servicer.py\n@@ -4,6 +4,7 @@\n \n from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\n from elasticdl.python.common.dtypes import dtype_numpy_to_tensor\n+from elasticdl.python.common.log_utils import default_logger as logger\n from elasticdl.python.common.tensor import Tensor, serialize_tensor\n \n \n@@ -23,7 +24,7 @@\n         self._optimizer = optimizer\n         self._lr_staleness_modulation = lr_staleness_modulation\n         self._use_async = use_async\n-        self._version = 0\n+        self._version_lock = threading.Lock()\n         self._lock = threading.Lock()\n \n     def pull_variable(self, request, _):\n@@ -69,5 +70,29 @@\n         return empty_pb2.Empty()\n \n     def push_gradient(self, request, _):\n-        # TODO: implement this RPC service\n+        if self._use_async:\n+            grad_vars = []\n+            for pb in request.gradients:\n+                tensor = Tensor.from_tensor_pb(pb)\n+                var = self._parameters.get_non_embedding_param(tensor.name)\n+                if var is None:\n+                    logger.warning(\n+                        \"Gradients with invalid name %s\" % tensor.name\n+                    )\n+                    continue\n+                grad = tensor.to_tf_tensor()\n+                grad_vars.append((grad, var))\n+\n+            self._optimizer.apply_gradients(grad_vars)\n+            with self._version_lock:\n+                self._parameters.version += 1\n+\n+            res = elasticdl_pb2.PushGradientResponse()\n+            res.accepted = True\n+            res.model_version = self._parameters.version\n+            return res\n+\n+        raise NotImplementedError(\n+            \"Updating parameters synchronously is not implemented.\"\n+        )\n         return elasticdl_pb2.PushGradientResponse()\n", "issue": "[PS-1] RPC service `push_gradient` async-SGD implementation\nAsync apply gradients to non-embedding parameters. Updating embedding parameters will be supported in issue #1306 .\n", "before_files": [{"content": "import tensorflow as tf\n\nfrom elasticdl.python.common.tensor import tensor_pb_to_ndarray\nfrom elasticdl.python.ps.embedding_table import create_embedding_table\n\n\nclass Parameters(object):\n    \"\"\"\n    There are two kinds of parameters:\n\n    1. non-embedding parameters, or dense tensor parameters. We save it\n       in a hashmap `non-embedding_params`, the key is the parameter name,\n       the value is a tf.Variable` object.\n    2. embedding parameters, or row-sparse parameters. We save it in a\n       hashmap `embedding_params`, the key is the embedding layer name,\n       the value is an `EmbeddingTable` object.\n\n    \"\"\"\n\n    def __init__(self):\n        self.version = 0\n        self.init_status = False\n        self.non_embedding_params = {}\n        self.embedding_params = {}\n\n    def get_embedding_param(self, name, indices):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        return self.embedding_params[name].get(indices)\n\n    def set_embedding_param(self, name, indices, values):\n        if name not in self.embedding_params:\n            raise ValueError(\n                \"Please initialize embedding param %s first!\", name\n            )\n        self.embedding_params[name].set(indices, values)\n\n    def init_from_model_pb(self, model_pb):\n        if not self.init_status:\n            tensors_pb = model_pb.param\n            embeddings_pb = model_pb.embedding_table_info\n            self._init_non_embedding_params(tensors_pb)\n            self._init_embedding_params(embeddings_pb)\n            self.version = model_pb.version\n            self.init_status = True\n\n    def _init_non_embedding_params(self, tensors_pb):\n        for pb in tensors_pb:\n            name = pb.name\n            arr = tensor_pb_to_ndarray(pb)\n            var = tf.Variable(name=name, initial_value=arr, trainable=True)\n            self.non_embedding_params[name] = var\n\n    def _init_embedding_params(self, embeddings_pb):\n        for pb in embeddings_pb:\n            self.embedding_params[pb.name] = create_embedding_table(pb)\n", "path": "elasticdl/python/ps/parameters.py"}, {"content": "import threading\n\nfrom google.protobuf import empty_pb2\n\nfrom elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc\nfrom elasticdl.python.common.dtypes import dtype_numpy_to_tensor\nfrom elasticdl.python.common.tensor import Tensor, serialize_tensor\n\n\nclass PserverServicer(elasticdl_pb2_grpc.PserverServicer):\n    \"\"\"PS service implementation\"\"\"\n\n    def __init__(\n        self,\n        parameters,\n        grads_to_wait,\n        optimizer,\n        lr_staleness_modulation=False,\n        use_async=False,\n    ):\n        self._parameters = parameters\n        self._grads_to_wait = grads_to_wait\n        self._optimizer = optimizer\n        self._lr_staleness_modulation = lr_staleness_modulation\n        self._use_async = use_async\n        self._version = 0\n        self._lock = threading.Lock()\n\n    def pull_variable(self, request, _):\n        \"\"\"\n        Response with all non-embedding parameters if initialized.\n        \"\"\"\n        res = elasticdl_pb2.PullVariableResponse()\n        if not self._parameters.init_status:\n            res.model_init_status = False\n            return res\n\n        # Only sync-SGD needs lock\n        # TODO: use a read-write lock to support multiple concurrent reads\n        if not self._use_async:\n            self._lock.acquire()\n        res.model.version = self._parameters.version\n        for name, var in self._parameters.non_embedding_params.items():\n            tensor = res.model.param.add()\n            tensor.name = name\n            tensor.dim.extend(var.shape.as_list())\n            var_values = var.numpy()\n            tensor.content = var_values.tobytes()\n            tensor.dtype = dtype_numpy_to_tensor(var_values.dtype)\n        if not self._use_async:\n            self._lock.release()\n        res.model_init_status = True\n        return res\n\n    def pull_embedding_vector(self, request, _):\n        ret = elasticdl_pb2.Tensor()\n        if not request.ids:\n            return ret\n        embedding_vectors = self._parameters.get_embedding_param(\n            request.name, request.ids\n        )\n        tensor = Tensor(values=embedding_vectors)\n        serialize_tensor(tensor, ret)\n        return ret\n\n    def push_model(self, request, _):\n        with self._lock:\n            self._parameters.init_from_model_pb(request)\n        return empty_pb2.Empty()\n\n    def push_gradient(self, request, _):\n        # TODO: implement this RPC service\n        return elasticdl_pb2.PushGradientResponse()\n", "path": "elasticdl/python/ps/servicer.py"}]}