[ray_client]: Implement object retain/release and Data Streaming API (#12818)

src/ray/protobuf/ray_client.proto

@@ -18,17 +18,24 @@ package ray.rpc;
 
 enum Type { DEFAULT = 0; }
 
+// An argument to a ClientTask.
 message Arg {
   enum Locality {
     INTERNED = 0;
     REFERENCE = 1;
   }
+
+  // The type of argument this is -- whether a data blob or a reference.
   Locality local = 1;
+  // The reference id, if a reference.
   bytes reference_id = 2;
+  // A data blob, if passed in-band.
   bytes data = 3;
+  // How to decode this data blob.
   Type type = 4;
 }
 
+// Represents one unit of work to be executed by the server.
 message ClientTask {
   enum RemoteExecType {
     FUNCTION = 0;
@@ -36,49 +43,69 @@ message ClientTask {
     METHOD = 2;
     STATIC_METHOD = 3;
   }
+  // Which type of work this request represents.
   RemoteExecType type = 1;
+  // A name parameter, if the payload can be called in more than one way (like a method on
+  // a payload object).
   string name = 2;
+  // A reference to the payload.
   bytes payload_id = 3;
+  // The parameters to pass to this call.
   repeated Arg args = 4;
-}
-
-message RemoteRef {
-  bytes id = 1;
-  bytes handle = 2;
+  // The ID of the client namespace associated with the Datapath stream making this
+  // request.
+  string client_id = 5;
 }
 
 message ClientTaskTicket {
-  RemoteRef return_ref = 1;
+  // A reference to the returned value from the execution.
+  bytes return_id = 1;
 }
 
+// Delivers data to the server
 message PutRequest {
+  // The data blob for the server to store.
   bytes data = 1;
 }
 
 message PutResponse {
-  RemoteRef ref = 1;
+  // The reference ID for the data that the server has stored.
+  bytes id = 1;
 }
 
+// Requests data from the server.
 message GetRequest {
-  bytes handle = 1;
+  // The reference ID for the requested object data
+  bytes id = 1;
+  // Length of time to wait for data to be available, in seconds. Zero is no timeout.
   float timeout = 2;
 }
 
 message GetResponse {
+  // Whether or not the data was successfully retrieved
   bool valid = 1;
+  // The data blob, on success
   bytes data = 2;
+  // An error blob (for example, an exception) on failure.
+  bytes error = 3;
 }
 
+// Waits for data to be ready on the server, with a timeout.
 message WaitRequest {
-  repeated bytes object_handles = 1;
+  // The IDs of the data to wait for ready status.
+  repeated bytes object_ids = 1;
+  // How many of the above ids to wait for before returning.
   int64 num_returns = 2;
+  // How long to wait for these IDs to become ready.
   double timeout = 3;
+  // The Client namespace associated with the Datapath stream that holds these IDs.
+  string client_id = 4;
 }
 
 message WaitResponse {
   bool valid = 1;
-  repeated RemoteRef ready_object_ids = 2;
-  repeated RemoteRef remaining_object_ids = 3;
+  repeated bytes ready_object_ids = 2;
+  repeated bytes remaining_object_ids = 3;
 }
 
 message ClusterInfoType {
@@ -108,18 +135,19 @@ message ClusterInfoResponse {
 
 message TerminateRequest {
   message ActorTerminate {
-    bytes handle = 1;
+    bytes id = 1;
     bool no_restart = 2;
   }
   message TaskObjectTerminate {
-    bytes handle = 1;
+    bytes id = 1;
     bool force = 2;
     bool recursive = 3;
   }
 
+  string client_id = 1;
   oneof terminate_type {
-    ActorTerminate actor = 1;
-    TaskObjectTerminate task_object = 2;
+    ActorTerminate actor = 2;
+    TaskObjectTerminate task_object = 3;
   }
 }
 
@@ -141,3 +169,40 @@ service RayletDriver {
   rpc ClusterInfo(ClusterInfoRequest) returns (ClusterInfoResponse) {
   }
 }
+
+message ReleaseRequest {
+  // The IDs to release from the server; the client connected on this stream no
+  // longer holds a reference to them.
+  repeated bytes ids = 1;
+}
+
+message ReleaseResponse {
+  // For each requested ID, whether or not it was released.
+  repeated bool ok = 2;
+}
+
+message DataRequest {
+  // An incrementing counter of request IDs on the Datapath,
+  // to match requests with responses asynchronously.
+  int32 req_id = 1;
+  oneof type {
+    GetRequest get = 2;
+    PutRequest put = 3;
+    ReleaseRequest release = 4;
+  }
+}
+
+message DataResponse {
+  // The request id that this response matches with.
+  int32 req_id = 1;
+  oneof type {
+    GetResponse get = 2;
+    PutResponse put = 3;
+    ReleaseResponse release = 4;
+  }
+}
+
+service RayletDataStreamer {
+  rpc Datapath(stream DataRequest) returns (stream DataResponse) {
+  }
+}