[core] Fix GCS FD usage increase regression.

BUILD.bazel

@@ -100,6 +100,9 @@ cc_grpc_library(
 # gcs rpc server and client.
 cc_library(
     name = "gcs_service_rpc",
+    srcs = [
+        "src/ray/rpc/gcs_server/gcs_rpc_client.cc",
+    ],
     hdrs = [
         "src/ray/rpc/gcs_server/gcs_rpc_client.h",
         "src/ray/rpc/gcs_server/gcs_rpc_server.h",

python/ray/tests/test_advanced_9.py

@@ -16,6 +16,7 @@
 from ray.experimental.internal_kv import _internal_kv_list
 from ray.tests.conftest import call_ray_start
 import subprocess
+import psutil
 
 
 @pytest.fixture
@@ -269,8 +270,6 @@ def test_gcs_connection_no_leak(ray_start_cluster):
     ray.init(cluster.address)
 
     def get_gcs_num_of_connections():
-        import psutil
-
         p = psutil.Process(gcs_server_pid)
         print(">>", p.num_fds())
         return p.num_fds()
@@ -438,6 +437,42 @@ def f():
         assert ray.get(f.remote())
 
 
+def test_gcs_fd_usage(shutdown_only):
+    ray.init(
+        _system_config={
+            "prestart_worker_first_driver": False,
+            "enable_worker_prestart": False,
+        },
+    )
+    gcs_process = ray._private.worker._global_node.all_processes["gcs_server"][0]
+    gcs_process = psutil.Process(gcs_process.process.pid)
+    print("GCS connections", len(gcs_process.connections()))
+
+    @ray.remote(runtime_env={"env_vars": {"Hello": "World"}})
+    class A:
+        def f(self):
+            import os
+
+            return os.environ.get("Hello")
+
+    # In case there are still some pre-start workers, consume all of them
+    aa = [A.remote() for _ in range(32)]
+    for a in aa:
+        assert ray.get(a.f.remote()) == "World"
+    base_fd_num = len(gcs_process.connections())
+    print("GCS connections", base_fd_num)
+
+    bb = [A.remote() for _ in range(4)]
+    for b in bb:
+        assert ray.get(b.f.remote()) == "World"
+    new_fd_num = len(gcs_process.connections())
+    print("GCS connections", new_fd_num)
+    # each worker has two connections:
+    #   GCS -> CoreWorker
+    #   CoreWorker -> GCS
+    assert (new_fd_num - base_fd_num) == len(bb) * 2
+
+
 if __name__ == "__main__":
     import pytest
     import os

src/ray/common/ray_config_def.h

@@ -215,7 +215,7 @@ RAY_CONFIG(int64_t, max_direct_call_object_size, 100 * 1024)
 // The max gRPC message size (the gRPC internal default is 4MB). We use a higher
 // limit in Ray to avoid crashing with many small inlined task arguments.
 // Keep in sync with GCS_STORAGE_MAX_SIZE in packaging.py.
-RAY_CONFIG(int64_t, max_grpc_message_size, 500 * 1024 * 1024)
+RAY_CONFIG(int64_t, max_grpc_message_size, 512 * 1024 * 1024)
 
 // Retry timeout for trying to create a gRPC server. Only applies if the number
 // of retries is non zero.

src/ray/gcs/gcs_client/gcs_client.cc

@@ -146,8 +146,8 @@ std::pair<std::string, int> GcsClient::GetGcsServerAddress() const {
 PythonGcsClient::PythonGcsClient(const GcsClientOptions &options) : options_(options) {}
 
 Status PythonGcsClient::Connect() {
-  auto arguments = PythonGrpcChannelArguments();
-  channel_ = rpc::BuildChannel(options_.gcs_address_, options_.gcs_port_, arguments);
+  channel_ =
+      rpc::GcsRpcClient::GetDefaultChannel(options_.gcs_address_, options_.gcs_port_);
   kv_stub_ = rpc::InternalKVGcsService::NewStub(channel_);
   runtime_env_stub_ = rpc::RuntimeEnvGcsService::NewStub(channel_);
   node_info_stub_ = rpc::NodeInfoGcsService::NewStub(channel_);

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -53,8 +53,8 @@ class GcsServerTest : public ::testing::Test {
 
     // Create gcs rpc client
     client_call_manager_.reset(new rpc::ClientCallManager(io_service_));
-    client_.reset(
-        new rpc::GcsRpcClient("0.0.0.0", gcs_server_->GetPort(), *client_call_manager_));
+    client_.reset(new rpc::GcsRpcClient(
+        "127.0.0.1", gcs_server_->GetPort(), *client_call_manager_));
   }
 
   void TearDown() override {

src/ray/gcs/pubsub/gcs_pub_sub.cc

@@ -15,6 +15,7 @@
 #include "ray/gcs/pubsub/gcs_pub_sub.h"
 
 #include "absl/strings/str_cat.h"
+#include "ray/rpc/gcs_server/gcs_rpc_client.h"
 #include "ray/rpc/grpc_client.h"
 
 namespace ray {
@@ -230,8 +231,7 @@ PythonGcsPublisher::PythonGcsPublisher(const std::string &gcs_address) {
 }
 
 Status PythonGcsPublisher::Connect() {
-  auto arguments = PythonGrpcChannelArguments();
-  channel_ = rpc::BuildChannel(gcs_address_, gcs_port_, arguments);
+  channel_ = rpc::GcsRpcClient::GetDefaultChannel(gcs_address_, gcs_port_);
   pubsub_stub_ = rpc::InternalPubSubGcsService::NewStub(channel_);
   return Status::OK();
 }

src/ray/rpc/gcs_server/gcs_rpc_client.cc

@@ -0,0 +1,63 @@
+// Copyright 2023 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/rpc/gcs_server/gcs_rpc_client.h"
+
+namespace ray {
+namespace rpc {
+grpc::ChannelArguments GetGcsRpcClientArguments() {
+  grpc::ChannelArguments arguments = CreateDefaultChannelArguments();
+  arguments.SetInt(GRPC_ARG_MAX_RECONNECT_BACKOFF_MS,
+                   ::RayConfig::instance().gcs_grpc_max_reconnect_backoff_ms());
+  arguments.SetInt(GRPC_ARG_MIN_RECONNECT_BACKOFF_MS,
+                   ::RayConfig::instance().gcs_grpc_min_reconnect_backoff_ms());
+  arguments.SetInt(GRPC_ARG_INITIAL_RECONNECT_BACKOFF_MS,
+                   ::RayConfig::instance().gcs_grpc_initial_reconnect_backoff_ms());
+  return arguments;
+}
+
+std::shared_ptr<grpc::Channel> GcsRpcClient::GetDefaultChannel(const std::string &address,
+                                                               int port) {
+  static std::shared_ptr<grpc::Channel> channel_;
+  static std::mutex mu_;
+
+  static std::string address_;
+  static int port_;
+  std::lock_guard<std::mutex> guard(mu_);
+  // Don't reuse channel if proxy or tls is set
+  // TODO: Reuse the channel even it's tls.
+  // Right now, if we do this, python/ray/serve/tests/test_grpc.py
+  // will fail.
+  if (::RayConfig::instance().grpc_enable_http_proxy() ||
+      ::RayConfig::instance().USE_TLS()) {
+    return BuildChannel(address, port, GetGcsRpcClientArguments());
+  }
+
+  if (channel_ == nullptr) {
+    address_ = address;
+    port_ = port;
+    channel_ = BuildChannel(address, port, GetGcsRpcClientArguments());
+  }
+
+  if (address_ == address && port_ == port) {
+    return channel_;
+  } else {
+    RAY_LOG(WARNING) << "Generate a new GCS channel: " << address << ":" << port
+                     << ". Potentially it will increase GCS socket numbers";
+    return BuildChannel(address, port, GetGcsRpcClientArguments());
+  }
+}
+
+}  // namespace rpc
+}  // namespace ray

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -168,6 +168,10 @@ class Executor {
 
 /// Client used for communicating with gcs server.
 class GcsRpcClient {
+ public:
+  static std::shared_ptr<grpc::Channel> GetDefaultChannel(const std::string &address,
+                                                          int port);
+
  public:
   /// Constructor. GcsRpcClient is not thread safe.
   ///
@@ -185,16 +189,7 @@ class GcsRpcClient {
         gcs_port_(port),
         io_context_(&client_call_manager.GetMainService()),
         timer_(std::make_unique<boost::asio::deadline_timer>(*io_context_)) {
-    grpc::ChannelArguments arguments = CreateDefaultChannelArguments();
-    arguments.SetInt(GRPC_ARG_MAX_RECONNECT_BACKOFF_MS,
-                     ::RayConfig::instance().gcs_grpc_max_reconnect_backoff_ms());
-    arguments.SetInt(GRPC_ARG_MIN_RECONNECT_BACKOFF_MS,
-                     ::RayConfig::instance().gcs_grpc_min_reconnect_backoff_ms());
-    arguments.SetInt(GRPC_ARG_INITIAL_RECONNECT_BACKOFF_MS,
-                     ::RayConfig::instance().gcs_grpc_initial_reconnect_backoff_ms());
-
-    channel_ = BuildChannel(address, port, arguments);
-
+    channel_ = GetDefaultChannel(address, port);
     // If not the reconnection will continue to work.
     auto deadline =
         std::chrono::system_clock::now() +