[core] (cgroups 12/n) Raylet will start worker processes in the application cgroup (#56549)

This PR stacks on #56522 .

For more details about the resource isolation project see
#54703.

This PR the makes the raylet move runtime_env and dashboard agents into
the system cgroup. Workers are now spawned inside the application
cgroup.

It introduces the following:
* I've added a new target `raylet_cgroup_types` which defines the type
used all functions that need to add a process to a cgroup.
* A new parameter is added to `NodeManager`, `WorkerPool`,
`AgentManager`, and `Process` constructors. The parameter is a callback
that will use the CgroupManager to add a process to the respective
cgroup.
* The callback is created in `main.cc`.
* `main.cc` owns CgroupManager because it needs to outlive the
`WorkerPool`.
* `process.c` calls the callback after fork() in the child process so
nothing else can happen in the forked process before it's moved into the
correct cgroup.
* Integration tests in python for end-to-end testing of cgroups with
system and application processes moved into their respective cgroups.
The tests are inside
`python/ray/tests/resource_isolation/test_resource_isolation_integration.py`
and have similar setup/teardown to the C++ integration tests introduced
in #55063.

---------

Signed-off-by: Ibrahim Rabbani <[email protected]>
Co-authored-by: Edward Oakes <[email protected]>

Author: israbbani

python/ray/tests/resource_isolation/test_resource_isolation_integration.py

@@ -289,25 +289,31 @@ Status CgroupManager::Initialize(int64_t system_reserved_cpu_weight,
                                     cpu_weight_constraint_.name_,
                                     std::to_string(application_cgroup_cpu_weight)));
   RegisterRemoveConstraint(application_cgroup_, cpu_weight_constraint_);
-
   return Status::OK();
 }
 
-Status CgroupManager::AddProcessToSystemCgroup(const std::string &pid) {
-  Status s = cgroup_driver_->AddProcessToCgroup(system_leaf_cgroup_, pid);
+Status CgroupManager::AddProcessToCgroup(const std::string &cgroup,
+                                         const std::string &pid) {
+  Status s = cgroup_driver_->AddProcessToCgroup(cgroup, pid);
   // TODO(#54703): Add link to OSS documentation once available.
-  RAY_CHECK(!s.IsNotFound()) << "Failed to move process " << pid << " into system cgroup "
-                             << system_leaf_cgroup_
-                             << " because the cgroup was not found. "
-                                "If resource isolation is enabled, Ray's cgroup "
-                                "hierarchy must not be modified "
-                                "while Ray is running.";
+  RAY_CHECK(!s.IsNotFound())
+      << "Failed to move process " << pid << " into cgroup " << cgroup
+      << " because the cgroup was not found. If resource isolation is enabled, Ray's "
+         "cgroup hierarchy must not be modified while Ray is running.";
   RAY_CHECK(!s.IsPermissionDenied())
-      << "Failed to move process " << pid << " into system cgroup " << system_leaf_cgroup_
+      << "Failed to move process " << pid << " into cgroup " << cgroup
       << " because Ray does not have read, write, and execute "
          "permissions for the cgroup. If resource isolation is enabled, Ray's cgroup "
          "hierarchy must not be modified while Ray is running.";
-
   return s;
 }
+
+Status CgroupManager::AddProcessToApplicationCgroup(const std::string &pid) {
+  return AddProcessToCgroup(application_leaf_cgroup_, pid);
+}
+
+Status CgroupManager::AddProcessToSystemCgroup(const std::string &pid) {
+  return AddProcessToCgroup(system_leaf_cgroup_, pid);
+}
+
 }  // namespace ray

src/ray/common/cgroup2/cgroup_manager.cc

