Forgot overwrite parameter in Ray client internal kv

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""

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."""

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:

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