update

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

Author: owenowenisme

src/ray/common/event_memory_monitor.cc

@@ -192,8 +192,8 @@ void EventMemoryMonitor::MonitoringThreadMain() {
 
       if (high_modified && IsEnabled()) {
         Disable();
-        kill_workers_callback_(
-            MemoryMonitorUtils::TakeSystemMemorySnapshot(cgroup_path_));
+        kill_workers_callback_(MemoryMonitorUtils::TakeSystemMemorySnapshot(cgroup_path_),
+                               MemoryMonitorInterface::kNull);
       }
     } else {
       RAY_LOG(ERROR) << absl::StrFormat(

src/ray/common/memory_monitor_interface.h

@@ -48,11 +48,17 @@ struct SystemMemorySnapshot {
 using ProcessesMemorySnapshot = absl::flat_hash_map<pid_t, int64_t>;
 
 /**
- * @brief Callback that runs at each monitoring interval.
+ * @brief Callback invoked by the memory monitor when worker killing should be triggered.
  *
- * \param system_memory snapshot of system memory information.
+ * \param system_memory snapshot of system memory information at the time the monitor
+ *        decided to fire.
+ * \param threshold_bytes the byte threshold against which the monitor evaluated
+ *        \p system_memory.used_bytes. For monitors that are event-driven rather than
+ *        threshold-driven (e.g. cgroup memory.events / PSI) this should be
+ *        MemoryMonitorInterface::kNull to indicate no single byte threshold applies.
  */
-using KillWorkersCallback = std::function<void(SystemMemorySnapshot system_memory)>;
+using KillWorkersCallback =
+    std::function<void(SystemMemorySnapshot system_memory, int64_t threshold_bytes)>;
 
 /**
  * @brief implementations of this interface monitors the memory usage of the node

src/ray/common/pressure_memory_monitor.cc

@@ -161,8 +161,8 @@ void PressureMemoryMonitor::MonitoringThreadMain() {
     if (fds[0].revents & POLLPRI) {
       if (IsEnabled()) {
         Disable();
-        kill_workers_callback_(
-            MemoryMonitorUtils::TakeSystemMemorySnapshot(cgroup_path_));
+        kill_workers_callback_(MemoryMonitorUtils::TakeSystemMemorySnapshot(cgroup_path_),
+                               MemoryMonitorInterface::kNull);
       }
     } else if (fds[0].revents & POLLERR) {
       RAY_LOG(ERROR) << "Got POLLERR while monitoring memory pressure. "

src/ray/common/tests/event_memory_monitor_test.cc

@@ -67,7 +67,7 @@ TEST_F(EventMemoryMonitorTest, TestNonexistentCgroupPathFailsGracefully) {
   std::string nonexistent_path = "/nonexistent/cgroup/path";
   StatusSetOr<std::unique_ptr<EventMemoryMonitor>, StatusT::IOError> result =
       EventMemoryMonitor::Create(std::move(nonexistent_path),
-                                 [](SystemMemorySnapshot) {});
+                                 [](SystemMemorySnapshot, int64_t) {});
 
   ASSERT_TRUE(result.has_error())
       << "Failed to catch invalid cgroup path when creating EventMemoryMonitor";
@@ -79,7 +79,8 @@ TEST_F(EventMemoryMonitorTest, TestMissingMemoryEventsFileFailsGracefully) {
   RAY_CHECK(empty_dir_or.ok()) << empty_dir_or.status().ToString();
   std::unique_ptr<TempDirectory> empty_dir = std::move(empty_dir_or.value());
   StatusSetOr<std::unique_ptr<EventMemoryMonitor>, StatusT::IOError> result =
-      EventMemoryMonitor::Create(empty_dir->GetPath(), [](SystemMemorySnapshot) {});
+      EventMemoryMonitor::Create(empty_dir->GetPath(),
+                                 [](SystemMemorySnapshot, int64_t) {});
 
   ASSERT_TRUE(result.has_error())
       << "Failed to catch invalid cgroup configuration when creating EventMemoryMonitor";
@@ -89,7 +90,7 @@ TEST_F(EventMemoryMonitorTest, TestMissingMemoryEventsFileFailsGracefully) {
 TEST_F(EventMemoryMonitorTest, TestSuccessfulCreationWithValidPath) {
   StatusSetOr<std::unique_ptr<EventMemoryMonitor>, StatusT::IOError> result =
       EventMemoryMonitor::Create(mock_cgroup_dir_->GetPath(),
-                                 [](SystemMemorySnapshot) {});
+                                 [](SystemMemorySnapshot, int64_t) {});
   ASSERT_TRUE(result.has_value())
       << "Failed to create EventMemoryMonitor: " << result.message();
   std::unique_ptr<EventMemoryMonitor> monitor = std::move(result.value());
@@ -100,7 +101,11 @@ TEST_F(EventMemoryMonitorTest, TestCallbackCalledWhenHighEventChanges) {
   WriteMemoryEventsFile(events_file_->GetPath(), 0, 0);
 
   auto callback_latch = std::make_shared<boost::latch>(1);
-  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot) {
+  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot,
+                                                  int64_t threshold_bytes) {
+    EXPECT_EQ(threshold_bytes, MemoryMonitorInterface::kNull)
+        << "EventMemoryMonitor is event-driven so it should report kNull as the "
+           "threshold bytes.";
     callback_latch->count_down();
   };
 
@@ -120,7 +125,7 @@ TEST_F(EventMemoryMonitorTest, TestNoCallbackWhenValuesUnchanged) {
   WriteMemoryEventsFile(events_file_->GetPath(), 0, 0);
 
   auto callback_latch = std::make_shared<boost::latch>(1);
-  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot) {
+  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot, int64_t) {
     callback_latch->count_down();
   };
 
@@ -139,7 +144,7 @@ TEST_F(EventMemoryMonitorTest, TestNoCallbackWhenIrrelevantEventChanges) {
   WriteMemoryEventsFile(events_file_->GetPath(), 0, 0);
 
   auto callback_latch = std::make_shared<boost::latch>(1);
-  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot) {
+  KillWorkersCallback callback = [callback_latch](SystemMemorySnapshot, int64_t) {
     callback_latch->count_down();
   };
 
@@ -161,17 +166,17 @@ TEST_F(EventMemoryMonitorTest, TestMultipleCallbacksOnMultipleChanges) {
   auto latch2 = std::make_shared<boost::latch>(1);
   auto latch3 = std::make_shared<boost::latch>(1);
   std::atomic<int> callback_count{0};
-  KillWorkersCallback callback =
-      [&callback_count, latch1, latch2, latch3](SystemMemorySnapshot) {
-        int count = ++callback_count;
-        if (count == 1) {
-          latch1->count_down();
-        } else if (count == 2) {
-          latch2->count_down();
-        } else if (count == 3) {
-          latch3->count_down();
-        }
-      };
+  KillWorkersCallback callback = [&callback_count, latch1, latch2, latch3](
+                                     SystemMemorySnapshot, int64_t) {
+    int count = ++callback_count;
+    if (count == 1) {
+      latch1->count_down();
+    } else if (count == 2) {
+      latch2->count_down();
+    } else if (count == 3) {
+      latch3->count_down();
+    }
+  };
 
   StatusSetOr<std::unique_ptr<EventMemoryMonitor>, StatusT::IOError> result =
       EventMemoryMonitor::Create(mock_cgroup_dir_->GetPath(), callback);

src/ray/common/tests/pressure_memory_monitor_test.cc

@@ -100,7 +100,7 @@ TEST_F(PressureMemoryMonitorTest, TestNonexistentCgroupPathFailsGracefully) {
   MemoryPsi psi = {.mode = "some", .stall_proportion = 0.5f, .stall_duration_s = 2};
   std::string nonexistent_path = "/nonexistent/cgroup/path";
   auto result = PressureMemoryMonitor::Create(
-      psi, std::move(nonexistent_path), [](const SystemMemorySnapshot &) {});
+      psi, std::move(nonexistent_path), [](SystemMemorySnapshot, int64_t) {});
 
   ASSERT_TRUE(result.has_error())
       << "Failed to catch invalid cgroup path when creating PressureMemoryMonitor";
@@ -111,7 +111,7 @@ TEST_F(PressureMemoryMonitorTest, TestMonitorCreationWritesTriggerStringToFile)
   MemoryPsi psi = {.mode = "some", .stall_proportion = 0.5f, .stall_duration_s = 2};
 
   auto result = PressureMemoryMonitor::Create(
-      psi, mock_cgroup_dir_->GetPath(), [](const SystemMemorySnapshot &) {});
+      psi, mock_cgroup_dir_->GetPath(), [](SystemMemorySnapshot, int64_t) {});
   ASSERT_TRUE(result.has_value())
       << "Failed to create PressureMemoryMonitor: " << result.message();
 
@@ -164,7 +164,7 @@ TEST_F(PressureMemoryMonitorTest,
   close(listener);
 
   std::shared_ptr<boost::latch> has_called_once = std::make_shared<boost::latch>(1);
-  auto kill_workers_callback = [has_called_once](const SystemMemorySnapshot &) {
+  auto kill_workers_callback = [has_called_once](SystemMemorySnapshot, int64_t) {
     has_called_once->count_down();
   };
   std::unique_ptr<PressureMemoryMonitor> monitor =

src/ray/common/tests/threshold_memory_monitor_test.cc

@@ -54,13 +54,16 @@ TEST_F(ThresholdMemoryMonitorTest, TestMonitorTriggerCanDetectMemoryUsage) {
   MakeThresholdMemoryMonitor(
       0 /*memory_usage_threshold_bytes*/,
       1 /*refresh_interval_ms*/,
-      [has_checked_once](SystemMemorySnapshot system_memory) {
+      [has_checked_once](SystemMemorySnapshot system_memory, int64_t threshold_bytes) {
         ASSERT_GT(system_memory.total_bytes, 0)
             << "Reported total bytes from cgroup is <= 0. Is the system memory snapshot "
                "taken correctly?";
         ASSERT_GE(system_memory.used_bytes, 0)
             << "Reported used bytes from cgroup is < 0. Is the system memory snapshot "
                "taken correctly?";
+        ASSERT_EQ(threshold_bytes, 0)
+            << "Threshold bytes propagated to the callback should match the value the "
+               "monitor was constructed with.";
         has_checked_once->count_down();
       },
       "" /*root_cgroup_path*/);
@@ -86,10 +89,14 @@ TEST_F(ThresholdMemoryMonitorTest,
   MakeThresholdMemoryMonitor(
       memory_usage_threshold_bytes,  // (70%)
       1 /*refresh_interval_ms*/,
-      [has_checked_once, cgroup_total_bytes](SystemMemorySnapshot system_memory) {
+      [has_checked_once, cgroup_total_bytes, memory_usage_threshold_bytes](
+          SystemMemorySnapshot system_memory, int64_t threshold_bytes) {
         ASSERT_EQ(system_memory.total_bytes, cgroup_total_bytes)
             << "Unexpected total bytes read from cgroup. Are we correctly reading memory "
                "from the cgroup?";
+        ASSERT_EQ(threshold_bytes, memory_usage_threshold_bytes)
+            << "Threshold bytes propagated to the callback should match the value the "
+               "monitor was constructed with.";
         has_checked_once->count_down();
       },
       cgroup_dir /*root_cgroup_path*/);
@@ -117,7 +124,7 @@ TEST_F(ThresholdMemoryMonitorTest,
   MakeThresholdMemoryMonitor(
       memory_usage_threshold_bytes,  // (70%)
       1 /*refresh_interval_ms*/,
-      [callback_triggered](SystemMemorySnapshot system_memory) {
+      [callback_triggered](SystemMemorySnapshot system_memory, int64_t threshold_bytes) {
         callback_triggered->store(true);
       },
       cgroup_dir /*root_cgroup_path*/);

src/ray/common/threshold_memory_monitor.cc

@@ -60,7 +60,8 @@ ThresholdMemoryMonitor::ThresholdMemoryMonitor(KillWorkersCallback kill_workers_
 
         if (is_usage_above_threshold && IsEnabled()) {
           Disable();
-          kill_workers_callback_(std::move(cur_memory_snapshot));
+          kill_workers_callback_(std::move(cur_memory_snapshot),
+                                 memory_usage_threshold_bytes_);
         }
       },
       monitor_interval_ms,

src/ray/raylet/node_manager.cc

@@ -3048,9 +3048,9 @@ std::optional<syncer::RaySyncMessage> NodeManager::CreateSyncMessage(
 
 // Picks the workers and kills the process if the memory usage is above the threshold.
 KillWorkersCallback NodeManager::CreateKillWorkersCallback() {
-  return [this](SystemMemorySnapshot system_memory_snapshot) {
+  return [this](SystemMemorySnapshot system_memory_snapshot, int64_t threshold_bytes) {
     io_service_.post(
-        [this, system_memory = std::move(system_memory_snapshot)]() {
+        [this, system_memory = std::move(system_memory_snapshot), threshold_bytes]() {
           ProcessesMemorySnapshot process_memory_snapshot =
               MemoryMonitorUtils::TakePerProcessMemorySnapshot();
           std::vector<std::shared_ptr<WorkerInterface>> workers =
@@ -3074,25 +3074,13 @@ KillWorkersCallback NodeManager::CreateKillWorkersCallback() {
             return;
           }
 
-          // Compute the memory usage threshold
-          int64_t total_memory_bytes = system_memory.total_bytes;
-          int64_t computed_threshold_bytes = MemoryMonitorUtils::GetMemoryThreshold(
-              total_memory_bytes,
-              RayConfig::instance().memory_usage_threshold(),
-              RayConfig::instance().min_memory_free_bytes(),
-              initial_config_.enable_resource_isolation,
-              *cgroup_manager_);
-          float computed_threshold_fraction =
-              static_cast<float>(computed_threshold_bytes) /
-              static_cast<float>(total_memory_bytes);
-
           std::string oom_kill_details = CreateOomKillMessageDetails(
               workers_to_kill_and_should_retry,
               self_node_id_,
               system_memory,
               store_client_->GetMemoryUsage().value_or("Not available"),
               process_memory_snapshot,
-              computed_threshold_fraction);
+              threshold_bytes);
           std::string oom_kill_suggestions =
               CreateOomKillMessageSuggestions(workers_to_kill_and_should_retry);
 
@@ -3162,7 +3150,7 @@ std::string NodeManager::CreateOomKillMessageDetails(
     const SystemMemorySnapshot &system_memory_snapshot,
     const std::string &object_store_memory_usage,
     const ProcessesMemorySnapshot &process_memory_snapshot,
-    float usage_threshold) const {
+    int64_t threshold_bytes) const {
   if (workers_to_kill.empty()) {
     return "";
   }
@@ -3175,6 +3163,25 @@ std::string NodeManager::CreateOomKillMessageDetails(
       "%.2f",
       static_cast<float>(system_memory_snapshot.total_bytes) / 1024 / 1024 / 1024);
 
+  // Render the threshold. For threshold-based monitors `threshold_bytes` is the
+  // exact byte value the monitor compared against. For event- or pressure-driven
+  // monitors (kNull) there is no single byte threshold to report, so we fall back
+  // to describing the triggering signal instead.
+  std::string threshold_str;
+  if (threshold_bytes == MemoryMonitorInterface::kNull) {
+    threshold_str = "cgroup memory event/pressure signal (no fixed threshold)";
+  } else {
+    float threshold_fraction =
+        system_memory_snapshot.total_bytes > 0
+            ? static_cast<float>(threshold_bytes) /
+                  static_cast<float>(system_memory_snapshot.total_bytes)
+            : 0.0f;
+    threshold_str =
+        absl::StrFormat("%.2fGB (%f)",
+                        static_cast<float>(threshold_bytes) / 1024 / 1024 / 1024,
+                        threshold_fraction);
+  }
+
   const auto &first_worker = workers_to_kill.front().first;
   std::string node_ip = first_worker->IpAddress();
 
@@ -3221,7 +3228,7 @@ std::string NodeManager::CreateOomKillMessageDetails(
 
   return absl::StrFormat(
       "Memory on the node (IP: %s, ID: %s) was %sGB / %sGB (%f), "
-      "which exceeds the memory usage threshold of %f; "
+      "which exceeds the memory usage threshold of %s; "
       "Object store memory usage: [%s]; "
       "Ray killed %d worker(s) based on the killing policy: "
       "[%s]; "
@@ -3233,7 +3240,7 @@ std::string NodeManager::CreateOomKillMessageDetails(
       used_bytes_gb,
       total_bytes_gb,
       usage_fraction,
-      usage_threshold,
+      threshold_str,
       absl::StrReplaceAll(object_store_memory_usage, {{"\n", "; "}}),
       workers_to_kill.size(),
       absl::StrJoin(worker_details, "; "),

src/ray/raylet/node_manager.h

@@ -772,7 +772,9 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
    * @param node_id The ID of the node.
    * @param system_memory_snapshot The snapshot of the system memory.
    * @param process_memory_snapshot The snapshot of the process memory.
-   * @param usage_threshold The memory limit.
+   * @param threshold_bytes The byte threshold used by the memory monitor to decide
+   *        that the kill should fire. Pass MemoryMonitorInterface::kNull if the
+   *        monitor is event-driven and has no single byte threshold.
    * @return The detail message for the workers that are killed due to memory running low.
    */
   std::string CreateOomKillMessageDetails(
@@ -782,7 +784,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
       const SystemMemorySnapshot &system_memory_snapshot,
       const std::string &object_store_memory_usage,
       const ProcessesMemorySnapshot &process_memory_snapshot,
-      float usage_threshold) const;
+      int64_t threshold_bytes) const;
 
   /**
    * @param workers_to_kill The workers to print the kill suggestions for.