From 470fda190a40fd21d2493d7b92794bed5f288f9e Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Mon, 11 Jan 2021 17:50:06 -0800
Subject: [PATCH] Forgot overwrite parameter in Ray client internal kv

---
 python/ray/tests/test_client.py         | 3 +++
 python/ray/util/client/api.py           | 8 ++++++--
 python/ray/util/client/server/server.py | 2 +-
 python/ray/util/client/worker.py        | 6 ++++--
 src/ray/protobuf/ray_client.proto       | 1 +
 5 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/python/ray/tests/test_client.py b/python/ray/tests/test_client.py
index 39606ed98..5368f9606 100644
--- a/python/ray/tests/test_client.py
+++ b/python/ray/tests/test_client.py
@@ -345,8 +345,11 @@ def test_internal_kv(ray_start_regular_shared):
     with ray_start_client_server() as ray:
         assert ray._internal_kv_initialized()
         assert not ray._internal_kv_put("apple", "b")
+        assert ray._internal_kv_put("apple", "asdf")
         assert ray._internal_kv_put("apple", "b")
         assert ray._internal_kv_get("apple") == b"b"
+        assert ray._internal_kv_put("apple", "asdf", overwrite=True)
+        assert ray._internal_kv_get("apple") == b"asdf"
         assert ray._internal_kv_list("a") == [b"apple"]
         ray._internal_kv_del("apple")
         assert ray._internal_kv_get("apple") == b""
diff --git a/python/ray/util/client/api.py b/python/ray/util/client/api.py
index c6a3a871b..22ae62cc2 100644
--- a/python/ray/util/client/api.py
+++ b/python/ray/util/client/api.py
@@ -240,9 +240,13 @@ class ClientAPI:
         """Hook for internal_kv._internal_kv_get."""
         return self.worker.internal_kv_get(as_bytes(key))
 
-    def _internal_kv_put(self, key: bytes, value: bytes) -> bool:
+    def _internal_kv_put(self,
+                         key: bytes,
+                         value: bytes,
+                         overwrite: bool = False) -> bool:
         """Hook for internal_kv._internal_kv_put."""
-        return self.worker.internal_kv_put(as_bytes(key), as_bytes(value))
+        return self.worker.internal_kv_put(
+            as_bytes(key), as_bytes(value), overwrite)
 
     def _internal_kv_del(self, key: bytes) -> None:
         """Hook for internal_kv._internal_kv_del."""
diff --git a/python/ray/util/client/server/server.py b/python/ray/util/client/server/server.py
index 9c75b9404..8ef8c963f 100644
--- a/python/ray/util/client/server/server.py
+++ b/python/ray/util/client/server/server.py
@@ -41,7 +41,7 @@ class RayletServicer(ray_client_pb2_grpc.RayletDriverServicer):
     def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse:
         with disable_client_hook():
             already_exists = ray.experimental.internal_kv._internal_kv_put(
-                request.key, request.value)
+                request.key, request.value, overwrite=request.overwrite)
         return ray_client_pb2.KVPutResponse(already_exists=already_exists)
 
     def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse:
diff --git a/python/ray/util/client/worker.py b/python/ray/util/client/worker.py
index e86690867..7ef15c717 100644
--- a/python/ray/util/client/worker.py
+++ b/python/ray/util/client/worker.py
@@ -260,8 +260,10 @@ class Worker:
         resp = self.server.KVGet(req, metadata=self.metadata)
         return resp.value
 
-    def internal_kv_put(self, key: bytes, value: bytes) -> bool:
-        req = ray_client_pb2.KVPutRequest(key=key, value=value)
+    def internal_kv_put(self, key: bytes, value: bytes,
+                        overwrite: bool) -> bool:
+        req = ray_client_pb2.KVPutRequest(
+            key=key, value=value, overwrite=overwrite)
         resp = self.server.KVPut(req, metadata=self.metadata)
         return resp.already_exists
 
diff --git a/src/ray/protobuf/ray_client.proto b/src/ray/protobuf/ray_client.proto
index 856815cbf..919e7a226 100644
--- a/src/ray/protobuf/ray_client.proto
+++ b/src/ray/protobuf/ray_client.proto
@@ -189,6 +189,7 @@ message KVGetResponse {
 message KVPutRequest {
   bytes key = 1;
   bytes value = 2;
+  bool overwrite = 3;
 }
 
 message KVPutResponse {