@@ -70,7 +70,7 @@ class CgroupManager : public CgroupManagerInterface {
   CgroupManager &operator=(CgroupManager &&);
 
   /**
-    Moves the process into the system leaf cgroup (@see
+    Moves the process into the application leaf cgroup (@see
     CgroupManagerInterface::kLeafCgroupName).
 
     To move the pid, the process must have read, write, and execute permissions for the
@@ -81,6 +81,25 @@ class CgroupManager : public CgroupManagerInterface {
     TODO(#54703): There currently is not a good way to signal to the caller that
     the method can cause a FATAL error. Revisit this once we've settled on a pattern.
 
+    NOTE: If the process does not have adequate cgroup permissions or the application leaf
+    cgroup does not exist, this will fail a RAY_CHECK.
+
+    @param pid of the process to move into the application leaf cgroup.
+
+    @return Status::OK if pid moved successfully.
+    @return Status::NotFound if the application cgroup does not exist.
+  */
+  Status AddProcessToApplicationCgroup(const std::string &pid) override;
+
+  /**
+    Moves the process into the system leaf cgroup (@see
+    CgroupManagerInterface::kLeafCgroupName).
+
+    To move the pid, the process must have read, write, and execute permissions for the
+      1) the cgroup the pid is currently in i.e. the source cgroup.
+      2) the system leaf cgroup i.e. the destination cgroup.
+      3) the lowest common ancestor of the source and destination cgroups.
+
     NOTE: If the process does not have adequate cgroup permissions or the system leaf
     cgroup does not exist, this will fail a RAY_CHECK.
 
@@ -107,6 +126,24 @@ class CgroupManager : public CgroupManagerInterface {
                 const std::string &node_id,
                 std::unique_ptr<CgroupDriverInterface> cgroup_driver);
 
+  /**
+    Moves the process into the specified cgroup.
+
+    To move the pid, the process must have read, write, and execute permissions for the
+      1) the cgroup the pid is currently in i.e. the source cgroup.
+      2) the system leaf cgroup i.e. the destination cgroup.
+      3) the lowest common ancestor of the source and destination cgroups.
+
+    NOTE: If the process does not have adequate cgroup permissions or the system leaf
+    cgroup does not exist, this will fail a RAY_CHECK.
+
+    @param pid of the process to move into the system leaf cgroup.
+
+    @return Status::OK if pid moved successfully.
+    @return Status::NotFound if the system cgroup does not exist.
+  */
+  Status AddProcessToCgroup(const std::string &cgroup, const std::string &pid);
+
   /**
     Performs the following operations:
 

src/ray/common/cgroup2/cgroup_manager.h

@@ -42,9 +42,26 @@ namespace ray {
 */
 class CgroupManagerInterface {
  public:
-  // TODO(#54703): These will be implemented in a later PR to move processes
-  // into a cgroup.
-  // virtual Status AddProcessToApplicationCgroup(int) = 0;
+  /*
+    Moves the process into the application leaf cgroup (@see kLeafCgroupName).
+
+    To move the pid, the process must have read, write, and execute permissions for the
+      1) the cgroup the pid is currently in i.e. the source cgroup.
+      2) the system leaf cgroup i.e. the destination cgroup.
+      3) the lowest common ancestor of the source and destination cgroups.
+
+    TODO(#54703): There currently is not a good way to signal to the caller that
+    the method can cause a FATAL error. Revisit this once we've settled on a pattern.
+
+    NOTE: If the process does not have adequate cgroup permissions or the application leaf
+    cgroup does not exist, this will fail a RAY_CHECK.
+
+    @param pid of the process to move into the system leaf cgroup.
+
+    @return Status::OK if pid moved successfully.
+    @return Status::NotFound if the application cgroup does not exist.
+  */
+  virtual Status AddProcessToApplicationCgroup(const std::string &pid) = 0;
 
   /**
     Moves the process into the system leaf cgroup (@see kLeafCgroupName).

src/ray/common/cgroup2/cgroup_manager_interface.h

@@ -41,4 +41,8 @@ Status CgroupManager::AddProcessToSystemCgroup(const std::string &pid) {
   return Status::OK();
 }
 
+Status CgroupManager::AddProcessToApplicationCgroup(const std::string &pid) {
+  return Status::OK();
+}
+
 }  // namespace ray

src/ray/common/cgroup2/noop_cgroup_manager.cc

@@ -221,7 +221,6 @@ ray_cc_library(
         "//src/ray/common:flatbuf_utils",
         "//src/ray/common:lease",
         "//src/ray/common:memory_monitor",
-        "//src/ray/common/cgroup2:cgroup_manager_interface",
         "//src/ray/core_worker:experimental_mutable_object_provider",
         "//src/ray/core_worker_rpc_client:core_worker_client_pool",
         "//src/ray/flatbuffers:node_manager_generated",

src/ray/raylet/BUILD.bazel

@@ -17,6 +17,7 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <utility>
 #include <vector>
 
 #include "ray/common/ray_config.h"
@@ -27,7 +28,7 @@
 namespace ray {
 namespace raylet {
 
-void AgentManager::StartAgent() {
+void AgentManager::StartAgent(AddProcessToCgroupHook add_to_cgroup) {
   std::vector<const char *> argv;
   argv.reserve(options_.agent_commands.size());
   for (const std::string &arg : options_.agent_commands) {
@@ -67,7 +68,8 @@ void AgentManager::StartAgent() {
               false,
               env,
               /*pipe_to_stdin*/
-              RayConfig::instance().enable_pipe_based_agent_to_parent_health_check());
+              RayConfig::instance().enable_pipe_based_agent_to_parent_health_check(),
+              std::move(add_to_cgroup));
   if (!process_.IsValid() || ec) {
     // The worker failed to start. This is a fatal error.
     RAY_LOG(FATAL) << "Failed to start agent " << options_.agent_name

src/ray/raylet/agent_manager.cc

@@ -33,6 +33,9 @@ namespace raylet {
 using DelayExecutorFn = std::function<std::shared_ptr<boost::asio::deadline_timer>(
     std::function<void()>, uint32_t)>;
 
+// TODO(#54703): Put this type in a separate target.
+using AddProcessToCgroupHook = std::function<void(const std::string &)>;
+
 // Manages a separate "Agent" process. In constructor (or the `StartAgent` method) it
 // starts a process with `agent_commands` plus some additional arguments.
 //
@@ -59,7 +62,8 @@ class AgentManager {
       Options options,
       DelayExecutorFn delay_executor,
       std::function<void(const rpc::NodeDeathInfo &)> shutdown_raylet_gracefully,
-      bool start_agent = true /* for test */)
+      bool start_agent = true /* for test */,
+      AddProcessToCgroupHook add_to_cgroup = [](const std::string &) {})
       : options_(std::move(options)),
         delay_executor_(std::move(delay_executor)),
         shutdown_raylet_gracefully_(std::move(shutdown_raylet_gracefully)),
@@ -71,13 +75,13 @@ class AgentManager {
       RAY_LOG(FATAL) << "AgentManager agent_commands must not be empty.";
     }
     if (start_agent) {
-      StartAgent();
+      StartAgent(std::move(add_to_cgroup));
     }
   }
   ~AgentManager();
 
  private:
