[Core] (Resource Isolation 12/n) Switch group killing policy to by time killing policy

doc/source/ray-core/doc_code/ray_oom_prevention.py

@@ -10,7 +10,7 @@
 # __last_task_start__
 import ray
 
-@ray.remote(max_retries=-1)
+@ray.remote(max_retries=0)
 def leaks_memory():
     chunks = []
     bits_to_allocate = 8 * 100 * 1024 * 1024  # ~100 MiB
@@ -67,11 +67,11 @@ def allocate(self, bytes_to_allocate: float) -> None:
     max_restarts=0, max_task_retries=0, name="second_actor"
 ).remote()
 
-# each task requests 0.3 of the system memory when the memory threshold is 0.4.
-allocate_bytes = get_additional_bytes_to_reach_memory_usage_pct(0.3)
+# We want total bytes to exceed 40% threshold to trigger the memory monitor.
+allocate_bytes = get_additional_bytes_to_reach_memory_usage_pct(0.5)
 
-first_actor_task = first_actor.allocate.remote(allocate_bytes)
-second_actor_task = second_actor.allocate.remote(allocate_bytes)
+first_actor_task = first_actor.allocate.remote(0.9 * allocate_bytes)
+second_actor_task = second_actor.allocate.remote(0.1* allocate_bytes)
 
 error_thrown = False
 try:

src/ray/common/BUILD.bazel

