diff --git a/src/ray/common/grpc_util.h b/src/ray/common/grpc_util.h index bc515046c..5095548ee 100644 --- a/src/ray/common/grpc_util.h +++ b/src/ray/common/grpc_util.h @@ -4,6 +4,7 @@ #include #include #include + #include #include "status.h" diff --git a/src/ray/common/task/task_util.h b/src/ray/common/task/task_util.h index bd6170663..31f894029 100644 --- a/src/ray/common/task/task_util.h +++ b/src/ray/common/task/task_util.h @@ -27,7 +27,7 @@ class TaskSpecBuilder { const TaskID &task_id, const Language &language, const std::vector &function_descriptor, const JobID &job_id, const TaskID &parent_task_id, uint64_t parent_counter, const TaskID &caller_id, - uint64_t num_returns, bool is_direct_call, + const rpc::Address &caller_address, uint64_t num_returns, bool is_direct_call, const std::unordered_map &required_resources, const std::unordered_map &required_placement_resources) { message_->set_type(TaskType::NORMAL_TASK); @@ -40,6 +40,7 @@ class TaskSpecBuilder { message_->set_parent_task_id(parent_task_id.Binary()); message_->set_parent_counter(parent_counter); message_->set_caller_id(caller_id.Binary()); + message_->mutable_caller_address()->CopyFrom(caller_address); message_->set_num_returns(num_returns); message_->set_is_direct_call(is_direct_call); message_->mutable_required_resources()->insert(required_resources.begin(), diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index b8fd80fa0..df9007cfc 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -11,15 +11,15 @@ namespace { void BuildCommonTaskSpec( ray::TaskSpecBuilder &builder, const JobID &job_id, const TaskID &task_id, const TaskID ¤t_task_id, const int task_index, const TaskID &caller_id, - const ray::RayFunction &function, const std::vector &args, - uint64_t num_returns, + const ray::rpc::Address &address, const ray::RayFunction &function, + const std::vector &args, uint64_t num_returns, const std::unordered_map &required_resources, const std::unordered_map &required_placement_resources, ray::TaskTransportType transport_type, std::vector *return_ids) { // Build common task spec. builder.SetCommonTaskSpec(task_id, function.GetLanguage(), function.GetFunctionDescriptor(), job_id, current_task_id, - task_index, caller_id, num_returns, + task_index, caller_id, address, num_returns, transport_type == ray::TaskTransportType::DIRECT, required_resources, required_placement_resources); // Set task arguments. @@ -124,16 +124,23 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, // instead of crashing. auto grpc_client = rpc::NodeManagerWorkerClient::make( node_ip_address, node_manager_port, *client_call_manager_); + ClientID raylet_id; raylet_client_ = std::unique_ptr(new RayletClient( std::move(grpc_client), raylet_socket, WorkerID::FromBinary(worker_context_.GetWorkerID().Binary()), (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(), - language_, core_worker_server_.GetPort())); + language_, &raylet_id, core_worker_server_.GetPort())); // Unfortunately the raylet client has to be constructed after the receivers. if (direct_task_receiver_ != nullptr) { direct_task_receiver_->Init(*raylet_client_); } + // Set our own address. + RAY_CHECK(!raylet_id.IsNil()); + rpc_address_.set_ip_address(node_ip_address); + rpc_address_.set_port(core_worker_server_.GetPort()); + rpc_address_.set_raylet_id(raylet_id.Binary()); + // Set timer to periodically send heartbeats containing active object IDs to the raylet. // If the heartbeat timeout is < 0, the heartbeats are disabled. if (RayConfig::instance().worker_heartbeat_timeout_milliseconds() >= 0) { @@ -166,8 +173,8 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, const TaskID task_id = TaskID::ForDriverTask(worker_context_.GetCurrentJobID()); builder.SetCommonTaskSpec( task_id, language_, empty_descriptor, worker_context_.GetCurrentJobID(), - TaskID::ComputeDriverTaskId(worker_context_.GetWorkerID()), 0, GetCallerId(), 0, - false, empty_resources, empty_resources); + TaskID::ComputeDriverTaskId(worker_context_.GetWorkerID()), 0, GetCallerId(), + rpc_address_, 0, false, empty_resources, empty_resources); std::shared_ptr data = std::make_shared(); data->mutable_task()->mutable_task_spec()->CopyFrom(builder.Build().GetMessage()); @@ -523,8 +530,8 @@ Status CoreWorker::SubmitTask(const RayFunction &function, // TODO(ekl) offload task building onto a thread pool for performance BuildCommonTaskSpec( builder, worker_context_.GetCurrentJobID(), task_id, - worker_context_.GetCurrentTaskID(), next_task_index, GetCallerId(), function, args, - task_options.num_returns, task_options.resources, {}, + worker_context_.GetCurrentTaskID(), next_task_index, GetCallerId(), rpc_address_, + function, args, task_options.num_returns, task_options.resources, {}, task_options.is_direct_call ? TaskTransportType::DIRECT : TaskTransportType::RAYLET, return_ids); if (task_options.is_direct_call) { @@ -546,10 +553,11 @@ Status CoreWorker::CreateActor(const RayFunction &function, const JobID job_id = worker_context_.GetCurrentJobID(); std::vector return_ids; TaskSpecBuilder builder; - BuildCommonTaskSpec( - builder, job_id, actor_creation_task_id, worker_context_.GetCurrentTaskID(), - next_task_index, GetCallerId(), function, args, 1, actor_creation_options.resources, - actor_creation_options.placement_resources, TaskTransportType::RAYLET, &return_ids); + BuildCommonTaskSpec(builder, job_id, actor_creation_task_id, + worker_context_.GetCurrentTaskID(), next_task_index, GetCallerId(), + rpc_address_, function, args, 1, actor_creation_options.resources, + actor_creation_options.placement_resources, + TaskTransportType::RAYLET, &return_ids); builder.SetActorCreationTaskSpec(actor_id, actor_creation_options.max_reconstructions, actor_creation_options.dynamic_worker_options, actor_creation_options.is_direct_call, @@ -588,8 +596,8 @@ Status CoreWorker::SubmitActorTask(const ActorID &actor_id, const RayFunction &f next_task_index, actor_handle->GetActorID()); BuildCommonTaskSpec(builder, actor_handle->CreationJobID(), actor_task_id, worker_context_.GetCurrentTaskID(), next_task_index, GetCallerId(), - function, args, num_returns, task_options.resources, {}, - transport_type, return_ids); + rpc_address_, function, args, num_returns, task_options.resources, + {}, transport_type, return_ids); const ObjectID new_cursor = return_ids->back(); actor_handle->SetActorTaskSpec(builder, transport_type, new_cursor); @@ -651,8 +659,8 @@ bool CoreWorker::AddActorHandle(std::unique_ptr actor_handle) { RAY_LOG(INFO) << "received notification on actor, state=" << static_cast(actor_data.state()) << ", actor_id: " << actor_id - << ", ip address: " << actor_data.ip_address() - << ", port: " << actor_data.port(); + << ", ip address: " << actor_data.address().ip_address() + << ", port: " << actor_data.address().port(); }; RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribe( diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 386e8f52d..c79ad3022 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -459,6 +459,9 @@ class CoreWorker { /// RPC server used to receive tasks to execute. rpc::GrpcServer core_worker_server_; + /// Address of our RPC server. + rpc::Address rpc_address_; + // Client to the GCS shared by core worker interfaces. std::shared_ptr gcs_client_; diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index ec0c95653..174ed6ab2 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -499,6 +499,7 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) { int64_t start_ms = current_time_ms(); const auto num_tasks = 10000 * 10; RAY_LOG(INFO) << "start creating " << num_tasks << " PushTaskRequests"; + rpc::Address address; for (int i = 0; i < num_tasks; i++) { TaskOptions options{1, false, resources}; std::vector return_ids; @@ -507,8 +508,8 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) { TaskSpecBuilder builder; builder.SetCommonTaskSpec(RandomTaskId(), function.GetLanguage(), function.GetFunctionDescriptor(), job_id, RandomTaskId(), 0, - RandomTaskId(), num_returns, /*is_direct*/ false, resources, - resources); + RandomTaskId(), address, num_returns, /*is_direct*/ false, + resources, resources); // Set task arguments. for (const auto &arg : args) { if (arg.IsPassedByReference()) { diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index 002db44c7..0c702d963 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -62,15 +62,16 @@ void CoreWorkerDirectActorTaskSubmitter::HandleActorUpdate( std::unique_lock guard(mutex_); actor_states_.erase(actor_id); actor_states_.emplace( - actor_id, - ActorStateData(actor_data.state(), actor_data.ip_address(), actor_data.port())); + actor_id, ActorStateData(actor_data.state(), actor_data.address().ip_address(), + actor_data.address().port())); if (actor_data.state() == ActorTableData::ALIVE) { // Check if this actor is the one that we're interested, if we already have // a connection to the actor, or have pending requests for it, we should // create a new connection. if (pending_requests_.count(actor_id) > 0 && rpc_clients_.count(actor_id) == 0) { - ConnectAndSendPendingTasks(actor_id, actor_data.ip_address(), actor_data.port()); + ConnectAndSendPendingTasks(actor_id, actor_data.address().ip_address(), + actor_data.address().port()); } } else { // Remove rpc client if it's dead or being reconstructed. diff --git a/src/ray/protobuf/common.proto b/src/ray/protobuf/common.proto index c2439dfbc..68daabc9a 100644 --- a/src/ray/protobuf/common.proto +++ b/src/ray/protobuf/common.proto @@ -27,6 +27,12 @@ enum TaskType { ACTOR_TASK = 2; } +message Address { + bytes raylet_id = 1; + string ip_address = 2; + int32 port = 3; +} + /// The task specification encapsulates all immutable information about the /// task. These fields are determined at submission time, converse to the /// `TaskExecutionSpec` may change at execution time. @@ -51,23 +57,25 @@ message TaskSpec { // Task ID of the caller. This is the same as parent_task_id for non-actors. // This is the actor ID (embedded in a nil task ID) for actors. bytes caller_id = 8; + /// Address of the caller. + Address caller_address = 9; // Task arguments. - repeated TaskArg args = 9; + repeated TaskArg args = 10; // Number of return objects. - uint64 num_returns = 10; + uint64 num_returns = 11; // Quantities of the different resources required by this task. - map required_resources = 11; + map required_resources = 12; // The resources required for placing this task on a node. If this is empty, // then the placement resources are equal to the required_resources. - map required_placement_resources = 12; + map required_placement_resources = 13; // Task specification for an actor creation task. // This field is only valid when `type == ACTOR_CREATION_TASK`. - ActorCreationTaskSpec actor_creation_task_spec = 13; + ActorCreationTaskSpec actor_creation_task_spec = 14; // Task specification for an actor task. // This field is only valid when `type == ACTOR_TASK`. - ActorTaskSpec actor_task_spec = 14; + ActorTaskSpec actor_task_spec = 15; // Whether this task is a direct call task. - bool is_direct_call = 15; + bool is_direct_call = 16; } // Argument in the task. diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index 54f35bee1..7bcb84e14 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -97,18 +97,16 @@ message ActorTableData { bytes actor_creation_dummy_object_id = 3; // The ID of the job that created the actor. bytes job_id = 4; - // The ID of the node manager that created the actor. - bytes node_manager_id = 5; // Current state of this actor. ActorState state = 6; // Max number of times this actor should be reconstructed. uint64 max_reconstructions = 7; // Remaining number of reconstructions. uint64 remaining_reconstructions = 8; - // The IP address of the node that the actor is running on. - string ip_address = 9; - // The port that the actor is listening on. - int32 port = 10; + // The address of the the actor. + Address address = 9; + // The address of the the actor's owner (parent). + Address owner_address = 10; // Whether direct actor call is used. bool is_direct_call = 11; // Whether the actor is persistent. diff --git a/src/ray/raylet/actor_registration.cc b/src/ray/raylet/actor_registration.cc index 622af4df7..a78c148b1 100644 --- a/src/ray/raylet/actor_registration.cc +++ b/src/ray/raylet/actor_registration.cc @@ -36,7 +36,7 @@ ActorRegistration::ActorRegistration(const ActorTableData &actor_table_data, } const ClientID ActorRegistration::GetNodeManagerId() const { - return ClientID::FromBinary(actor_table_data_.node_manager_id()); + return ClientID::FromBinary(actor_table_data_.address().raylet_id()); } const ObjectID ActorRegistration::GetActorCreationDependency() const { diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs index 525628802..94bda2f89 100644 --- a/src/ray/raylet/format/node_manager.fbs +++ b/src/ray/raylet/format/node_manager.fbs @@ -142,8 +142,8 @@ table RegisterClientRequest { } table RegisterClientReply { - // The IDs of the GPUs that are reserved for this worker. - gpu_ids: [int]; + // GCS ClientID of the local node manager. + raylet_id: string; } table RegisterNodeManagerRequest { diff --git a/src/ray/raylet/lineage_cache_test.cc b/src/ray/raylet/lineage_cache_test.cc index 688bc1761..7250a7627 100644 --- a/src/ray/raylet/lineage_cache_test.cc +++ b/src/ray/raylet/lineage_cache_test.cc @@ -138,9 +138,10 @@ class LineageCacheTest : public ::testing::Test { static inline Task ExampleTask(const std::vector &arguments, uint64_t num_returns) { TaskSpecBuilder builder; + rpc::Address address; builder.SetCommonTaskSpec(RandomTaskId(), Language::PYTHON, {"", "", ""}, JobID::Nil(), - RandomTaskId(), 0, RandomTaskId(), num_returns, false, {}, - {}); + RandomTaskId(), 0, RandomTaskId(), address, num_returns, + false, {}, {}); for (const auto &arg : arguments) { builder.AddByRefArg(arg); } diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 549f6d97b..dbe6a9e30 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -969,6 +969,20 @@ void NodeManager::ProcessRegisterClientRequestMessage( auto worker = std::make_shared(worker_id, message->worker_pid(), language, message->port(), client, client_call_manager_); Status status; + flatbuffers::FlatBufferBuilder fbb; + auto reply = ray::protocol::CreateRegisterClientReply( + fbb, to_flatbuf(fbb, gcs_client_->client_table().GetLocalClientId())); + fbb.Finish(reply); + client->WriteMessageAsync( + static_cast(protocol::MessageType::RegisterClientReply), fbb.GetSize(), + fbb.GetBufferPointer(), [this, client](const ray::Status &status) { + if (!status.ok()) { + RAY_LOG(WARNING) + << "Failed to send RegisterClientReply to client, so disconnecting"; + ProcessDisconnectClientMessage(client); + } + }); + if (message->is_worker()) { // Register the new worker. if (worker_pool_.RegisterWorker(std::move(worker)).ok()) { @@ -2079,6 +2093,8 @@ std::shared_ptr NodeManager::CreateActorTableDataFromCreationTas actor_info_ptr->set_remaining_reconstructions(task_spec.MaxActorReconstructions()); actor_info_ptr->set_is_direct_call(task_spec.IsDirectCall()); actor_info_ptr->set_is_detached(task_spec.IsDetachedActor()); + actor_info_ptr->mutable_owner_address()->CopyFrom( + task_spec.GetMessage().caller_address()); } else { // If we've already seen this actor, it means that this actor was reconstructed. // Thus, its previous state must be RECONSTRUCTING. @@ -2098,14 +2114,12 @@ std::shared_ptr NodeManager::CreateActorTableDataFromCreationTas actor_info_ptr->remaining_reconstructions() - 1); } - // Set the ip address & port, which could change after reconstruction. - actor_info_ptr->set_ip_address( - gcs_client_->client_table().GetLocalClient().node_manager_address()); - actor_info_ptr->set_port(port); - // Set the new fields for the actor's state to indicate that the actor is // now alive on this node manager. - actor_info_ptr->set_node_manager_id( + actor_info_ptr->mutable_address()->set_ip_address( + gcs_client_->client_table().GetLocalClient().node_manager_address()); + actor_info_ptr->mutable_address()->set_port(port); + actor_info_ptr->mutable_address()->set_raylet_id( gcs_client_->client_table().GetLocalClientId().Binary()); actor_info_ptr->set_state(ActorTableData::ALIVE); return actor_info_ptr; diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index 429910e99..80ef91517 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/src/ray/raylet/raylet_client.cc @@ -204,7 +204,7 @@ ray::Status RayletConnection::AtomicRequestReply( RayletClient::RayletClient(std::shared_ptr grpc_client, const std::string &raylet_socket, const WorkerID &worker_id, bool is_worker, const JobID &job_id, const Language &language, - int port) + ClientID *raylet_id, int port) : grpc_client_(std::move(grpc_client)), worker_id_(worker_id), is_worker_(is_worker), @@ -220,8 +220,13 @@ RayletClient::RayletClient(std::shared_ptr gr fbb.Finish(message); // Register the process ID with the raylet. // NOTE(swang): If raylet exits and we are registered as a worker, we will get killed. - auto status = conn_->WriteMessage(MessageType::RegisterClientRequest, &fbb); + std::unique_ptr reply; + auto status = conn_->AtomicRequestReply(MessageType::RegisterClientRequest, + MessageType::RegisterClientReply, reply, &fbb); RAY_CHECK_OK_PREPEND(status, "[RayletClient] Unable to register worker with raylet."); + auto reply_message = + flatbuffers::GetRoot(reply.get()); + *raylet_id = ClientID::FromBinary(reply_message->raylet_id()->str()); } ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) { diff --git a/src/ray/raylet/raylet_client.h b/src/ray/raylet/raylet_client.h index 38b46d56f..aa36d4c02 100644 --- a/src/ray/raylet/raylet_client.h +++ b/src/ray/raylet/raylet_client.h @@ -87,11 +87,14 @@ class RayletClient : public WorkerLeaseInterface { /// \param is_worker Whether this client is a worker. If it is a worker, an /// additional message will be sent to register as one. /// \param job_id The ID of the driver. This is non-nil if the client is a driver. - /// \return The connection information. + /// \param language Language of the worker. + /// \param raylet_id This will be populated with the local raylet's ClientID. + /// \param port The port that the worker will listen on for gRPC requests, if + /// any. RayletClient(std::shared_ptr grpc_client, const std::string &raylet_socket, const WorkerID &worker_id, bool is_worker, const JobID &job_id, const Language &language, - int port = -1); + ClientID *raylet_id, int port = -1); ray::Status Disconnect() { return conn_->Disconnect(); }; diff --git a/src/ray/raylet/task_dependency_manager_test.cc b/src/ray/raylet/task_dependency_manager_test.cc index 2ae727434..5431f5d0a 100644 --- a/src/ray/raylet/task_dependency_manager_test.cc +++ b/src/ray/raylet/task_dependency_manager_test.cc @@ -75,9 +75,10 @@ class TaskDependencyManagerTest : public ::testing::Test { static inline Task ExampleTask(const std::vector &arguments, uint64_t num_returns) { TaskSpecBuilder builder; + rpc::Address address; builder.SetCommonTaskSpec(RandomTaskId(), Language::PYTHON, {"", "", ""}, JobID::Nil(), - RandomTaskId(), 0, RandomTaskId(), num_returns, false, {}, - {}); + RandomTaskId(), 0, RandomTaskId(), address, num_returns, + false, {}, {}); for (const auto &arg : arguments) { builder.AddByRefArg(arg); } diff --git a/src/ray/rpc/worker/core_worker_client.h b/src/ray/rpc/worker/core_worker_client.h index 4cd74af1f..763f76f9d 100644 --- a/src/ray/rpc/worker/core_worker_client.h +++ b/src/ray/rpc/worker/core_worker_client.h @@ -84,6 +84,8 @@ class CoreWorkerClientInterface { const ClientCallback &callback) { return Status::NotImplemented(""); } + + virtual ~CoreWorkerClientInterface(){}; }; /// Client used for communicating with a remote worker server.