-  void StartAgent();
+  void StartAgent(AddProcessToCgroupHook add_to_cgroup);
 
  private:
   const Options options_;

src/ray/raylet/agent_manager.h

@@ -266,6 +266,10 @@ int main(int argc, char *argv[]) {
   gflags::ShutDownCommandLineFlags();
 
   std::unique_ptr<ray::CgroupManager> cgroup_manager;
+  AddProcessToCgroupHook add_process_to_cgroup_hook = [](const std::string &) {};
+  AddProcessToCgroupHook add_process_to_application_cgroup_hook =
+      [](const std::string &) {};
+  AddProcessToCgroupHook add_process_to_system_cgroup_hook = [](const std::string &) {};
 
   // TODO(#54703): Link OSS documentation once it's available in the error messages.
   if (enable_resource_isolation) {
@@ -321,6 +325,27 @@ int main(int argc, char *argv[]) {
             s.ToString());
       }
     }
+    add_process_to_application_cgroup_hook =
+        [&cgroup_mgr = *cgroup_manager](const std::string &pid) {
+          ray::Status s = cgroup_mgr.AddProcessToApplicationCgroup(pid);
+          if (!s.ok()) {
+            RAY_LOG(WARNING) << absl::StrFormat(
+                "Failed to move process %s into the application cgroup with error %s.",
+                pid,
+                s.ToString());
+          }
+        };
+
+    add_process_to_system_cgroup_hook = [&cgroup_mgr =
+                                             *cgroup_manager](const std::string &pid) {
+      ray::Status s = cgroup_mgr.AddProcessToSystemCgroup(pid);
+      if (!s.ok()) {
+        RAY_LOG(WARNING) << absl::StrFormat(
+            "Failed to move process %s into the system cgroup with error %s.",
+            pid,
+            s.ToString());
+      }
+    };
   }
 
   // Configuration for the node manager.
@@ -647,7 +672,8 @@ int main(int argc, char *argv[]) {
         /*starting_worker_timeout_callback=*/
         [&] { cluster_lease_manager->ScheduleAndGrantLeases(); },
         node_manager_config.ray_debugger_external,
-        /*get_time=*/[]() { return absl::Now(); });
+        /*get_time=*/[]() { return absl::Now(); },
+        std::move(add_process_to_application_cgroup_hook));
 
     client_call_manager = std::make_unique<ray::rpc::ClientCallManager>(
         main_service, /*record_stats=*/true);
@@ -950,7 +976,7 @@ int main(int argc, char *argv[]) {
             std::move(raylet_client_factory),
             /*check_signals=*/nullptr),
         shutdown_raylet_gracefully,
-        std::move(cgroup_manager));
+        std::move(add_process_to_system_cgroup_hook));
 
     // Initialize the node manager.
     raylet = std::make_unique<ray::raylet::Raylet>(main_service,