Move the pending resource load pull off the GCS main io_context - #65024

Signed-off-by: yicheng <yicheng@anyscale.com>

Author: yicheng

src/ray/gcs/BUILD.bazel

@@ -91,6 +91,20 @@ ray_cc_library(
     ],
 )
 
+ray_cc_library(
+    name = "gcs_resource_load_puller",
+    srcs = ["gcs_resource_load_puller.cc"],
+    hdrs = ["gcs_resource_load_puller.h"],
+    deps = [
+        "//src/ray/asio:instrumented_io_context",
+        "//src/ray/common:id",
+        "//src/ray/protobuf:gcs_cc_proto",
+        "//src/ray/raylet_rpc_client:raylet_client_pool",
+        "//src/ray/util:logging",
+        "@com_google_absl//absl/container:flat_hash_set",
+    ],
+)
+
 ray_cc_library(
     name = "gcs_resource_manager",
     srcs = ["gcs_resource_manager.cc"],
@@ -238,6 +252,7 @@ ray_cc_library(
     name = "gcs_server_io_context_policy",
     hdrs = ["gcs_server_io_context_policy.h"],
     deps = [
+        ":gcs_resource_load_puller",
         ":gcs_task_manager",
         "//src/ray/observability:ray_event_recorder",
         "//src/ray/pubsub:gcs_publisher",
@@ -427,6 +442,7 @@ ray_cc_library(
         ":gcs_placement_group_manager",
         ":gcs_placement_group_scheduler",
         ":gcs_pubsub_handler",
+        ":gcs_resource_load_puller",
         ":gcs_resource_manager",
         ":gcs_runtime_env_handler",
         ":gcs_server_io_context_policy",

src/ray/gcs/gcs_resource_load_puller.cc

@@ -0,0 +1,68 @@
+// Copyright 2026 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/gcs/gcs_resource_load_puller.h"
+
+#include <utility>
+#include <vector>
+
+#include "ray/util/logging.h"
+
+namespace ray {
+namespace gcs {
+
+GcsResourceLoadPuller::GcsResourceLoadPuller(
+    instrumented_io_context &pull_io_context,
+    instrumented_io_context &main_io_context,
+    rpc::RayletClientPool &raylet_client_pool,
+    std::function<void(rpc::ResourcesData)> apply_on_main)
+    : pull_io_context_(pull_io_context),
+      main_io_context_(main_io_context),
+      raylet_client_pool_(raylet_client_pool),
+      apply_on_main_(std::move(apply_on_main)) {}
+
+void GcsResourceLoadPuller::Pull(std::vector<rpc::Address> raylet_addresses) {
+  RAY_CHECK(pull_io_context_.get_executor().running_in_this_thread());
+  absl::flat_hash_set<NodeID> current_node_ids;
+  current_node_ids.reserve(raylet_addresses.size());
+  for (const auto &address : raylet_addresses) {
+    current_node_ids.insert(NodeID::FromBinary(address.node_id()));
+  }
+  for (const auto &node_id : pulled_node_ids_) {
+    if (!current_node_ids.contains(node_id)) {
+      raylet_client_pool_.Disconnect(node_id);
+    }
+  }
+  pulled_node_ids_ = std::move(current_node_ids);
+
+  for (const auto &address : raylet_addresses) {
+    auto raylet_client = raylet_client_pool_.GetOrConnectByAddress(address);
+    raylet_client->GetResourceLoad(
+        [this](const Status &status, rpc::GetResourceLoadReply &&reply) {
+          if (!status.ok()) {
+            RAY_LOG_EVERY_N(WARNING, 10)
+                << "Failed to get the resource load: " << status.ToString();
+            return;
+          }
+          main_io_context_.post(
+              [this, resources = std::move(*reply.mutable_resources())]() mutable {
+                apply_on_main_(std::move(resources));
+              },
+              "GcsResourceLoadPuller.apply");
+        });
+  }
+}
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/gcs_resource_load_puller.h

@@ -0,0 +1,52 @@
+// Copyright 2026 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.
+
+#pragma once
+
+#include <functional>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "ray/asio/instrumented_io_context.h"
+#include "ray/common/id.h"
+#include "ray/raylet_rpc_client/raylet_client_pool.h"
+#include "src/ray/protobuf/gcs.pb.h"
+
+namespace ray {
+namespace gcs {
+
+/// Pulls every raylet's pending resource requests (autoscaler bookkeeping) on
+/// a dedicated io_context, and posts each reply back to the main io_context,
+/// where the consumers live.
+class GcsResourceLoadPuller {
+ public:
+  GcsResourceLoadPuller(instrumented_io_context &pull_io_context,
+                        instrumented_io_context &main_io_context,
+                        rpc::RayletClientPool &raylet_client_pool,
+                        std::function<void(rpc::ResourcesData)> apply_on_main);
+
+  void Pull(std::vector<rpc::Address> raylet_addresses);
+
+ private:
+  instrumented_io_context &pull_io_context_;
+  instrumented_io_context &main_io_context_;
+  rpc::RayletClientPool &raylet_client_pool_;
+  std::function<void(rpc::ResourcesData)> apply_on_main_;
+  /// Node ids from the last Pull(), diffed against the next snapshot to remove
+  /// dead raylets' pooled clients.
+  absl::flat_hash_set<NodeID> pulled_node_ids_;
+};
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/gcs_server.cc

@@ -131,6 +131,20 @@ GcsServer::GcsServer(const ray::gcs::GcsServerConfig &config,
                   });
             });
       }),
+      resource_load_pull_client_call_manager_(
+          io_context_provider_.GetIOContext<GcsResourceLoadPuller>(),
+          /*record_stats=*/true,
+          config.node_ip_address,
+          ClusterID::Nil(),
+          /*num_threads=*/1),
+      resource_load_pull_raylet_client_pool_([this](const rpc::Address &addr) {
+        // GetResourceLoad is not retryable, so the unavailable-timeout callback
+        // can never fire; the puller's snapshot diff evicts instead.
+        return std::make_shared<ray::rpc::RayletClient>(
+            addr,
+            this->resource_load_pull_client_call_manager_,
+            /*raylet_unavailable_timeout_callback=*/[]() {});
+      }),
       event_aggregator_client_call_manager_(
           io_context_provider_.GetIOContext<observability::RayEventRecorder>(),
           /*record_stats=*/true,
@@ -154,6 +168,8 @@ GcsServer::GcsServer(const ray::gcs::GcsServerConfig &config,
           io_context_provider_.GetIOContext<pubsub::GcsPublisher>())),
       observability_pubsub_periodical_runner_(PeriodicalRunner::Create(
           io_context_provider_.GetIOContext<pubsub::ObservabilityPublisher>())),
+      resource_load_pull_periodical_runner_(PeriodicalRunner::Create(
+          io_context_provider_.GetIOContext<GcsResourceLoadPuller>())),
       periodical_runner_(
           PeriodicalRunner::Create(io_context_provider_.GetDefaultIOContext())),
       is_started_(false),
@@ -500,35 +516,35 @@ void GcsServer::InitGcsResourceManager(const GcsInitData &gcs_init_data) {
       *gcs_resource_manager_,
       RayConfig::instance().gcs_max_active_rpcs_per_handler()));
 
-  periodical_runner_->RunFnPeriodically(
+  resource_load_puller_ = std::make_unique<GcsResourceLoadPuller>(
+      io_context_provider_.GetIOContext<GcsResourceLoadPuller>(),
+      io_context_provider_.GetDefaultIOContext(),
+      resource_load_pull_raylet_client_pool_,
+      /*apply_on_main=*/
+      [this](rpc::ResourcesData resources) {
+        // TODO(vitsai): Remove duplicate reporting to GcsResourceManager
+        // after verifying that non-autoscaler paths are taken care of.
+        // Currently, GcsResourceManager aggregates reporting from different
+        // sources at different intervals, leading to an obviously inconsistent
+        // view.
+        //
+        // Once autoscaler is completely moved to the new mode of consistent
+        // per-node reporting, remove this if it is not needed anymore.
+        gcs_resource_manager_->UpdateResourceLoads(resources);
+        gcs_autoscaler_state_manager_->UpdateResourceLoadAndUsage(std::move(resources));
+      });
+  resource_load_pull_periodical_runner_->RunFnPeriodically(
       [this] {
-        for (const auto &alive_node : gcs_node_manager_->GetAllAliveNodes()) {
-          auto remote_address = rpc::RayletClientPool::GenerateRayletAddress(
+        const auto alive_nodes = gcs_node_manager_->GetAllAliveNodes();
+        std::vector<rpc::Address> raylet_addresses;
+        raylet_addresses.reserve(alive_nodes.size());
+        for (const auto &alive_node : alive_nodes) {
+          raylet_addresses.push_back(rpc::RayletClientPool::GenerateRayletAddress(
               alive_node.first,
               alive_node.second->node_manager_address(),
-              alive_node.second->node_manager_port());
-          auto raylet_client = raylet_client_pool_.GetOrConnectByAddress(remote_address);
-
-          // GetResourceLoad will also get usage. Historically it didn't.
-          raylet_client->GetResourceLoad([this](auto &status, auto &&load_and_usage) {
-            if (status.ok()) {
-              // TODO(vitsai): Remove duplicate reporting to GcsResourceManager
-              // after verifying that non-autoscaler paths are taken care of.
-              // Currently, GcsResourceManager aggregates reporting from different
-              // sources at different intervals, leading to an obviously inconsistent
-              // view.
-              //
-              // Once autoscaler is completely moved to the new mode of consistent
-              // per-node reporting, remove this if it is not needed anymore.
-              gcs_resource_manager_->UpdateResourceLoads(load_and_usage.resources());
-              gcs_autoscaler_state_manager_->UpdateResourceLoadAndUsage(
-                  std::move(*load_and_usage.mutable_resources()));
-            } else {
-              RAY_LOG_EVERY_N(WARNING, 10)
-                  << "Failed to get the resource load: " << status.ToString();
-            }
-          });
+              alive_node.second->node_manager_port()));
         }
+        resource_load_puller_->Pull(std::move(raylet_addresses));
       },
       RayConfig::instance().gcs_pull_resource_loads_period_milliseconds(),
       "RayletLoadPulled");

src/ray/gcs/gcs_server.h

@@ -260,11 +260,14 @@ class GcsServer {
   rpc::ClientCallManager client_call_manager_;
   rpc::RayletClientPool raylet_client_pool_;
   rpc::CoreWorkerClientPool worker_client_pool_;
+  rpc::ClientCallManager resource_load_pull_client_call_manager_;
+  rpc::RayletClientPool resource_load_pull_raylet_client_pool_;
   std::shared_ptr<ClusterResourceScheduler> cluster_resource_scheduler_;
   std::unique_ptr<gcs::GcsTableStorage> gcs_table_storage_;
   /// gcs_resource_manager_ depends on cluster_lease_manager_.
   std::unique_ptr<GcsResourceManager> gcs_resource_manager_;
   std::unique_ptr<GcsAutoscalerStateManager> gcs_autoscaler_state_manager_;
+  std::unique_ptr<GcsResourceLoadPuller> resource_load_puller_;
   /// A publisher for publishing gcs messages (control-plane pubsub channels).
   std::unique_ptr<pubsub::GcsPublisher> gcs_publisher_;
   /// Publisher for observability pubsub (logs, errors, dashboard resource JSON).
@@ -304,6 +307,8 @@ class GcsServer {
   /// gRPC based pubsub's periodical runner.
   std::shared_ptr<PeriodicalRunner> pubsub_periodical_runner_;
   std::shared_ptr<PeriodicalRunner> observability_pubsub_periodical_runner_;
+  /// The resource load pull's periodical runner.
+  std::shared_ptr<PeriodicalRunner> resource_load_pull_periodical_runner_;
   /// The runner to run function periodically.
   std::shared_ptr<PeriodicalRunner> periodical_runner_;
   /// GCS service state flag, which is used for unit tests.

src/ray/gcs/gcs_server_io_context_policy.h

@@ -21,6 +21,7 @@
 
 #include "ray/gcs/gcs_kv_manager.h"
 #include "ray/gcs/gcs_node_manager.h"
+#include "ray/gcs/gcs_resource_load_puller.h"
 #include "ray/gcs/gcs_task_manager.h"
 #include "ray/observability/ray_event_recorder.h"
 #include "ray/pubsub/gcs_publisher.h"
@@ -64,6 +65,8 @@ struct GcsServerIOContextPolicy {
       return IndexOf("internal_kv_io_context");
     } else if constexpr (std::is_same_v<T, GcsNodeManager>) {
       return IndexOf("node_manager_io_context");
+    } else if constexpr (std::is_same_v<T, GcsResourceLoadPuller>) {
+      return IndexOf("resource_load_pull_io_context");
     } else {
       // default io context
       return -1;
@@ -74,7 +77,7 @@ struct GcsServerIOContextPolicy {
   // and a complete set of those returned from GetDedicatedIOContextIndex. Or you
   // can get runtime crashes when accessing a missing name, or get leaks by
   // creating unused threads.
-  constexpr static std::array<IOContextMetadata, 7> kAllDedicatedIOContexts{{
+  constexpr static std::array<IOContextMetadata, 8> kAllDedicatedIOContexts{{
       // task_io_context only runs GcsTaskManager, which ingests and serves
       // task-state events (observability) and drops events under load by design.
       // It is not on the GCS control plane, so a backlog here (e.g. under a
@@ -99,6 +102,9 @@ struct GcsServerIOContextPolicy {
       {"node_manager_io_context",
        /*enable_lag_probe=*/true,
        /*used_for_health_check=*/true},
+      {"resource_load_pull_io_context",
+       /*enable_lag_probe=*/true,
+       /*used_for_health_check=*/false},
   }};
 
   // Returns int (not size_t) to match GetDedicatedIOContextIndex's return type and

src/ray/gcs/tests/BUILD.bazel

@@ -1,5 +1,21 @@
 load("//bazel:ray.bzl", "ray_cc_library", "ray_cc_test")
 
+ray_cc_test(
+    name = "gcs_resource_load_puller_test",
+    size = "small",
+    srcs = ["gcs_resource_load_puller_test.cc"],
+    tags = ["team:core"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "//src/ray/asio:instrumented_io_context",
+        "//src/ray/common:test_utils",
+        "//src/ray/gcs:gcs_resource_load_puller",
+        "//src/ray/raylet_rpc_client:fake_raylet_client",
+        "//src/ray/raylet_rpc_client:raylet_client_pool",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 ray_cc_test(
     name = "gcs_function_manager_test",
     srcs = ["gcs_function_manager_test.cc"],