@@ -135,6 +135,8 @@ ray_cc_library(
     ],
     deps = [
         ":memory_monitor_interface",
+        ":ray_config",
+        "//src/ray/common/cgroup2:cgroup_manager_interface",
         "//src/ray/util:compat",
         "//src/ray/util:logging",
         "@boost//:algorithm",

src/ray/common/cgroup2/BUILD.bazel

@@ -58,7 +58,10 @@ ray_cc_library(
         "cgroup_manager.h",
         "scoped_cgroup_operation.h",
     ],
-    visibility = [":__subpackages__"],
+    visibility = [
+        ":__subpackages__",
+        "//src/ray/common/tests:__pkg__",
+    ],
     deps = [
         ":cgroup_driver_interface",
         ":cgroup_manager_interface",
@@ -74,7 +77,11 @@ ray_cc_library(
     hdrs = [
         "noop_cgroup_manager.h",
     ],
-    visibility = ["//visibility:public"],
+    visibility = [
+        ":__subpackages__",
+        "//src/ray/common/tests:__pkg__",
+        "//src/ray/raylet/tests:__pkg__",
+    ],
     deps = [
         ":cgroup_driver_interface",
         ":cgroup_manager_interface",
@@ -117,7 +124,10 @@ ray_cc_library(
     target_compatible_with = [
         "@platforms//os:linux",
     ],
-    visibility = [":__subpackages__"],
+    visibility = [
+        ":__subpackages__",
+        "//src/ray/common/tests:__pkg__",
+    ],
     deps = [
         ":cgroup_driver_interface",
         "//src/ray/common:status",

src/ray/common/memory_monitor_utils.cc

@@ -22,6 +22,7 @@
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "ray/common/memory_monitor_interface.h"
+#include "ray/common/ray_config.h"
 #include "ray/util/logging.h"
 
 namespace ray {
@@ -292,23 +293,59 @@ int64_t MemoryMonitorUtils::NullableMin(int64_t left, int64_t right) {
   }
 }
 
-int64_t MemoryMonitorUtils::GetMemoryThreshold(int64_t total_memory_bytes,
-                                               float usage_threshold,
-                                               int64_t min_memory_free_bytes) {
+int64_t MemoryMonitorUtils::GetMemoryThreshold(
+    int64_t total_memory_bytes,
+    float usage_threshold,
+    int64_t min_memory_free_bytes,
+    bool resource_isolation_enabled,
+    const CgroupManagerInterface &cgroup_manager) {
   RAY_CHECK_GE(total_memory_bytes, MemoryMonitorInterface::kNull);
   RAY_CHECK_GE(min_memory_free_bytes, MemoryMonitorInterface::kNull);
-  RAY_CHECK_GE(usage_threshold, 0);
-  RAY_CHECK_LE(usage_threshold, 1);
+  RAY_CHECK_GE(usage_threshold, 0)
+      << "Invalid configuration: usage_threshold must be >= 0";
+  RAY_CHECK_LE(usage_threshold, 1)
+      << "Invalid configuration: usage_threshold must be <= 1";
 
-  int64_t threshold_fraction = (int64_t)(total_memory_bytes * usage_threshold);
+  int64_t resolved_memory_threshold_bytes;
+  int64_t threshold_fraction = static_cast<int64_t>(total_memory_bytes * usage_threshold);
 
   if (min_memory_free_bytes > MemoryMonitorInterface::kNull) {
     int64_t threshold_absolute = total_memory_bytes - min_memory_free_bytes;
     RAY_CHECK_GE(threshold_absolute, 0);
-    return std::max(threshold_fraction, threshold_absolute);
+    resolved_memory_threshold_bytes = std::max(threshold_fraction, threshold_absolute);
   } else {
-    return threshold_fraction;
+    resolved_memory_threshold_bytes = threshold_fraction;
+  }
+
+  if (resource_isolation_enabled) {
+    StatusOr<std::string> user_memory_max_bytes_or =
+        cgroup_manager.GetUserCgroupConstraintValue("memory.max");
+    RAY_CHECK(user_memory_max_bytes_or.ok()) << absl::StrFormat(
+        "Failed to get user cgroup memory limit when setting up memory monitor: %s",
+        user_memory_max_bytes_or.ToString());
+    std::string user_memory_max_bytes_str = user_memory_max_bytes_or.value();
+
+    if (!user_memory_max_bytes_str.empty() &&
+        std::all_of(user_memory_max_bytes_str.begin(),
+                    user_memory_max_bytes_str.end(),
+                    ::isdigit)) {
+      int64_t user_memory_max_bytes = std::stoll(user_memory_max_bytes_str);
+      int64_t reaction_buffer_bytes =
+          std::min(static_cast<int64_t>(total_memory_bytes *
+                                        kDefaultThresholdMonitorReactionBufferProportion),
+                   RayConfig::instance().max_threshold_monitor_reaction_buffer_bytes());
+      resolved_memory_threshold_bytes = user_memory_max_bytes - reaction_buffer_bytes;
+      RAY_CHECK_GE(resolved_memory_threshold_bytes, 0) << absl::StrFormat(
+          "Available user task memory is less than the kill memory buffer bytes: "
+          "%d < %d. This means the available memory for user proceses is likely "
+          "less than 5%% of total memory. Please consider decreasing the proportion "
+          "of reserved system memory if it was custom set.",
+          user_memory_max_bytes,
+          reaction_buffer_bytes);
+    }
   }
+
+  return resolved_memory_threshold_bytes;
 }
 
 int64_t MemoryMonitorUtils::GetProcessUsedMemoryBytes(

src/ray/common/memory_monitor_utils.h

@@ -21,6 +21,7 @@
 #include <tuple>
 #include <vector>
 
+#include "ray/common/cgroup2/cgroup_manager_interface.h"
 #include "ray/common/memory_monitor_interface.h"
 #include "ray/util/compat.h"
 
@@ -69,11 +70,18 @@ class MemoryMonitorUtils {
    * @param usage_threshold A value in [0-1] to indicate the max usage.
    * @param min_memory_free_bytes The min amount of free space to maintain before it is
    *        exceeding the threshold.
+   * @param resource_isolation_enabled Whether resource isolation is enabled. Used
+   *        to determine if the threshold should be calculated based on the cgroup
+   * constraints.
+   * @param cgroup_manager The cgroup manager to fetch the upper bound memory constraints
+   * from.
    * @return The memory threshold.
    */
   static int64_t GetMemoryThreshold(int64_t total_memory_bytes,
                                     float usage_threshold,
-                                    int64_t min_memory_free_bytes);
+                                    int64_t min_memory_free_bytes,
+                                    bool resource_isolation_enabled,
+                                    const CgroupManagerInterface &cgroup_manager);
 
   /**
    * @brief Gets the used memory for a process from the process memory snapshot.
@@ -221,6 +229,8 @@ class MemoryMonitorUtils {
   static constexpr char kCgroupsV2MemoryStatActiveFileKey[] = "active_file";
   static constexpr char kProcDirectory[] = "/proc";
   static constexpr char kCommandlinePath[] = "cmdline";
+
+  static constexpr double kDefaultThresholdMonitorReactionBufferProportion = 0.05;
 };
 
 }  // namespace ray

src/ray/common/ray_config_def.h

@@ -72,9 +72,10 @@ RAY_CONFIG(uint64_t, raylet_check_gc_period_milliseconds, 100)
 /// it will start killing processes to free up the space.
 /// Note: when resource isolation is enabled, the memory usage threshold is set to
 /// total memory - system reserved memory (can be specified in ray start) -
-/// kill_memory_buffer_bytes. Notice that the formula does not account for object store
-/// memory in system reserved memory. To configure the usage threshold, please adjust the
-/// system reserved memory in ray start command instead. Ranging from [0, 1]
+/// max(5% of total memory, max_threshold_monitor_reaction_buffer_bytes).
+/// Notice that the formula does not account for object store memory in system reserved
+/// memory. To configure the usage threshold, please adjust the system reserved memory in
+/// ray start command instead. Ranging from [0, 1]
 RAY_CONFIG(float, memory_usage_threshold, 0.95)
 
 /// The interval between runs of the memory usage monitor.
@@ -90,9 +91,24 @@ RAY_CONFIG(uint64_t, memory_monitor_refresh_ms, 250)
 /// means 6.4 GB of the memory will not be usable.
 RAY_CONFIG(int64_t, min_memory_free_bytes, (int64_t)-1)
 
-/// The amount of memory to free under the memory usage threshold when
-/// killing workers via the worker killing policy.
-RAY_CONFIG(uint64_t, kill_memory_buffer_bytes, 3ULL * 1024 * 1024 * 1024)  // 3GB
+/// The maximum amount of memory to free under the memory usage threshold when
+/// killing workers via the worker killing policy. The system will by default
+/// free up to 5% of total memory under the threshold, capping at
+/// max_kill_memory_buffer_bytes.
+RAY_CONFIG(int64_t, max_kill_memory_buffer_bytes, 3ULL * 1024 * 1024 * 1024)  // 3GiB cap
+
+/// The threshold monitor is poll based and may miss memory bursts occurring between
+/// polls. This is the maximum buffer size that can be subtracted from memory.max to give
+/// the threshold monitor time to react before memory max is reached under resource
+/// isolation. The system will by default provide 5% of total memory as reaction buffer,
+/// capping at max_threshold_monitor_reaction_buffer_bytes.
+RAY_CONFIG(int64_t,
+           max_threshold_monitor_reaction_buffer_bytes,
+           2LL * 1024 * 1024 * 1024)  // 2GiB
+
+/// When true, use the legacy group-by-owner worker killing policy instead of the
+/// default time-based policy.
+RAY_CONFIG(bool, worker_killing_policy_by_group, false)
 
 /// The reserved memory bytes for system processes
 /// enforced via cgroup memory.min constraint which guarantees

src/ray/common/tests/BUILD.bazel

@@ -112,7 +112,11 @@ ray_cc_test(
         "//src/ray/common:id",
         "//src/ray/common:memory_monitor_test_fixture",
         "//src/ray/common:memory_monitor_utils",
+        "//src/ray/common:ray_config",
+        "//src/ray/common/cgroup2:cgroup_manager",
         "//src/ray/common/cgroup2:cgroup_test_utils",
+        "//src/ray/common/cgroup2:fake_cgroup_driver",
+        "//src/ray/common/cgroup2:noop_cgroup_manager",
         "//src/ray/util:process",
         "@boost//:filesystem",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -135,7 +139,9 @@ ray_cc_test(
     deps = [
         "//src/ray/common:memory_monitor_interface",
         "//src/ray/common:memory_monitor_test_fixture",
+        "//src/ray/common:memory_monitor_utils",
         "//src/ray/common:threshold_memory_monitor",
+        "//src/ray/common/cgroup2:noop_cgroup_manager",
         "@boost//:thread",
     ],
 )

src/ray/common/tests/memory_monitor_utils_test.cc

@@ -18,11 +18,16 @@
 #include <fstream>
 #include <memory>
 #include <string>
+#include <unordered_map>
 
 #include "gtest/gtest.h"
+#include "ray/common/cgroup2/cgroup_manager.h"
 #include "ray/common/cgroup2/cgroup_test_utils.h"
+#include "ray/common/cgroup2/fake_cgroup_driver.h"
+#include "ray/common/cgroup2/noop_cgroup_manager.h"
 #include "ray/common/id.h"
 #include "ray/common/memory_monitor_test_fixture.h"
+#include "ray/common/ray_config.h"
 #include "ray/util/process.h"
 
 namespace ray {
@@ -228,25 +233,82 @@ TEST_F(MemoryMonitorUtilsTest, TestCgroupNonexistentUsageFileReturnskNull) {
 }
 
 TEST_F(MemoryMonitorUtilsTest, TestGetMemoryThresholdTakeGreaterOfTheTwoValues) {
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0.5, 0), 100);
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0.5, 60), 50);
+  NoopCgroupManager noop_cgroup_manager;
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 0.5, 0, false, noop_cgroup_manager),
+      100);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 0.5, 60, false, noop_cgroup_manager),
+      50);
 
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 1, 10), 100);
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 1, 100), 100);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 1, 10, false, noop_cgroup_manager),
+      100);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 1, 100, false, noop_cgroup_manager),
+      100);
 
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0.1, 100), 10);
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0, 10), 90);
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0, 100), 0);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 0.1, 100, false, noop_cgroup_manager),
+      10);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 0, 10, false, noop_cgroup_manager), 90);
+  ASSERT_EQ(
+      MemoryMonitorUtils::GetMemoryThreshold(100, 0, 100, false, noop_cgroup_manager), 0);
 
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 0, MemoryMonitorInterface::kNull),
+  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(
+                100, 0, MemoryMonitorInterface::kNull, false, noop_cgroup_manager),
             0);
-  ASSERT_EQ(
-      MemoryMonitorUtils::GetMemoryThreshold(100, 0.5, MemoryMonitorInterface::kNull),
-      50);
-  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(100, 1, MemoryMonitorInterface::kNull),
+  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(
+                100, 0.5, MemoryMonitorInterface::kNull, false, noop_cgroup_manager),
+            50);
+  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(
+                100, 1, MemoryMonitorInterface::kNull, false, noop_cgroup_manager),
             100);
 }
 
+TEST_F(
+    MemoryMonitorUtilsTest,
+    TestGetMemoryThresholdWithResourceIsolationUsesUpperBoundConstraintToComputeThreshold) {
+  // Create a fake cgroup directory using MockCgroupMemoryUsage.
+  std::string cgroup_dir = MockCgroupMemoryUsage(
+      /*total_bytes=*/16LL * 1024 * 1024 * 1024,
+      /*current_bytes=*/5LL * 1024 * 1024 * 1024,
+      /*inactive_file_bytes=*/100 * 1024 * 1024,
+      /*active_file_bytes=*/50 * 1024 * 1024);
+
+  // Create a CgroupManager backed by a fake cgroup driver on the mock cgroup directory.
+  std::shared_ptr<std::unordered_map<std::string, FakeCgroup>> cgroups =
+      std::make_shared<std::unordered_map<std::string, FakeCgroup>>();
+  cgroups->emplace(cgroup_dir, FakeCgroup{cgroup_dir, {5}, {}, {"cpu", "memory"}, {}});
+  std::unique_ptr<FakeCgroupDriver> driver = FakeCgroupDriver::Create(cgroups);
+
+  int64_t user_memory_max_bytes = 10LL * 1024 * 1024 * 1024;  // 10 GB
+  StatusOr<std::unique_ptr<CgroupManager>> result =
+      CgroupManager::Create(cgroup_dir,
+                            "node_id_123",
+                            /*system_reserved_cpu_weight=*/100,
+                            /*system_memory_bytes_min=*/1LL * 1024 * 1024 * 1024,
+                            /*system_memory_bytes_low=*/1LL * 1024 * 1024 * 1024,
+                            /*user_memory_high_bytes=*/user_memory_max_bytes,
+                            user_memory_max_bytes,
+                            std::move(driver));
+  std::unique_ptr<CgroupManager> cgroup_manager = std::move(result.value());
+
+  // Reaction buffer defaults to 5% of total memory. If
+  // kDefaultThresholdMonitorReactionBufferProportion is changed, this should be changed
+  // accordingly.
+  int64_t expected_threshold =
+      user_memory_max_bytes - static_cast<int64_t>(16LL * 1024 * 1024 * 1024 * 0.05);
+  ASSERT_EQ(MemoryMonitorUtils::GetMemoryThreshold(
+                /*total_memory_bytes=*/16LL * 1024 * 1024 * 1024,
+                /*usage_threshold=*/0.5,
+                /*min_memory_free_bytes=*/MemoryMonitorInterface::kNull,
+                /*resource_isolation_enabled=*/true,
+                *cgroup_manager),
+            expected_threshold);
+}
+
 TEST_F(MemoryMonitorUtilsTest, TestGetPidsFromDirOnlyReturnsNumericFilenames) {
   std::string proc_dir = UniqueID::FromRandom().Hex();
   boost::filesystem::create_directory(proc_dir);

src/ray/common/tests/threshold_memory_monitor_test.cc

@@ -23,6 +23,7 @@
 #include <utility>
 
 #include "gtest/gtest.h"
+#include "ray/common/cgroup2/noop_cgroup_manager.h"
 #include "ray/common/memory_monitor_interface.h"
 #include "ray/common/memory_monitor_test_fixture.h"
 #include "ray/common/memory_monitor_utils.h"
@@ -79,8 +80,9 @@ TEST_F(ThresholdMemoryMonitorTest,
 
   std::shared_ptr<boost::latch> has_checked_once = std::make_shared<boost::latch>(1);
 
-  int64_t memory_usage_threshold_bytes =
-      MemoryMonitorUtils::GetMemoryThreshold(cgroup_total_bytes, 0.7f, -1);
+  NoopCgroupManager noop_cgroup_manager;
+  int64_t memory_usage_threshold_bytes = MemoryMonitorUtils::GetMemoryThreshold(
+      cgroup_total_bytes, 0.7f, -1, false, noop_cgroup_manager);
   MakeThresholdMemoryMonitor(
       memory_usage_threshold_bytes,  // (70%)
       1 /*refresh_interval_ms*/,
@@ -109,8 +111,9 @@ TEST_F(ThresholdMemoryMonitorTest,
   std::shared_ptr<std::atomic<bool>> callback_triggered =
       std::make_shared<std::atomic<bool>>(false);
 
-  int64_t memory_usage_threshold_bytes =
-      MemoryMonitorUtils::GetMemoryThreshold(cgroup_total_bytes, 0.7f, -1);
+  NoopCgroupManager noop_cgroup_manager;
+  int64_t memory_usage_threshold_bytes = MemoryMonitorUtils::GetMemoryThreshold(
+      cgroup_total_bytes, 0.7f, -1, false, noop_cgroup_manager);
   MakeThresholdMemoryMonitor(
       memory_usage_threshold_bytes,  // (70%)
       1 /*refresh_interval_ms*/,