[Core] (Resource Isolation 14/n) Introduce user slice threshold monitor for resource isolation mode (ray-project#63067)

## Description
This PR introduces the threshold memory monitor portion of ray's new
resource isolation enabled memory monitoring system. Specifically, it
introduces the bottom half of the monitoring system depiction below:
<img width="2168" height="1696" alt="image"
src="https://github.com/user-attachments/assets/b85fc87f-9731-485f-8eac-dfb56ddddf00"
/>



The core of this PR introduces changes to the following parts of the
memory monitoring components:
* `threshold_memory_monitor` where we modify the memory monitor to
monitor the _total user application memory_ (defined by user application
memory defined in the diagram + total object store memory) usage when
resource isolation is enabled.
* `kill worker callback` where we pass in the _total user application
memory_ usage to the killing policy instead of the system wide usage, so
that the killing policy will determine the amount of memory to free
based on the user application usage.
* `memory_monitor_utils` where we add the ability for us to monitor the
total memory usage of the user applications


### Why monitor the _total user application usage_ with a threshold
monitor
The overarching design of the new monitoring system can be found in the
description for ray-project#62705. The TLDR
is that resource isolation aims to provide stronger protection for the
critical ray system processes through isolation the system process
resources from the available user process resources. We accomplish this
through the use of cgroups by placing upper bound resource constraints
on the user application.

However, the cgroup memory upper bound constraint alone is insufficient
for handling cases with heavy object store usage. Because object store
is shared memory that's shared between the raylet (since it does the
allocations) within the system slice and the user workers, it's possible
for portions of object store memory usage to be accounted in either the
system cgroup or the user cgroup.

In the best case scenario, all of object store is accounted within the
user cgroup, and the user cgroup upper bound constraint will be able to
effectively restrict the user slice from using more memory than is
allocated to it. In the worst case scenario, some to all of object store
memory may be accounted within the system cgroup. In this scenario, the
amount of memory that the user is actually using may not be reflected by
the amount of memory that the user cgroup can see as shown in the
diagram below.
<img width="923" height="370" alt="image"
src="https://github.com/user-attachments/assets/04d22466-f541-441c-92b7-876a662b8fc7"
/>



Thus, we bring the threshold monitor to the rescue. The threshold
monitor will be responsible for catching special scenarios like above by
checking that the actual user usage (approximated by the sum of the user
heap memory usage and the shared memory usage of both the system and
user slice) is less than the upper bound we've allocated as available
user memory.


### The road ahead
This PR only introduces the bottom half of the resource isolation
monitoring system shown in the first diagram. The last step for resource
isolation after this PR will be to introduce the top half to complete
our resource isolation memory monitoring system.

---------

Signed-off-by: davik <davik@anyscale.com>
Co-authored-by: davik <davik@anyscale.com>

Author: 2 people

src/ray/common/BUILD.bazel

