Pass through caller address when submitting a task (#6143)

* Add RpcAddress, set in actor table data

* Pass through task caller address

* RpcAddress -> Address

* update

* fix

* lint

* fix cc tests

src/ray/common/grpc_util.h

@@ -4,6 +4,7 @@
 #include <google/protobuf/map.h>
 #include <google/protobuf/repeated_field.h>
 #include <grpcpp/grpcpp.h>
+
 #include <sstream>
 #include "status.h"
 

src/ray/common/task/task_util.h

@@ -27,7 +27,7 @@ class TaskSpecBuilder {
       const TaskID &task_id, const Language &language,
       const std::vector<std::string> &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<std::string, double> &required_resources,
       const std::unordered_map<std::string, double> &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(),

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 &current_task_id, const int task_index, const TaskID &caller_id,
-    const ray::RayFunction &function, const std::vector<ray::TaskArg> &args,
-    uint64_t num_returns,
+    const ray::rpc::Address &address, const ray::RayFunction &function,
+    const std::vector<ray::TaskArg> &args, uint64_t num_returns,
     const std::unordered_map<std::string, double> &required_resources,
     const std::unordered_map<std::string, double> &required_placement_resources,
     ray::TaskTransportType transport_type, std::vector<ObjectID> *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<RayletClient>(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<gcs::TaskTableData> data = std::make_shared<gcs::TaskTableData>();
     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<ObjectID> 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<ActorHandle> actor_handle) {
 
       RAY_LOG(INFO) << "received notification on actor, state="
                     << static_cast<int>(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(

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::RedisGcsClient> gcs_client_;
 

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<ObjectID> 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()) {

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -62,15 +62,16 @@ void CoreWorkerDirectActorTaskSubmitter::HandleActorUpdate(
   std::unique_lock<std::mutex> 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.

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<string, double> required_resources = 11;
+  map<string, double> 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<string, double> required_placement_resources = 12;
+  map<string, double> 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.

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.

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 {

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 {

src/ray/raylet/lineage_cache_test.cc

@@ -138,9 +138,10 @@ class LineageCacheTest : public ::testing::Test {
 static inline Task ExampleTask(const std::vector<ObjectID> &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);
   }

src/ray/raylet/node_manager.cc

@@ -969,6 +969,20 @@ void NodeManager::ProcessRegisterClientRequestMessage(
   auto worker = std::make_shared<Worker>(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<int64_t>(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<ActorTableData> 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<ActorTableData> 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;

src/ray/raylet/raylet_client.cc

@@ -204,7 +204,7 @@ ray::Status RayletConnection::AtomicRequestReply(
 RayletClient::RayletClient(std::shared_ptr<ray::rpc::NodeManagerWorkerClient> 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<ray::rpc::NodeManagerWorkerClient> 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<uint8_t[]> 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<ray::protocol::RegisterClientReply>(reply.get());
+  *raylet_id = ClientID::FromBinary(reply_message->raylet_id()->str());
 }
 
 ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) {

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<ray::rpc::NodeManagerWorkerClient> 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(); };
 

src/ray/raylet/task_dependency_manager_test.cc

@@ -75,9 +75,10 @@ class TaskDependencyManagerTest : public ::testing::Test {
 static inline Task ExampleTask(const std::vector<ObjectID> &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);
   }

src/ray/rpc/worker/core_worker_client.h

@@ -84,6 +84,8 @@ class CoreWorkerClientInterface {
       const ClientCallback<WorkerLeaseGrantedReply> &callback) {
     return Status::NotImplemented("");
   }
+
+  virtual ~CoreWorkerClientInterface(){};
 };
 
 /// Client used for communicating with a remote worker server.