[Core] Added ability to specify different IP addresses for a core worker and its raylet. (#7985)

src/ray/core_worker/core_worker.cc

@@ -313,13 +313,13 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
   // so that the worker (java/python .etc) can retrieve and handle the error
   // instead of crashing.
   auto grpc_client = rpc::NodeManagerWorkerClient::make(
-      options_.node_ip_address, options_.node_manager_port, *client_call_manager_);
+      options_.raylet_ip_address, options_.node_manager_port, *client_call_manager_);
   ClientID local_raylet_id;
   local_raylet_client_ = std::shared_ptr<raylet::RayletClient>(new raylet::RayletClient(
       io_service_, std::move(grpc_client), options_.raylet_socket, GetWorkerID(),
       (options_.worker_type == ray::WorkerType::WORKER),
       worker_context_.GetCurrentJobID(), options_.language, &local_raylet_id,
-      core_worker_server_.GetPort()));
+      options_.node_ip_address, core_worker_server_.GetPort()));
   connected_ = true;
 
   // Set our own address.

src/ray/core_worker/core_worker.h

@@ -84,6 +84,8 @@ struct CoreWorkerOptions {
   std::string node_ip_address;
   /// Port of the local raylet.
   int node_manager_port;
+  /// IP address of the raylet.
+  std::string raylet_ip_address;
   /// The name of the driver.
   std::string driver_name;
   /// The stdout file of this process.

src/ray/core_worker/lib/java/io_ray_runtime_RayNativeRuntime.cc

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 #include "ray/core_worker/lib/java/io_ray_runtime_RayNativeRuntime.h"
+
 #include <jni.h>
+
 #include <sstream>
+
 #include "ray/common/id.h"
 #include "ray/core_worker/core_worker.h"
 #include "ray/core_worker/lib/java/jni_utils.h"
@@ -37,7 +40,6 @@ inline ray::gcs::GcsClientOptions ToGcsClientOptions(JNIEnv *env,
 extern "C" {
 #endif
 
-
 JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
     JNIEnv *env, jclass, jint workerMode, jstring nodeIpAddress, jint nodeManagerPort,
     jstring driverName, jstring storeSocket, jstring rayletSocket, jbyteArray jobId,
@@ -112,6 +114,7 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
       false,                                         // install_failure_signal_handler
       JavaStringToNativeString(env, nodeIpAddress),  // node_ip_address
       static_cast<int>(nodeManagerPort),             // node_manager_port
+      JavaStringToNativeString(env, nodeIpAddress),  // raylet_ip_address
       JavaStringToNativeString(env, driverName),     // driver_name
       "",                                            // stdout_file
       "",                                            // stderr_file
@@ -135,7 +138,7 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeRunTaskExecuto
 }
 
 JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeShutdown(JNIEnv *env,
-                                                                            jclass o) {
+                                                                           jclass o) {
   ray::CoreWorkerProcess::Shutdown();
 }
 

src/ray/core_worker/test/core_worker_test.cc

@@ -263,6 +263,7 @@ class CoreWorkerTest : public ::testing::Test {
           true,                           // install_failure_signal_handler
           "127.0.0.1",                    // node_ip_address
           node_manager_port,              // node_manager_port
+          "127.0.0.1",                    // raylet_ip_address
           "core_worker_test",             // driver_name
           "",                             // stdout_file
           "",                             // stderr_file

src/ray/core_worker/test/mock_worker.cc

@@ -45,6 +45,7 @@ class MockWorker {
         true,                // install_failure_signal_handler
         "127.0.0.1",         // node_ip_address
         node_manager_port,   // node_manager_port
+        "127.0.0.1",         // raylet_ip_address
         "",                  // driver_name
         "",                  // stdout_file
         "",                  // stderr_file

src/ray/gcs/pb_util.h

@@ -16,6 +16,7 @@
 #define RAY_GCS_PB_UTIL_H
 
 #include <memory>
+
 #include "ray/common/id.h"
 #include "ray/common/task/task_spec.h"
 #include "ray/protobuf/gcs.pb.h"
@@ -29,17 +30,17 @@ namespace gcs {
 /// \param job_id The ID of job that need to be registered or updated.
 /// \param is_dead Whether the driver of this job is dead.
 /// \param timestamp The UNIX timestamp of corresponding to this event.
-/// \param node_manager_address Address of the node this job was started on.
+/// \param driver_ip_address IP address of the driver that started this job.
 /// \param driver_pid Process ID of the driver running this job.
 /// \return The job table data created by this method.
 inline std::shared_ptr<ray::rpc::JobTableData> CreateJobTableData(
     const ray::JobID &job_id, bool is_dead, int64_t timestamp,
-    const std::string &node_manager_address, int64_t driver_pid) {
+    const std::string &driver_ip_address, int64_t driver_pid) {
   auto job_info_ptr = std::make_shared<ray::rpc::JobTableData>();
   job_info_ptr->set_job_id(job_id.Binary());
   job_info_ptr->set_is_dead(is_dead);
   job_info_ptr->set_timestamp(timestamp);
-  job_info_ptr->set_node_manager_address(node_manager_address);
+  job_info_ptr->set_driver_ip_address(driver_ip_address);
   job_info_ptr->set_driver_pid(driver_pid);
   return job_info_ptr;
 }

src/ray/gcs/redis_accessor.cc

@@ -13,7 +13,9 @@
 // limitations under the License.
 
 #include "ray/gcs/redis_accessor.h"
+
 #include <boost/none.hpp>
+
 #include "ray/gcs/pb_util.h"
 #include "ray/gcs/redis_gcs_client.h"
 #include "ray/util/logging.h"
@@ -304,7 +306,7 @@ Status RedisJobInfoAccessor::AsyncMarkFinished(const JobID &job_id,
                                                const StatusCallback &callback) {
   std::shared_ptr<JobTableData> data_ptr =
       CreateJobTableData(job_id, /*is_dead*/ true, /*time_stamp*/ std::time(nullptr),
-                         /*node_manager_address*/ "", /*driver_pid*/ -1);
+                         /*driver_ip_address*/ "", /*driver_pid*/ -1);
   return DoAsyncAppend(data_ptr, callback);
 }
 

src/ray/gcs/test/gcs_test_util.h

@@ -15,6 +15,9 @@
 #ifndef RAY_GCS_TEST_UTIL_H
 #define RAY_GCS_TEST_UTIL_H
 
+#include <memory>
+#include <utility>
+
 #include "src/ray/common/task/task.h"
 #include "src/ray/common/task/task_util.h"
 #include "src/ray/common/test_util.h"
@@ -23,9 +26,6 @@
 #include "src/ray/gcs/gcs_server/gcs_node_manager.h"
 #include "src/ray/util/asio_util.h"
 
-#include <memory>
-#include <utility>
-
 namespace ray {
 
 struct Mocker {
@@ -64,7 +64,7 @@ struct Mocker {
     job_table_data->set_job_id(job_id.Binary());
     job_table_data->set_is_dead(false);
     job_table_data->set_timestamp(std::time(nullptr));
-    job_table_data->set_node_manager_address("127.0.0.1");
+    job_table_data->set_driver_ip_address("127.0.0.1");
     job_table_data->set_driver_pid(5667L);
     return job_table_data;
   }

src/ray/gcs/test/redis_job_info_accessor_test.cc

@@ -31,7 +31,7 @@ class RedisJobInfoAccessorTest : public AccessorTestBase<JobID, JobTableData> {
       JobID job_id = JobID::FromInt(i);
       std::shared_ptr<JobTableData> job_data_ptr =
           CreateJobTableData(job_id, /*is_dead*/ false, /*timestamp*/ 1,
-                             /*node_manager_address*/ "", /*driver_pid*/ i);
+                             /*driver_ip_address*/ "", /*driver_pid*/ i);
       id_to_data_[job_id] = job_data_ptr;
     }
   }

src/ray/protobuf/gcs.proto

@@ -256,8 +256,8 @@ message JobTableData {
   bool is_dead = 2;
   // The UNIX timestamp corresponding to this event (job added or removed).
   int64 timestamp = 3;
-  // IP of the node this job was started on.
-  string node_manager_address = 4;
+  // IP address of the driver that started this job.
+  string driver_ip_address = 4;
   // Process ID of the driver running this job.
   int64 driver_pid = 5;
 }

src/ray/raylet/format/node_manager.fbs

@@ -151,6 +151,8 @@ table RegisterClientRequest {
   // Language of this worker.
   // TODO(hchen): Use `Language` in `common.proto`.
   language: int;
+  // IP address of this worker.
+  ip_address: string;
   // Port that this worker is listening on.
   port: int;
 }

src/ray/raylet/node_manager.cc

@@ -1098,8 +1098,9 @@ void NodeManager::ProcessRegisterClientRequestMessage(
   Language language = static_cast<Language>(message->language());
   WorkerID worker_id = from_flatbuf<WorkerID>(*message->worker_id());
   pid_t pid = message->worker_pid();
-  auto worker = std::make_shared<Worker>(worker_id, language, message->port(), client,
-                                         client_call_manager_);
+  std::string worker_ip_address = string_from_flatbuf(*message->ip_address());
+  auto worker = std::make_shared<Worker>(worker_id, language, worker_ip_address,
+                                         message->port(), client, client_call_manager_);
   if (message->is_worker()) {
     // Register the new worker.
     if (worker_pool_.RegisterWorker(worker, pid).ok()) {
@@ -1117,9 +1118,8 @@ void NodeManager::ProcessRegisterClientRequestMessage(
     Status status = worker_pool_.RegisterDriver(worker);
     if (status.ok()) {
       local_queues_.AddDriverTaskId(driver_task_id);
-      auto job_data_ptr =
-          gcs::CreateJobTableData(job_id, /*is_dead*/ false, std::time(nullptr),
-                                  initial_config_.node_manager_address, pid);
+      auto job_data_ptr = gcs::CreateJobTableData(
+          job_id, /*is_dead*/ false, std::time(nullptr), worker_ip_address, pid);
       RAY_CHECK_OK(gcs_client_->Jobs().AsyncAdd(job_data_ptr, nullptr));
     }
   }
@@ -1260,8 +1260,8 @@ void NodeManager::ProcessDisconnectClientMessage(
 
     // Publish the worker failure.
     auto worker_failure_data_ptr = gcs::CreateWorkerFailureData(
-        self_node_id_, worker->WorkerId(), initial_config_.node_manager_address,
-        worker->Port(), time(nullptr), intentional_disconnect);
+        self_node_id_, worker->WorkerId(), worker->IpAddress(), worker->Port(),
+        time(nullptr), intentional_disconnect);
     RAY_CHECK_OK(gcs_client_->Workers().AsyncReportWorkerFailure(worker_failure_data_ptr,
                                                                  nullptr));
   }
@@ -1687,8 +1687,7 @@ void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest
                                                       ClientID spillback_to,
                                                       std::string address, int port) {
           if (worker != nullptr) {
-            reply->mutable_worker_address()->set_ip_address(
-                initial_config_.node_manager_address);
+            reply->mutable_worker_address()->set_ip_address(worker->IpAddress());
             reply->mutable_worker_address()->set_port(worker->Port());
             reply->mutable_worker_address()->set_worker_id(worker->WorkerId().Binary());
             reply->mutable_worker_address()->set_raylet_id(self_node_id_.Binary());

src/ray/raylet/raylet_client.cc

@@ -165,7 +165,8 @@ raylet::RayletClient::RayletClient(
     boost::asio::io_service &io_service,
     std::shared_ptr<rpc::NodeManagerWorkerClient> grpc_client,
     const std::string &raylet_socket, const WorkerID &worker_id, bool is_worker,
-    const JobID &job_id, const Language &language, ClientID *raylet_id, int port)
+    const JobID &job_id, const Language &language, ClientID *raylet_id,
+    const std::string &ip_address, int port)
     : grpc_client_(std::move(grpc_client)), worker_id_(worker_id), job_id_(job_id) {
   // For C++14, we could use std::make_unique
   conn_ = std::unique_ptr<raylet::RayletConnection>(
@@ -174,7 +175,7 @@ raylet::RayletClient::RayletClient(
   flatbuffers::FlatBufferBuilder fbb;
   auto message = protocol::CreateRegisterClientRequest(
       fbb, is_worker, to_flatbuf(fbb, worker_id), getpid(), to_flatbuf(fbb, job_id),
-      language, port);
+      language, fbb.CreateString(ip_address), port);
   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.

src/ray/raylet/raylet_client.h

@@ -154,13 +154,14 @@ class RayletClient : public PinObjectsInterface,
   /// \param job_id The ID of the driver. This is non-nil if the client is a driver.
   /// \param language Language of the worker.
   /// \param raylet_id This will be populated with the local raylet's ClientID.
+  /// \param ip_address The IP address of the worker.
   /// \param port The port that the worker will listen on for gRPC requests, if
   /// any.
   RayletClient(boost::asio::io_service &io_service,
                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,
-               ClientID *raylet_id, int port = -1);
+               ClientID *raylet_id, const std::string &ip_address, int port = -1);
 
   /// Connect to the raylet via grpc only.
   ///

src/ray/raylet/worker.cc

@@ -26,11 +26,13 @@ namespace ray {
 namespace raylet {
 
 /// A constructor responsible for initializing the state of a worker.
-Worker::Worker(const WorkerID &worker_id, const Language &language, int port,
+Worker::Worker(const WorkerID &worker_id, const Language &language,
+               const std::string &ip_address, int port,
                std::shared_ptr<ClientConnection> connection,
                rpc::ClientCallManager &client_call_manager)
     : worker_id_(worker_id),
       language_(language),
+      ip_address_(ip_address),
       port_(port),
       connection_(connection),
       dead_(false),
@@ -39,7 +41,7 @@ Worker::Worker(const WorkerID &worker_id, const Language &language, int port,
       is_detached_actor_(false) {
   if (port_ > 0) {
     rpc::Address addr;
-    addr.set_ip_address("127.0.0.1");
+    addr.set_ip_address(ip_address_);
     addr.set_port(port_);
     rpc_client_ = std::unique_ptr<rpc::CoreWorkerClient>(
         new rpc::CoreWorkerClient(addr, client_call_manager_));
@@ -67,6 +69,8 @@ void Worker::SetProcess(Process proc) {
 
 Language Worker::GetLanguage() const { return language_; }
 
+const std::string Worker::IpAddress() const { return ip_address_; }
+
 int Worker::Port() const { return port_; }
 
 void Worker::AssignTaskId(const TaskID &task_id) { assigned_task_id_ = task_id; }

src/ray/raylet/worker.h

@@ -38,7 +38,8 @@ class Worker {
  public:
   /// A constructor that initializes a worker object.
   /// NOTE: You MUST manually set the worker process.
-  Worker(const WorkerID &worker_id, const Language &language, int port,
+  Worker(const WorkerID &worker_id, const Language &language,
+         const std::string &ip_address, int port,
          std::shared_ptr<ClientConnection> connection,
          rpc::ClientCallManager &client_call_manager);
   /// A destructor responsible for freeing all worker state.
@@ -54,6 +55,7 @@ class Worker {
   Process GetProcess() const;
   void SetProcess(Process proc);
   Language GetLanguage() const;
+  const std::string IpAddress() const;
   int Port() const;
   void AssignTaskId(const TaskID &task_id);
   const TaskID &GetAssignedTaskId() const;
@@ -131,6 +133,8 @@ class Worker {
   Process proc_;
   /// The language type of this worker.
   Language language_;
+  /// IP address of this worker.
+  std::string ip_address_;
   /// Port that this worker listens on.
   /// If port <= 0, this indicates that the worker will not listen to a port.
   int port_;

src/ray/raylet/worker_pool_test.cc

@@ -115,7 +115,7 @@ class WorkerPoolTest : public ::testing::Test {
         ClientConnection::Create(client_handler, message_handler, std::move(socket),
                                  "worker", {}, error_message_type_);
     std::shared_ptr<Worker> worker = std::make_shared<Worker>(
-        WorkerID::FromRandom(), language, -1, client, client_call_manager_);
+        WorkerID::FromRandom(), language, "127.0.0.1", -1, client, client_call_manager_);
     if (!proc.IsNull()) {
       worker->SetProcess(proc);
     }