@@ -100,6 +100,7 @@ ray_cc_library(
         "//src/ray/common/tests:__pkg__",
     ],
     deps = [
+        ":memory_monitor_utils",
         "//src/ray/common/cgroup2:cgroup_test_utils",
         "//src/ray/util:logging",
         "@com_google_googletest//:gtest",

src/ray/common/cgroup2/cgroup_manager.cc

@@ -392,6 +392,8 @@ Status CgroupManager::AddProcessToSystemCgroup(const std::string &pid) {
   return AddProcessToCgroup(system_leaf_cgroup_, pid);
 }
 
+std::string CgroupManager::GetSystemCgroupPath() const { return system_cgroup_; }
+
 std::string CgroupManager::GetUserCgroupPath() const { return user_cgroup_; }
 
 StatusOr<std::string> CgroupManager::GetSystemCgroupConstraintValue(

src/ray/common/cgroup2/cgroup_manager.h

@@ -113,6 +113,11 @@ class CgroupManager : public CgroupManagerInterface {
   */
   Status AddProcessToSystemCgroup(const std::string &pid) override;
 
+  /**
+    @return the path to the system cgroup.
+  */
+  std::string GetSystemCgroupPath() const override;
+
   /**
     @return the path to the user cgroup.
   */

src/ray/common/cgroup2/cgroup_manager_interface.h

@@ -76,6 +76,11 @@ class CgroupManagerInterface {
   */
   virtual Status AddProcessToSystemCgroup(const std::string &pid) = 0;
 
+  /**
+    @return the path to the system cgroup.
+  */
+  virtual std::string GetSystemCgroupPath() const = 0;
+
   /**
     @return the path to the user cgroup.
   */

src/ray/common/cgroup2/linux_cgroup_manager_factory.cc

@@ -61,8 +61,8 @@ std::unique_ptr<CgroupManagerInterface> CgroupManagerFactory::Create(
   int64_t system_memory_bytes_low = system_reserved_memory_bytes;
 
   // Compute user memory limits from proportions
-  SystemMemorySnapshot memory_snapshot =
-      MemoryMonitorUtils::TakeSystemMemorySnapshot(cgroup_path);
+  MemoryUsageSnapshot memory_snapshot =
+      MemoryMonitorUtils::TakeSystemMemoryUsageSnapshot(cgroup_path);
   int64_t total_memory_bytes = memory_snapshot.total_bytes;
   float user_memory_proportion_high = RayConfig::instance().user_memory_proportion_high();
   float user_memory_proportion_max = RayConfig::instance().user_memory_proportion_max();

src/ray/common/cgroup2/noop_cgroup_manager.h

@@ -39,6 +39,8 @@ class NoopCgroupManager : public CgroupManagerInterface {
     return Status::OK();
   }
 
+  std::string GetSystemCgroupPath() const override { return ""; }
+
   std::string GetUserCgroupPath() const override { return ""; }
 
   StatusOr<std::string> GetSystemCgroupConstraintValue(

src/ray/common/memory_monitor_interface.h

@@ -18,6 +18,7 @@
 
 #include <cstdint>
 #include <functional>
+#include <optional>
 #include <ostream>
 
 #include "absl/container/flat_hash_map.h"
@@ -28,20 +29,40 @@ namespace ray {
 /**
  * @brief A snapshot of aggregated memory usage across the system.
  */
-struct SystemMemorySnapshot {
+struct MemoryUsageSnapshot {
+  /// The current memory usage.
   int64_t used_bytes;
 
-  /// The total memory available on the system. >= used_bytes.
+  /// The total memory available on the system.
   int64_t total_bytes;
 
   friend std::ostream &operator<<(std::ostream &os,
-                                  const SystemMemorySnapshot &memory_snapshot) {
+                                  const MemoryUsageSnapshot &memory_snapshot) {
     os << "Used bytes: " << memory_snapshot.used_bytes
        << ", Total bytes: " << memory_snapshot.total_bytes;
     return os;
   }
 };
 
+/**
+ * @brief A snapshot of memory usage within a cgroup.
+ */
+struct CgroupMemorySnapshot {
+  /// size of non-file-backed region mappings within the cgroup in bytes.
+  /// This is an approximation of heap usage for the cgroup.
+  int64_t anon_memory_bytes;
+
+  /// size of shared memory mappings within the cgroup in bytes.
+  int64_t shmem_memory_bytes;
+
+  friend std::ostream &operator<<(std::ostream &os,
+                                  const CgroupMemorySnapshot &memory_snapshot) {
+    os << "Anon memory bytes: " << memory_snapshot.anon_memory_bytes
+       << ", Shmem memory bytes: " << memory_snapshot.shmem_memory_bytes;
+    return os;
+  }
+};
+
 /**
  * @brief A snapshot of per-process memory usage.
  */

src/ray/common/memory_monitor_test_fixture.cc

@@ -14,6 +14,7 @@
 
 #include "ray/common/memory_monitor_test_fixture.h"
 
+#include "ray/common/memory_monitor_utils.h"
 #include "ray/util/logging.h"
 
 namespace ray {
@@ -41,31 +42,82 @@ std::string MemoryMonitorTestFixture::MockProcMemoryUsage(pid_t pid,
   return proc_dir;
 }
 
-std::string MemoryMonitorTestFixture::MockCgroupMemoryUsage(int64_t total_bytes,
-                                                            int64_t current_bytes,
-                                                            int64_t inactive_file_bytes,
-                                                            int64_t active_file_bytes) {
+std::string MemoryMonitorTestFixture::MockCgroupv2MemoryUsage(
+    int64_t total_bytes,
+    int64_t current_bytes,
+    std::optional<int64_t> anon_memory_bytes,
+    std::optional<int64_t> shmem_memory_bytes,
+    int64_t inactive_file_bytes,
+    int64_t active_file_bytes) {
   auto temp_dir_or = TempDirectory::Create();
   RAY_CHECK(temp_dir_or.ok()) << "Failed to create temp directory: "
                               << temp_dir_or.status().message();
   mock_cgroup_dirs_.push_back(std::move(temp_dir_or.value()));
 
   const std::string &cgroup_path = mock_cgroup_dirs_.back()->GetPath();
 
-  mock_cgroup_files_.push_back(std::make_unique<TempFile>(cgroup_path + "/memory.max"));
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV2MemoryMaxPath));
   mock_cgroup_files_.back()->AppendLine(std::to_string(total_bytes) + "\n");
 
-  mock_cgroup_files_.push_back(
-      std::make_unique<TempFile>(cgroup_path + "/memory.current"));
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV2MemoryUsagePath));
   mock_cgroup_files_.back()->AppendLine(std::to_string(current_bytes) + "\n");
 
-  mock_cgroup_files_.push_back(std::make_unique<TempFile>(cgroup_path + "/memory.stat"));
-  mock_cgroup_files_.back()->AppendLine("anon 123456\n");
-  mock_cgroup_files_.back()->AppendLine("inactive_file " +
-                                        std::to_string(inactive_file_bytes) + "\n");
-  mock_cgroup_files_.back()->AppendLine("active_file " +
-                                        std::to_string(active_file_bytes) + "\n");
-  mock_cgroup_files_.back()->AppendLine("some_other_key 789\n");
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV2MemoryStatPath));
+  if (anon_memory_bytes.has_value()) {
+    mock_cgroup_files_.back()->AppendLine(
+        std::string(MemoryMonitorUtils::kCgroupsV2MemoryAnonKey) + " " +
+        std::to_string(*anon_memory_bytes) + "\n");
+  }
+  if (shmem_memory_bytes.has_value()) {
+    mock_cgroup_files_.back()->AppendLine(
+        std::string(MemoryMonitorUtils::kCgroupsV2MemoryShmemKey) + " " +
+        std::to_string(*shmem_memory_bytes) + "\n");
+  }
+  mock_cgroup_files_.back()->AppendLine(
+      std::string(MemoryMonitorUtils::kCgroupsV2MemoryStatInactiveFileKey) + " " +
+      std::to_string(inactive_file_bytes) + "\n");
+  mock_cgroup_files_.back()->AppendLine(
+      std::string(MemoryMonitorUtils::kCgroupsV2MemoryStatActiveFileKey) + " " +
+      std::to_string(active_file_bytes) + "\n");
+
+  return cgroup_path;
+}
+
+std::string MemoryMonitorTestFixture::MockCgroupv1MemoryUsage(int64_t total_bytes,
+                                                              int64_t current_bytes,
+                                                              int64_t inactive_file_bytes,
+                                                              int64_t active_file_bytes) {
+  auto temp_dir_or = TempDirectory::Create();
+  RAY_CHECK(temp_dir_or.ok()) << "Failed to create temp directory: "
+                              << temp_dir_or.status().message();
+  mock_cgroup_dirs_.push_back(std::move(temp_dir_or.value()));
+
+  const std::string &cgroup_path = mock_cgroup_dirs_.back()->GetPath();
+
+  auto memory_dir_or = TempDirectory::Create(cgroup_path + "/memory");
+  RAY_CHECK(memory_dir_or.ok())
+      << "Failed to create temp directory: " << memory_dir_or.status().message();
+  mock_cgroup_dirs_.push_back(std::move(memory_dir_or.value()));
+
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV1MemoryMaxPath));
+  mock_cgroup_files_.back()->AppendLine(std::to_string(total_bytes) + "\n");
+
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV1MemoryUsagePath));
+  mock_cgroup_files_.back()->AppendLine(std::to_string(current_bytes) + "\n");
+
+  mock_cgroup_files_.push_back(std::make_unique<TempFile>(
+      cgroup_path + "/" + MemoryMonitorUtils::kCgroupsV1MemoryStatPath));
+  mock_cgroup_files_.back()->AppendLine(
+      std::string(MemoryMonitorUtils::kCgroupsV1MemoryStatInactiveFileKey) + " " +
+      std::to_string(inactive_file_bytes) + "\n");
+  mock_cgroup_files_.back()->AppendLine(
+      std::string(MemoryMonitorUtils::kCgroupsV1MemoryStatActiveFileKey) + " " +
+      std::to_string(active_file_bytes) + "\n");
 
   return cgroup_path;
 }

