[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

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,10 @@ ray_cc_library(
     hdrs = [
         "noop_cgroup_manager.h",
     ],
-    visibility = ["//visibility:public"],
+    visibility = [
+        ":__subpackages__",
+        "//src/ray/common/tests:__pkg__",
+    ],
     deps = [
         ":cgroup_driver_interface",
         ":cgroup_manager_interface",
@@ -117,7 +123,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,57 @@ 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);
+      resolved_memory_threshold_bytes =
+          user_memory_max_bytes -
+          RayConfig::instance().threshold_monitor_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. Please consider using a host with more memory. If the current host "
+          "memory size must be kept, please adjust the threshold monitor reaction buffer "
+          "size.",
+          user_memory_max_bytes,
+          RayConfig::instance().threshold_monitor_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.

src/ray/common/ray_config_def.h

@@ -94,6 +94,18 @@ RAY_CONFIG(int64_t, min_memory_free_bytes, (int64_t)-1)
 /// killing workers via the worker killing policy.
 RAY_CONFIG(uint64_t, kill_memory_buffer_bytes, 3ULL * 1024 * 1024 * 1024)  // 3GB
 
+/// The threshold monitor is poll based and may miss memory bursts occurring between
+/// polls. This buffer is subtracted from memory.max to give the threshold monitor
+/// time to react before memory max is reached under resource isolation.
+RAY_CONFIG(int64_t,
+           threshold_monitor_reaction_buffer_bytes,
+           2LL * 1024 * 1024 * 1024)  // 2GB
+
+/// Whether to use the group-by-owner worker killing policy instead of the
+/// default time-based worker killing policy. When true, workers killed by the
+/// by group killing policy (legacy 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
 /// that the system processes' memory will not be reclaimed under any conditions.

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());
+
+  // When resource isolation is enabled, GetMemoryThreshold uses
+  // memory.max - threshold_monitor_reaction_buffer_bytes regardless of other parameters.
+  int64_t expected_threshold =
+      user_memory_max_bytes -
+      RayConfig::instance().threshold_monitor_reaction_buffer_bytes();
+  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*/,

src/ray/common/threshold_memory_monitor_factory.cc

@@ -38,46 +38,16 @@ std::unique_ptr<MemoryMonitorInterface> MemoryMonitorFactory::Create(
     return std::make_unique<NoopMemoryMonitor>();
   }
 
-  float usage_threshold = RayConfig::instance().memory_usage_threshold();
-  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 total_memory_bytes = MemoryMonitorUtils::TakeSystemMemorySnapshot(
                                    MemoryMonitorInterface::kDefaultCgroupPath)
                                    .total_bytes;
   memory_usage_threshold_bytes = MemoryMonitorUtils::GetMemoryThreshold(
-      total_memory_bytes, usage_threshold, RayConfig::instance().min_memory_free_bytes());
+      total_memory_bytes,
+      RayConfig::instance().memory_usage_threshold(),
+      RayConfig::instance().min_memory_free_bytes(),
+      resource_isolation_enabled,
+      cgroup_manager);
 
-  if (resource_isolation_enabled) {
-    std::string user_cgroup_path = cgroup_manager.GetUserCgroupPath();
-    StatusOr<std::string> result =
-        cgroup_manager.GetUserCgroupConstraintValue("memory.max");
-    RAY_CHECK(result.ok()) << absl::StrFormat(
-        "Failed to get user cgroup memory limit when setting up memory monitor: %s",
-        result.ToString());
-    std::string user_memory_max_bytes_str = result.value();
-    RAY_CHECK(!user_memory_max_bytes_str.empty()) << absl::StrFormat(
-        "Failed to get memory.max constraint value from user cgroup %s. "
-        "Please check that the cgroup path for resource isolation is correct.",
-        user_cgroup_path);
-
-    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);
-      memory_usage_threshold_bytes =
-          user_memory_max_bytes -
-          static_cast<int64_t>(RayConfig::instance().kill_memory_buffer_bytes());
-      RAY_CHECK_GE(memory_usage_threshold_bytes, 0) << absl::StrFormat(
-          "Available user task memory is less than the kill memory buffer bytes: "
-          "%d < %d. Please consider using a host with more memory. If the current host "
-          "memory size must be kept, please adjust the kill memory buffer size.",
-          user_memory_max_bytes,
-          RayConfig::instance().kill_memory_buffer_bytes());
-    }
-  }
   return std::make_unique<ThresholdMemoryMonitor>(std::move(kill_workers_callback),
                                                   memory_usage_threshold_bytes,
                                                   monitor_interval_ms);