wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-11 11:24:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rpc] refactor GRPC client (#6637)

Browse Source 
* refactor RPC client

* remove unused code

* format

* fix

* resolve comments

* format

* update

* fix

* fix python pb build failure

* lint
This commit is contained in:
Zhijun Fu
2019-12-31 22:28:25 -08:00
committed by Eric Liang
parent f4b29dae9c
commit 91a98d2295
16 changed files with 225 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/dashboard/dashboard.py
+1 -1
View File
@@ -454,7 +454,7 @@ class RayletStats(threading.Thread):
node_id = node["NodeID"]
stub = self.stubs[node_id]
reply = stub.GetNodeStats(
node_manager_pb2.NodeStatsRequest(), timeout=2)
node_manager_pb2.GetNodeStatsRequest(), timeout=2)
replies[node["NodeManagerAddress"]] = reply
with self._raylet_stats_lock:
for address, reply in replies.items():
python/ray/scripts/scripts.py
+1 -1
View File
@@ -860,7 +860,7 @@ def stat(address):
channel = grpc.insecure_channel(raylet_address)
stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel)
reply = stub.GetNodeStats(
node_manager_pb2.NodeStatsRequest(), timeout=2.0)
node_manager_pb2.GetNodeStatsRequest(), timeout=2.0)
print(reply)
python/ray/tests/test_metrics.py
+1 -1
View File
@@ -29,7 +29,7 @@ def test_worker_stats(shutdown_only):
for _ in range(num_retry):
try:
reply = stub.GetNodeStats(
node_manager_pb2.NodeStatsRequest(), timeout=timeout)
node_manager_pb2.GetNodeStatsRequest(), timeout=timeout)
break
except grpc.RpcError:
continue
src/ray/core_worker/test/direct_task_transport_test.cc
+3 -3
View File
@@ -78,7 +78,7 @@ class MockRayletClient : public WorkerLeaseInterface {
ray::Status RequestWorkerLease(
const ray::TaskSpecification &resource_spec,
const rpc::ClientCallback<rpc::WorkerLeaseReply> &callback) override {
const rpc::ClientCallback<rpc::RequestWorkerLeaseReply> &callback) override {
num_workers_requested += 1;
callbacks.push_back(callback);
return Status::OK();
@@ -87,7 +87,7 @@ class MockRayletClient : public WorkerLeaseInterface {
// Trigger reply to RequestWorkerLease.
bool GrantWorkerLease(const std::string &address, int port,
const ClientID &retry_at_raylet_id) {
rpc::WorkerLeaseReply reply;
rpc::RequestWorkerLeaseReply reply;
if (!retry_at_raylet_id.IsNil()) {
reply.mutable_retry_at_raylet_address()->set_ip_address(address);
reply.mutable_retry_at_raylet_address()->set_port(port);
@@ -112,7 +112,7 @@ class MockRayletClient : public WorkerLeaseInterface {
int num_workers_requested = 0;
int num_workers_returned = 0;
int num_workers_disconnected = 0;
std::list<rpc::ClientCallback<rpc::WorkerLeaseReply>> callbacks = {};
std::list<rpc::ClientCallback<rpc::RequestWorkerLeaseReply>> callbacks = {};
};
TEST(TestMemoryStore, TestPromoteToPlasma) {
src/ray/core_worker/transport/direct_task_transport.cc
+1 -1
View File
@@ -108,7 +108,7 @@ void CoreWorkerDirectTaskSubmitter::RequestNewWorkerIfNeeded(
auto status = lease_client->RequestWorkerLease(
resource_spec,
[this, lease_client, task_id, scheduling_key](
const Status &status, const rpc::WorkerLeaseReply &reply) mutable {
const Status &status, const rpc::RequestWorkerLeaseReply &reply) mutable {
absl::MutexLock lock(&mu_);
pending_lease_requests_.erase(scheduling_key);
if (status.ok()) {
src/ray/protobuf/node_manager.proto
+6 -6
View File
@@ -5,12 +5,12 @@ package ray.rpc;
import "src/ray/protobuf/common.proto";
// Request a worker from the raylet with the specified resources.
message WorkerLeaseRequest {
message RequestWorkerLeaseRequest {
// TaskSpec containing the requested resources.
TaskSpec resource_spec = 1;
}
message WorkerLeaseReply {
message RequestWorkerLeaseReply {
// Address of the leased worker. If this is empty, then the request should be
// retried at the provided raylet address.
Address worker_address = 1;
@@ -45,7 +45,7 @@ message ForwardTaskRequest {
message ForwardTaskReply {
}
message NodeStatsRequest {
message GetNodeStatsRequest {
}
message WorkerStats {
@@ -81,7 +81,7 @@ message ViewData {
repeated Measure measures = 2;
}
message NodeStatsReply {
message GetNodeStatsReply {
repeated WorkerStats workers_stats = 1;
repeated ViewData view_data = 2;
map<string, double> available_resources = 3;
@@ -92,11 +92,11 @@ message NodeStatsReply {
// Service for inter-node-manager communication.
service NodeManagerService {
// Request a worker from the raylet.
rpc RequestWorkerLease(WorkerLeaseRequest) returns (WorkerLeaseReply);
rpc RequestWorkerLease(RequestWorkerLeaseRequest) returns (RequestWorkerLeaseReply);
// Release a worker back to its raylet.
rpc ReturnWorker(ReturnWorkerRequest) returns (ReturnWorkerReply);
// Forward a task and its uncommitted lineage to the remote node manager.
rpc ForwardTask(ForwardTaskRequest) returns (ForwardTaskReply);
// Get the current node stats.
rpc GetNodeStats(NodeStatsRequest) returns (NodeStatsReply);
rpc GetNodeStats(GetNodeStatsRequest) returns (GetNodeStatsReply);
}
src/ray/raylet/node_manager.cc
+4 -4
View File
@@ -1497,8 +1497,8 @@ void NodeManager::NewSchedulerSchedulePendingTasks() {
DispatchScheduledTasksToWorkers();
}
void NodeManager::HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request,
rpc::WorkerLeaseReply *reply,
void NodeManager::HandleWorkerLeaseRequest(const rpc::RequestWorkerLeaseRequest &request,
rpc::RequestWorkerLeaseReply *reply,
rpc::SendReplyCallback send_reply_callback) {
rpc::Task task_message;
task_message.mutable_task_spec()->CopyFrom(request.resource_spec());
@@ -2924,8 +2924,8 @@ std::string compact_tag_string(const opencensus::stats::ViewDescriptor &view,
return result.str();
}
void NodeManager::HandleNodeStatsRequest(const rpc::NodeStatsRequest &request,
rpc::NodeStatsReply *reply,
void NodeManager::HandleNodeStatsRequest(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
for (const auto &driver : worker_pool_.GetAllDrivers()) {
auto worker_stats = reply->add_workers_stats();
src/ray/raylet/node_manager.h
+5 -5
View File
@@ -4,7 +4,7 @@
#include <boost/asio/steady_timer.hpp>
// clang-format off
#include "ray/rpc/client_call.h"
#include "ray/rpc/grpc_client.h"
#include "ray/rpc/node_manager/node_manager_server.h"
#include "ray/rpc/node_manager/node_manager_client.h"
#include "ray/common/task/task.h"
@@ -531,8 +531,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
bool success);
/// Handle a `WorkerLease` request.
void HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request,
rpc::WorkerLeaseReply *reply,
void HandleWorkerLeaseRequest(const rpc::RequestWorkerLeaseRequest &request,
rpc::RequestWorkerLeaseReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Handle a `ReturnWorker` request.
@@ -546,8 +546,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
rpc::SendReplyCallback send_reply_callback) override;
/// Handle a `NodeStats` request.
void HandleNodeStatsRequest(const rpc::NodeStatsRequest &request,
rpc::NodeStatsReply *reply,
void HandleNodeStatsRequest(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Push an error to the driver if this node is full of actors and so we are
src/ray/raylet/raylet_client.cc
+2 -2
View File
@@ -408,8 +408,8 @@ Status raylet::RayletClient::ReportActiveObjectIDs(
Status raylet::RayletClient::RequestWorkerLease(
const TaskSpecification &resource_spec,
const rpc::ClientCallback<rpc::WorkerLeaseReply> &callback) {
rpc::WorkerLeaseRequest request;
const rpc::ClientCallback<rpc::RequestWorkerLeaseReply> &callback) {
rpc::RequestWorkerLeaseRequest request;
request.mutable_resource_spec()->CopyFrom(resource_spec.GetMessage());
return grpc_client_->RequestWorkerLease(request, callback);
}
src/ray/raylet/raylet_client.h
+3 -2
View File
@@ -40,7 +40,7 @@ class WorkerLeaseInterface {
/// \return ray::Status
virtual ray::Status RequestWorkerLease(
const ray::TaskSpecification &resource_spec,
const ray::rpc::ClientCallback<ray::rpc::WorkerLeaseReply> &callback) = 0;
const ray::rpc::ClientCallback<ray::rpc::RequestWorkerLeaseReply> &callback) = 0;
/// Returns a worker to the raylet.
/// \param worker_port The local port of the worker on the raylet node.
@@ -242,7 +242,8 @@ class RayletClient : public WorkerLeaseInterface {
/// Implements WorkerLeaseInterface.
ray::Status RequestWorkerLease(
const ray::TaskSpecification &resource_spec,
const ray::rpc::ClientCallback<ray::rpc::WorkerLeaseReply> &callback) override;
const ray::rpc::ClientCallback<ray::rpc::RequestWorkerLeaseReply> &callback)
override;
/// Implements WorkerLeaseInterface.
ray::Status ReturnWorker(int worker_port, const WorkerID &worker_id,
src/ray/rpc/gcs_server/gcs_rpc_client.h
+35 -118
View File
@@ -6,7 +6,7 @@
#include <grpcpp/grpcpp.h>
#include "src/ray/protobuf/gcs_service.pb.h"
#include "src/ray/rpc/client_call.h"
#include "src/ray/rpc/grpc_client.h"
namespace ray {
namespace rpc {
@@ -22,149 +22,66 @@ class GcsRpcClient {
GcsRpcClient(const std::string &address, const int port,
ClientCallManager &client_call_manager)
: client_call_manager_(client_call_manager) {
std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
job_info_stub_ = JobInfoGcsService::NewStub(channel);
actor_info_stub_ = ActorInfoGcsService::NewStub(channel);
node_info_stub_ = NodeInfoGcsService::NewStub(channel);
object_info_stub_ = ObjectInfoGcsService::NewStub(channel);
job_info_grpc_client_ = std::unique_ptr<GrpcClient<JobInfoGcsService>>(
new GrpcClient<JobInfoGcsService>(address, port, client_call_manager));
actor_info_grpc_client_ = std::unique_ptr<GrpcClient<ActorInfoGcsService>>(
new GrpcClient<ActorInfoGcsService>(address, port, client_call_manager));
node_info_grpc_client_ = std::unique_ptr<GrpcClient<NodeInfoGcsService>>(
new GrpcClient<NodeInfoGcsService>(address, port, client_call_manager));
object_info_grpc_client_ = std::unique_ptr<GrpcClient<ObjectInfoGcsService>>(
new GrpcClient<ObjectInfoGcsService>(address, port, client_call_manager));
};
/// Add job info to gcs server.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void AddJob(const AddJobRequest &request, const ClientCallback<AddJobReply> &callback) {
client_call_manager_.CreateCall<JobInfoGcsService, AddJobRequest, AddJobReply>(
*job_info_stub_, &JobInfoGcsService::Stub::PrepareAsyncAddJob, request, callback);
}
VOID_RPC_CLIENT_METHOD(JobInfoGcsService, AddJob, request, callback,
job_info_grpc_client_)
/// Mark job as finished to gcs server.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void MarkJobFinished(const MarkJobFinishedRequest &request,
const ClientCallback<MarkJobFinishedReply> &callback) {
client_call_manager_
.CreateCall<JobInfoGcsService, MarkJobFinishedRequest, MarkJobFinishedReply>(
*job_info_stub_, &JobInfoGcsService::Stub::PrepareAsyncMarkJobFinished,
request, callback);
}
VOID_RPC_CLIENT_METHOD(JobInfoGcsService, MarkJobFinished, request, callback,
job_info_grpc_client_)
/// Get actor data from GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void GetActorInfo(const GetActorInfoRequest &request,
const ClientCallback<GetActorInfoReply> &callback) {
client_call_manager_
.CreateCall<ActorInfoGcsService, GetActorInfoRequest, GetActorInfoReply>(
*actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncGetActorInfo,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, GetActorInfo, request, callback,
actor_info_grpc_client_)
/// Register an actor to GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void RegisterActorInfo(const RegisterActorInfoRequest &request,
const ClientCallback<RegisterActorInfoReply> &callback) {
client_call_manager_.CreateCall<ActorInfoGcsService, RegisterActorInfoRequest,
RegisterActorInfoReply>(
*actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncRegisterActorInfo,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, RegisterActorInfo, request, callback,
actor_info_grpc_client_)
/// Update actor info in GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void UpdateActorInfo(const UpdateActorInfoRequest &request,
const ClientCallback<UpdateActorInfoReply> &callback) {
client_call_manager_
.CreateCall<ActorInfoGcsService, UpdateActorInfoRequest, UpdateActorInfoReply>(
*actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncUpdateActorInfo,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, UpdateActorInfo, request, callback,
actor_info_grpc_client_)
/// Register a node to GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void RegisterNode(const RegisterNodeRequest &request,
const ClientCallback<RegisterNodeReply> &callback) {
client_call_manager_
.CreateCall<NodeInfoGcsService, RegisterNodeRequest, RegisterNodeReply>(
*node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncRegisterNode,
request, callback);
}
VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, RegisterNode, request, callback,
node_info_grpc_client_)
/// Unregister a node from GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void UnregisterNode(const UnregisterNodeRequest &request,
const ClientCallback<UnregisterNodeReply> &callback) {
client_call_manager_
.CreateCall<NodeInfoGcsService, UnregisterNodeRequest, UnregisterNodeReply>(
*node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncUnregisterNode,
request, callback);
}
VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, UnregisterNode, request, callback,
node_info_grpc_client_)
/// Get information of all nodes from GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void GetAllNodeInfo(const GetAllNodeInfoRequest &request,
const ClientCallback<GetAllNodeInfoReply> &callback) {
client_call_manager_
.CreateCall<NodeInfoGcsService, GetAllNodeInfoRequest, GetAllNodeInfoReply>(
*node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncGetAllNodeInfo,
request, callback);
}
VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, GetAllNodeInfo, request, callback,
node_info_grpc_client_)
/// Get object's locations from GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void GetObjectLocations(const GetObjectLocationsRequest &request,
const ClientCallback<GetObjectLocationsReply> &callback) {
client_call_manager_.CreateCall<ObjectInfoGcsService, GetObjectLocationsRequest,
GetObjectLocationsReply>(
*object_info_stub_, &ObjectInfoGcsService::Stub::PrepareAsyncGetObjectLocations,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, GetObjectLocations, request, callback,
object_info_grpc_client_)
/// Add location of object to GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void AddObjectLocation(const AddObjectLocationRequest &request,
const ClientCallback<AddObjectLocationReply> &callback) {
client_call_manager_.CreateCall<ObjectInfoGcsService, AddObjectLocationRequest,
AddObjectLocationReply>(
*object_info_stub_, &ObjectInfoGcsService::Stub::PrepareAsyncAddObjectLocation,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, AddObjectLocation, request, callback,
object_info_grpc_client_)
/// Remove location of object to GCS Service.
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server.
void RemoveObjectLocation(const RemoveObjectLocationRequest &request,
const ClientCallback<RemoveObjectLocationReply> &callback) {
client_call_manager_.CreateCall<ObjectInfoGcsService, RemoveObjectLocationRequest,
RemoveObjectLocationReply>(
*object_info_stub_, &ObjectInfoGcsService::Stub::PrepareAsyncRemoveObjectLocation,
request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, RemoveObjectLocation, request, callback,
object_info_grpc_client_)
private:
/// The gRPC-generated stub.
std::unique_ptr<JobInfoGcsService::Stub> job_info_stub_;
std::unique_ptr<ActorInfoGcsService::Stub> actor_info_stub_;
std::unique_ptr<NodeInfoGcsService::Stub> node_info_stub_;
std::unique_ptr<ObjectInfoGcsService::Stub> object_info_stub_;
std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;
std::unique_ptr<GrpcClient<ActorInfoGcsService>> actor_info_grpc_client_;
std::unique_ptr<GrpcClient<NodeInfoGcsService>> node_info_grpc_client_;
std::unique_ptr<GrpcClient<ObjectInfoGcsService>> object_info_grpc_client_;
/// The `ClientCallManager` used for managing requests.
ClientCallManager &client_call_manager_;
src/ray/rpc/grpc_client.h
+88
View File
@@ -0,0 +1,88 @@
#ifndef RAY_RPC_GRPC_CLIENT_H
#define RAY_RPC_GRPC_CLIENT_H
#include <grpcpp/grpcpp.h>
#include <boost/asio.hpp>
#include "ray/common/grpc_util.h"
#include "ray/common/status.h"
#include "ray/rpc/client_call.h"
namespace ray {
namespace rpc {
// This macro wraps the logic to call a specific RPC method of a service,
// to make it easier to implement a new RPC client.
#define INVOKE_RPC_CALL(SERVICE, METHOD, request, callback, rpc_client) \
({ \
rpc_client->CallMethod<METHOD##Request, METHOD##Reply>( \
&SERVICE::Stub::PrepareAsync##METHOD, request, callback); \
})
// Define a void RPC client method.
#define VOID_RPC_CLIENT_METHOD(SERVICE, METHOD, request, callback, rpc_client) \
void METHOD(const METHOD##Request &request, \
const ClientCallback<METHOD##Reply> &callback) { \
RAY_UNUSED(INVOKE_RPC_CALL(SERVICE, METHOD, request, callback, rpc_client)); \
}
// Define a RPC client method that returns ray::Status.
#define RPC_CLIENT_METHOD(SERVICE, METHOD, request, callback, rpc_client) \
ray::Status METHOD(const METHOD##Request &request, \
const ClientCallback<METHOD##Reply> &callback) { \
return INVOKE_RPC_CALL(SERVICE, METHOD, request, callback, rpc_client); \
}
template <class GrpcService>
class GrpcClient {
public:
GrpcClient(const std::string &address, const int port, ClientCallManager &call_manager)
: client_call_manager_(call_manager) {
std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
stub_ = GrpcService::NewStub(channel);
}
GrpcClient(const std::string &address, const int port, ClientCallManager &call_manager,
int num_threads)
: client_call_manager_(call_manager) {
grpc::ResourceQuota quota;
quota.SetMaxThreads(num_threads);
grpc::ChannelArguments argument;
argument.SetResourceQuota(quota);
std::shared_ptr<grpc::Channel> channel =
grpc::CreateCustomChannel(address + ":" + std::to_string(port),
grpc::InsecureChannelCredentials(), argument);
stub_ = GrpcService::NewStub(channel);
}
/// Create a new `ClientCall` and send request.
///
/// \tparam Request Type of the request message.
/// \tparam Reply Type of the reply message.
///
/// \param[in] prepare_async_function Pointer to the gRPC-generated
/// `FooService::Stub::PrepareAsyncBar` function.
/// \param[in] request The request message.
/// \param[in] callback The callback function that handles reply.
///
/// \return Status.
template <class Request, class Reply>
ray::Status CallMethod(
const PrepareAsyncFunction<GrpcService, Request, Reply> prepare_async_function,
const Request &request, const ClientCallback<Reply> &callback) {
auto call = client_call_manager_.CreateCall<GrpcService, Request, Reply>(
*stub_, prepare_async_function, request, callback);
return call->GetStatus();
}
private:
ClientCallManager &client_call_manager_;
/// The gRPC-generated stub.
std::unique_ptr<typename GrpcService::Stub> stub_;
};
} // namespace rpc
} // namespace ray
#endif
src/ray/rpc/node_manager/node_manager_client.h
+20 -37
View File
@@ -6,7 +6,7 @@
#include <grpcpp/grpcpp.h>
#include "ray/common/status.h"
#include "ray/rpc/client_call.h"
#include "ray/rpc/grpc_client.h"
#include "ray/util/logging.h"
#include "src/ray/protobuf/node_manager.grpc.pb.h"
#include "src/ray/protobuf/node_manager.pb.h"
@@ -25,33 +25,28 @@ class NodeManagerClient {
NodeManagerClient(const std::string &address, const int port,
ClientCallManager &client_call_manager)
: client_call_manager_(client_call_manager) {
std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
stub_ = NodeManagerService::NewStub(channel);
grpc_client_ = std::unique_ptr<GrpcClient<NodeManagerService>>(
new GrpcClient<NodeManagerService>(address, port, client_call_manager));
};
/// Forward a task and its uncommitted lineage.
///
/// \param[in] request The request message.
/// \param[in] callback The callback function that handles reply.
void ForwardTask(const ForwardTaskRequest &request,
const ClientCallback<ForwardTaskReply> &callback) {
client_call_manager_
.CreateCall<NodeManagerService, ForwardTaskRequest, ForwardTaskReply>(
*stub_, &NodeManagerService::Stub::PrepareAsyncForwardTask, request,
callback);
}
VOID_RPC_CLIENT_METHOD(NodeManagerService, ForwardTask, request, callback, grpc_client_)
/// Get current node stats.
void GetNodeStats(const ClientCallback<NodeStatsReply> &callback) {
NodeStatsRequest request;
client_call_manager_.CreateCall<NodeManagerService, NodeStatsRequest, NodeStatsReply>(
*stub_, &NodeManagerService::Stub::PrepareAsyncGetNodeStats, request, callback);
VOID_RPC_CLIENT_METHOD(NodeManagerService, GetNodeStats, request, callback,
grpc_client_)
void GetNodeStats(const ClientCallback<GetNodeStatsReply> &callback) {
GetNodeStatsRequest request;
GetNodeStats(request, callback);
}
private:
/// The gRPC-generated stub.
std::unique_ptr<NodeManagerService::Stub> stub_;
/// The RPC client.
std::unique_ptr<GrpcClient<NodeManagerService>> grpc_client_;
/// The `ClientCallManager` used for managing requests.
ClientCallManager &client_call_manager_;
@@ -74,22 +69,11 @@ class NodeManagerWorkerClient
}
/// Request a worker lease.
ray::Status RequestWorkerLease(const WorkerLeaseRequest &request,
const ClientCallback<WorkerLeaseReply> &callback) {
auto call = client_call_manager_
.CreateCall<NodeManagerService, WorkerLeaseRequest, WorkerLeaseReply>(
*stub_, &NodeManagerService::Stub::PrepareAsyncRequestWorkerLease,
request, callback);
return call->GetStatus();
}
RPC_CLIENT_METHOD(NodeManagerService, RequestWorkerLease, request, callback,
grpc_client_)
ray::Status ReturnWorker(const ReturnWorkerRequest &request,
const ClientCallback<ReturnWorkerReply> &callback) {
auto call = client_call_manager_.CreateCall<NodeManagerService, ReturnWorkerRequest,
ReturnWorkerReply>(
*stub_, &NodeManagerService::Stub::PrepareAsyncReturnWorker, request, callback);
return call->GetStatus();
}
/// Return a worker lease.
RPC_CLIENT_METHOD(NodeManagerService, ReturnWorker, request, callback, grpc_client_)
private:
/// Constructor.
@@ -100,13 +84,12 @@ class NodeManagerWorkerClient
NodeManagerWorkerClient(const std::string &address, const int port,
ClientCallManager &client_call_manager)
: client_call_manager_(client_call_manager) {
std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
stub_ = NodeManagerService::NewStub(channel);
grpc_client_ = std::unique_ptr<GrpcClient<NodeManagerService>>(
new GrpcClient<NodeManagerService>(address, port, client_call_manager));
};
/// The gRPC-generated stub.
std::unique_ptr<NodeManagerService::Stub> stub_;
/// The RPC client.
std::unique_ptr<GrpcClient<NodeManagerService>> grpc_client_;
/// The `ClientCallManager` used for managing requests.
ClientCallManager &client_call_manager_;
src/ray/rpc/node_manager/node_manager_server.h
+6 -6
View File
@@ -24,8 +24,8 @@ class NodeManagerServiceHandler {
/// \param[out] reply The reply message.
/// \param[in] send_reply_callback The callback to be called when the request is done.
virtual void HandleWorkerLeaseRequest(const WorkerLeaseRequest &request,
WorkerLeaseReply *reply,
virtual void HandleWorkerLeaseRequest(const RequestWorkerLeaseRequest &request,
RequestWorkerLeaseReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandleReturnWorker(const ReturnWorkerRequest &request,
@@ -36,8 +36,8 @@ class NodeManagerServiceHandler {
ForwardTaskReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandleNodeStatsRequest(const NodeStatsRequest &request,
NodeStatsReply *reply,
virtual void HandleNodeStatsRequest(const GetNodeStatsRequest &request,
GetNodeStatsReply *reply,
SendReplyCallback send_reply_callback) = 0;
};
@@ -62,7 +62,7 @@ class NodeManagerGrpcService : public GrpcService {
// Initialize the factory for requests.
std::unique_ptr<ServerCallFactory> request_worker_lease_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
WorkerLeaseRequest, WorkerLeaseReply>(
RequestWorkerLeaseRequest, RequestWorkerLeaseReply>(
service_, &NodeManagerService::AsyncService::RequestRequestWorkerLease,
service_handler_, &NodeManagerServiceHandler::HandleWorkerLeaseRequest, cq,
main_service_));
@@ -83,7 +83,7 @@ class NodeManagerGrpcService : public GrpcService {
std::unique_ptr<ServerCallFactory> node_stats_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
NodeStatsRequest, NodeStatsReply>(
GetNodeStatsRequest, GetNodeStatsReply>(
service_, &NodeManagerService::AsyncService::RequestGetNodeStats,
service_handler_, &NodeManagerServiceHandler::HandleNodeStatsRequest, cq,
main_service_));
src/ray/rpc/object_manager/object_manager_client.h
+17 -29
View File
@@ -11,7 +11,7 @@
#include "ray/util/logging.h"
#include "src/ray/protobuf/object_manager.grpc.pb.h"
#include "src/ray/protobuf/object_manager.pb.h"
#include "src/ray/rpc/client_call.h"
#include "src/ray/rpc/grpc_client.h"
namespace ray {
namespace rpc {
@@ -30,16 +30,10 @@ class ObjectManagerClient {
push_rr_index_ = rand() % num_connections_;
pull_rr_index_ = rand() % num_connections_;
freeobjects_rr_index_ = rand() % num_connections_;
stubs_.reserve(num_connections_);
grpc_clients_.reserve(num_connections_);
for (int i = 0; i < num_connections_; i++) {
grpc::ResourceQuota quota;
quota.SetMaxThreads(num_connections_);
grpc::ChannelArguments argument;
argument.SetResourceQuota(quota);
std::shared_ptr<grpc::Channel> channel =
grpc::CreateCustomChannel(address + ":" + std::to_string(port),
grpc::InsecureChannelCredentials(), argument);
stubs_.push_back(ObjectManagerService::NewStub(channel));
grpc_clients_.emplace_back(new GrpcClient<ObjectManagerService>(
address, port, client_call_manager, num_connections_));
}
};
@@ -47,43 +41,37 @@ class ObjectManagerClient {
///
/// \param request The request message.
/// \param callback The callback function that handles reply from server
void Push(const PushRequest &request, const ClientCallback<PushReply> &callback) {
client_call_manager_.CreateCall<ObjectManagerService, PushRequest, PushReply>(
*stubs_[push_rr_index_++ % num_connections_],
&ObjectManagerService::Stub::PrepareAsyncPush, request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectManagerService, Push, request, callback,
grpc_clients_[push_rr_index_++ % num_connections_])
/// Pull object from remote object manager
///
/// \param request The request message
/// \param callback The callback function that handles reply from server
void Pull(const PullRequest &request, const ClientCallback<PullReply> &callback) {
client_call_manager_.CreateCall<ObjectManagerService, PullRequest, PullReply>(
*stubs_[pull_rr_index_++ % num_connections_],
&ObjectManagerService::Stub::PrepareAsyncPull, request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectManagerService, Pull, request, callback,
grpc_clients_[pull_rr_index_++ % num_connections_])
/// Tell remote object manager to free objects
///
/// \param request The request message
/// \param callback The callback function that handles reply
void FreeObjects(const FreeObjectsRequest &request,
const ClientCallback<FreeObjectsReply> &callback) {
client_call_manager_
.CreateCall<ObjectManagerService, FreeObjectsRequest, FreeObjectsReply>(
*stubs_[freeobjects_rr_index_++ % num_connections_],
&ObjectManagerService::Stub::PrepareAsyncFreeObjects, request, callback);
}
VOID_RPC_CLIENT_METHOD(ObjectManagerService, FreeObjects, request, callback,
grpc_clients_[freeobjects_rr_index_++ % num_connections_])
private:
/// To optimize object manager performance we create multiple concurrent
/// GRPC connections, and use these connections in a round-robin way.
int num_connections_;
/// Current connection index for `Push`.
std::atomic<unsigned int> push_rr_index_;
/// Current connection index for `Pull`.
std::atomic<unsigned int> pull_rr_index_;
/// Current connection index for `FreeObjects`.
std::atomic<unsigned int> freeobjects_rr_index_;
/// The gRPC-generated stub.
std::vector<std::unique_ptr<ObjectManagerService::Stub>> stubs_;
/// The RPC clients.
std::vector<std::unique_ptr<GrpcClient<ObjectManagerService>>> grpc_clients_;
/// The `ClientCallManager` used for managing requests.
ClientCallManager &client_call_manager_;
src/ray/rpc/worker/core_worker_client.h
+32 -74
View File
@@ -11,7 +11,7 @@
#include "absl/base/thread_annotations.h"
#include "absl/hash/hash.h"
#include "ray/common/status.h"
#include "ray/rpc/client_call.h"
#include "ray/rpc/grpc_client.h"
#include "ray/util/logging.h"
#include "src/ray/protobuf/core_worker.grpc.pb.h"
#include "src/ray/protobuf/core_worker.pb.h"
@@ -148,19 +148,21 @@ class CoreWorkerClient : public std::enable_shared_from_this<CoreWorkerClient>,
CoreWorkerClient(const std::string &address, const int port,
ClientCallManager &client_call_manager)
: client_call_manager_(client_call_manager) {
std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
stub_ = CoreWorkerService::NewStub(channel);
grpc_client_ = std::unique_ptr<GrpcClient<CoreWorkerService>>(
new GrpcClient<CoreWorkerService>(address, port, client_call_manager));
};
ray::Status AssignTask(const AssignTaskRequest &request,
const ClientCallback<AssignTaskReply> &callback) override {
auto call = client_call_manager_
.CreateCall<CoreWorkerService, AssignTaskRequest, AssignTaskReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncAssignTask, request,
callback);
return call->GetStatus();
}
RPC_CLIENT_METHOD(CoreWorkerService, AssignTask, request, callback, grpc_client_)
RPC_CLIENT_METHOD(CoreWorkerService, DirectActorCallArgWaitComplete, request, callback,
grpc_client_)
RPC_CLIENT_METHOD(CoreWorkerService, GetObjectStatus, request, callback, grpc_client_)
RPC_CLIENT_METHOD(CoreWorkerService, KillActor, request, callback, grpc_client_)
RPC_CLIENT_METHOD(CoreWorkerService, GetCoreWorkerStats, request, callback,
grpc_client_)
ray::Status PushActorTask(std::unique_ptr<PushTaskRequest> request,
const ClientCallback<PushTaskReply> &callback) override {
@@ -182,51 +184,7 @@ class CoreWorkerClient : public std::enable_shared_from_this<CoreWorkerClient>,
const ClientCallback<PushTaskReply> &callback) override {
request->set_sequence_number(-1);
request->set_client_processed_up_to(-1);
auto call = client_call_manager_
.CreateCall<CoreWorkerService, PushTaskRequest, PushTaskReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncPushTask, *request,
callback);
return call->GetStatus();
}
ray::Status DirectActorCallArgWaitComplete(
const DirectActorCallArgWaitCompleteRequest &request,
const ClientCallback<DirectActorCallArgWaitCompleteReply> &callback) override {
auto call = client_call_manager_.CreateCall<CoreWorkerService,
DirectActorCallArgWaitCompleteRequest,
DirectActorCallArgWaitCompleteReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncDirectActorCallArgWaitComplete,
request, callback);
return call->GetStatus();
}
virtual ray::Status GetObjectStatus(
const GetObjectStatusRequest &request,
const ClientCallback<GetObjectStatusReply> &callback) override {
auto call = client_call_manager_.CreateCall<CoreWorkerService, GetObjectStatusRequest,
GetObjectStatusReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncGetObjectStatus, request, callback);
return call->GetStatus();
}
virtual ray::Status KillActor(const KillActorRequest &request,
const ClientCallback<KillActorReply> &callback) override {
auto call = client_call_manager_
.CreateCall<CoreWorkerService, KillActorRequest, KillActorReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncKillActor, request,
callback);
return call->GetStatus();
}
virtual ray::Status GetCoreWorkerStats(
const GetCoreWorkerStatsRequest &request,
const ClientCallback<GetCoreWorkerStatsReply> &callback) override {
auto call =
client_call_manager_.CreateCall<CoreWorkerService, GetCoreWorkerStatsRequest,
GetCoreWorkerStatsReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncGetCoreWorkerStats, request,
callback);
return call->GetStatus();
return INVOKE_RPC_CALL(CoreWorkerService, PushTask, *request, callback, grpc_client_);
}
/// Send as many pending tasks as possible. This method is thread-safe.
@@ -249,21 +207,21 @@ class CoreWorkerClient : public std::enable_shared_from_this<CoreWorkerClient>,
request->set_client_processed_up_to(max_finished_seq_no_);
rpc_bytes_in_flight_ += task_size;
client_call_manager_.CreateCall<CoreWorkerService, PushTaskRequest, PushTaskReply>(
*stub_, &CoreWorkerService::Stub::PrepareAsyncPushTask, *request,
[this, this_ptr, seq_no, task_size, callback](Status status,
const rpc::PushTaskReply &reply) {
{
std::lock_guard<std::mutex> lock(mutex_);
if (seq_no > max_finished_seq_no_) {
max_finished_seq_no_ = seq_no;
}
rpc_bytes_in_flight_ -= task_size;
RAY_CHECK(rpc_bytes_in_flight_ >= 0);
}
SendRequests();
callback(status, reply);
});
auto rpc_callback = [this, this_ptr, seq_no, task_size, callback](
Status status, const rpc::PushTaskReply &reply) {
{
std::lock_guard<std::mutex> lock(mutex_);
if (seq_no > max_finished_seq_no_) {
max_finished_seq_no_ = seq_no;
}
rpc_bytes_in_flight_ -= task_size;
RAY_CHECK(rpc_bytes_in_flight_ >= 0);
}
SendRequests();
callback(status, reply);
};
INVOKE_RPC_CALL(CoreWorkerService, PushTask, *request, rpc_callback, grpc_client_);
}
if (!send_queue_.empty()) {
@@ -275,8 +233,8 @@ class CoreWorkerClient : public std::enable_shared_from_this<CoreWorkerClient>,
/// Protects against unsafe concurrent access from the callback thread.
std::mutex mutex_;
/// The gRPC-generated stub.
std::unique_ptr<CoreWorkerService::Stub> stub_;
/// The RPC client.
std::unique_ptr<GrpcClient<CoreWorkerService>> grpc_client_;
/// The `ClientCallManager` used for managing requests.
ClientCallManager &client_call_manager_;