src/ray/common/memory_monitor_test_fixture.h

@@ -18,6 +18,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -74,14 +75,33 @@ class MemoryMonitorTestFixture : public ::testing::Test {
    *
    * @param total_bytes The value to write to memory.max (total memory limit).
    * @param current_bytes The value to write to memory.current (current usage).
+   * @param anon_memory_bytes The anon to write to memory.stat (anonymous memory usage).
+   * @param shmem_memory_bytes The shmem to write to memory.stat (shared memory usage).
    * @param inactive_file_bytes The inactive_file value in memory.stat.
    * @param active_file_bytes The active_file value in memory.stat.
    * @return The path to the created mock cgroup directory.
    */
-  std::string MockCgroupMemoryUsage(int64_t total_bytes,
-                                    int64_t current_bytes,
-                                    int64_t inactive_file_bytes,
-                                    int64_t active_file_bytes);
+  std::string MockCgroupv2MemoryUsage(int64_t total_bytes,
+                                      int64_t current_bytes,
+                                      std::optional<int64_t> anon_memory_bytes,
+                                      std::optional<int64_t> shmem_memory_bytes,
+                                      int64_t inactive_file_bytes,
+                                      int64_t active_file_bytes);
+
+  /**
+   * @brief Sets up a mock cgroup v1 directory for emulating memory usage and populates
+   * the files with the provided mock memory values.
+   *
+   * @param total_bytes The value to write to memory/memory.limit_in_bytes.
+   * @param current_bytes The value to write to memory/memory.usage_in_bytes.
+   * @param inactive_file_bytes The total_inactive_file value in memory/memory.stat.
+   * @param active_file_bytes The total_active_file value in memory/memory.stat.
+   * @return The path to the created mock cgroup root directory.
+   */
+  std::string MockCgroupv1MemoryUsage(int64_t total_bytes,
+                                      int64_t current_bytes,
+                                      int64_t inactive_file_bytes,
+                                      int64_t active_file_bytes);
 
  private:
   std::vector<std::unique_ptr<TempDirectory>> mock_proc_dirs_;