ray-project · dentiny · Dec 23, 2024 · Dec 24, 2024 · Dec 24, 2024 · Dec 24, 2024
diff --git a/src/ray/common/cgroup/BUILD b/src/ray/common/cgroup/BUILD
@@ -9,3 +9,20 @@ ray_cc_library(
         "@com_google_absl//absl/strings:str_format",
     ],
 )
+
+ray_cc_library(
+    name = "cgroup_context",
+    hdrs = ["cgroup_context.h"],
+)
+
+ray_cc_library(
+    name = "cgroup_utils",
+    srcs = ["cgroup_utils.cc"],
+    hdrs = ["cgroup_utils.h"],
+    deps = [
+        ":cgroup_context",
+        "//src/ray/util",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings",
+    ],
+)
diff --git a/src/ray/common/cgroup/cgroup_context.h b/src/ray/common/cgroup/cgroup_context.h
@@ -0,0 +1,42 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <unistd.h>
+
+#include <cstdint>
+#include <string>
+
+namespace ray {
+
+// Context used to setup cgroupv2 for a task / actor.
+struct PhysicalModeExecutionContext {
+  // Directory for cgroup, which is applied to application process.
+  //
+  // TODO(hjiang): Revisit if we could save some CPU/mem with string view.
+  std::string cgroup_directory;
+  // A unique id to uniquely identity a certain task / actor attempt.
+  std::string id;
+  // PID for the process.
+  pid_t pid;
+
+  // Memory-related spec.
+  //
+  // Unit: bytes. Corresponds to cgroup V2 `memory.max`, which enforces hard cap on max
+  // memory consumption. 0 means no limit.
+  uint64_t max_memory = 0;
+};
+
+}  // namespace ray
diff --git a/src/ray/common/cgroup/cgroup_utils.cc b/src/ray/common/cgroup/cgroup_utils.cc
@@ -0,0 +1,161 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ray/common/cgroup/cgroup_utils.h"
+
+#ifndef __linux__
+namespace ray {
+bool CgroupV2Setup::SetupCgroupV2ForContext(const PhysicalModeExecutionContext &ctx) {
+  return false;
+}
+/*static*/ bool CgroupV2Setup::CleanupCgroupV2ForContext(
+    const PhysicalModeExecutionContext &ctx) {
+  return false;
+}
+}  // namespace ray
+#else  // __linux__
+
+#include <sys/stat.h>
+
+#include <fstream>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
+#include "ray/util/logging.h"
+
+namespace ray {
+
+namespace {
+
+// Owner can read and write.
+constexpr int kCgroupV2FilePerm = 0600;
+
+// There're two types of memory cgroup constraints:
+// 1. For those with limit capped, they will be created a dedicated cgroup;
+// 2. For those without limit specified, they will be added to the default cgroup.
+static constexpr std::string_view kDefaultCgroupV2Id = "default_cgroup_id";
+
+// Open a cgroup path and append write [content] into the file.
+void OpenCgroupV2FileAndAppend(std::string_view path, std::string_view content) {
+  std::ofstream out_file{path.data(), std::ios::out | std::ios::app};
+  out_file << content;
+}
+
+bool CreateNewCgroupV2(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.id.empty());
+  RAY_CHECK_NE(ctx.id, kDefaultCgroupV2Id);
+  RAY_CHECK_GT(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.id);
+  int ret_code = mkdir(cgroup_folder.data(), kCgroupV2FilePerm);
+  if (ret_code != 0) {
+    return false;
+  }
+
+  const std::string procs_path = absl::StrFormat("%s/cgroup.procs", cgroup_folder);
+  OpenCgroupV2FileAndAppend(procs_path, absl::StrFormat("%d", ctx.pid));
+
+  // Add max memory into cgroup.
+  const std::string max_memory_path = absl::StrFormat("%s/memory.max", cgroup_folder);
+  OpenCgroupV2FileAndAppend(max_memory_path, absl::StrFormat("%d", ctx.max_memory));
+
+  return true;
+}
+
+bool UpdateDefaultCgroupV2(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.id.empty());
+  RAY_CHECK_EQ(ctx.id, kDefaultCgroupV2Id);
+  RAY_CHECK_EQ(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.id);
+  int ret_code = mkdir(cgroup_folder.data(), kCgroupV2FilePerm);
+  if (ret_code != 0) {
+    return false;
+  }
+
+  const std::string procs_path = absl::StrFormat("%s/cgroup.procs", cgroup_folder);
+  OpenCgroupV2FileAndAppend(procs_path, absl::StrFormat("%d", ctx.pid));
+
+  return true;
+}
+
+bool DeleteCgroupV2(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.id.empty());
+  RAY_CHECK_NE(ctx.id, kDefaultCgroupV2Id);
+  RAY_CHECK_GT(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.id);
+  return rmdir(cgroup_folder.data()) == 0;
+}
+
+void PlaceProcessIntoDefaultCgroup(const PhysicalModeExecutionContext &ctx) {
+  const std::string procs_path =
+      absl::StrFormat("%s/%s/cgroup.procs", ctx.cgroup_directory, kDefaultCgroupV2Id);
+  {
+    std::ofstream out_file{procs_path.data(), std::ios::out};
+    out_file << ctx.pid;
+  }
+
+  return;
+}
+
+}  // namespace
+
+/*static*/ std::unique_ptr<CgroupV2Setup> CgroupV2Setup::New(
+    PhysicalModeExecutionContext ctx) {
+  if (!CgroupV2Setup::SetupCgroupV2ForContext(ctx)) {
+    return nullptr;
+  }
+  return std::unique_ptr<CgroupV2Setup>(new CgroupV2Setup(std::move(ctx)));
+}
+
+CgroupV2Setup::~CgroupV2Setup() {
+  if (!CleanupCgroupV2ForContext(ctx_)) {
+    RAY_LOG(ERROR) << "Fails to cleanup cgroup for execution context with id " << ctx_.id;
+  }
+}
+
+/*static*/ bool CgroupV2Setup::SetupCgroupV2ForContext(
+    const PhysicalModeExecutionContext &ctx) {
+  // Create a new cgroup if max memory specified.
+  if (ctx.max_memory > 0) {
+    return CreateNewCgroupV2(ctx);
+  }
+
+  // Update default cgroup if no max resource specified.
+  return UpdateDefaultCgroupV2(ctx);
+}
+
+/*static*/ bool CgroupV2Setup::CleanupCgroupV2ForContext(
+    const PhysicalModeExecutionContext &ctx) {
+  // Delete the dedicated cgroup if max memory specified.
+  if (ctx.max_memory > 0) {
+    PlaceProcessIntoDefaultCgroup(ctx);
+    return DeleteCgroupV2(ctx);
+  }
+
+  // If pid already in default cgroup, no action needed.
+  return true;
+}
+
+}  // namespace ray
+
+#endif  // __linux__
diff --git a/src/ray/common/cgroup/cgroup_utils.h b/src/ray/common/cgroup/cgroup_utils.h
@@ -0,0 +1,61 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Util functions to setup cgroup.
+
+#pragma once
+
+#include <memory>
+#include <string_view>
+#include <utility>
+
+#include "ray/common/cgroup/cgroup_context.h"
+
+namespace ray {
+
+// A util class which sets up cgroup at construction, and cleans up at destruction.
+// On ctor, creates a cgroup v2 if necessary based on the context. Then puts `ctx.pid`
+// into this cgroup.
+// On dtor, puts `ctx.pid` into the default cgroup, and remove this cgroup v2 if any.
+//
+// Precondition:
+// 1. rw permission for cgroup has been validated.
+// 2. Cgroup folder (i.e. default application cgroup folder) has been properly setup.
+// See README under this folder for more details.
+class CgroupV2Setup {
+ public:
+  // A failed construction returns nullptr.
+  static std::unique_ptr<CgroupV2Setup> New(PhysicalModeExecutionContext ctx);
+
+  ~CgroupV2Setup();
+
+  CgroupV2Setup(const CgroupV2Setup &) = delete;
+  CgroupV2Setup &operator=(const CgroupV2Setup &) = delete;
+  CgroupV2Setup(CgroupV2Setup &&) = delete;
+  CgroupV2Setup &operator=(CgroupV2Setup &&) = delete;
+
+ private:
+  CgroupV2Setup(PhysicalModeExecutionContext ctx) : ctx_(std::move(ctx)) {}
+
+  // Setup cgroup based on the given [ctx]. Return whether the setup succeeds or not.
+  static bool SetupCgroupV2ForContext(const PhysicalModeExecutionContext &ctx);
+
+  // Cleanup cgroup based on the given [ctx]. Return whether the cleanup succeds or not.
+  static bool CleanupCgroupV2ForContext(const PhysicalModeExecutionContext &ctx);
+
+  // Execution context for current cgroup v2 setup.
+  PhysicalModeExecutionContext ctx_;
+};
+
+}  // namespace ray
diff --git a/src/ray/raylet/local_task_manager.cc b/src/ray/raylet/local_task_manager.cc
@@ -387,6 +387,16 @@ void LocalTaskManager::DispatchScheduledTasksToWorkers() {
                 const std::shared_ptr<WorkerInterface> worker,
                 PopWorkerStatus status,
                 const std::string &runtime_env_setup_error_message) -> bool {
+              // TODO(hjiang): After getting the ready-to-use worker and task id, we're
+              // able to get physical execution context.
+              //
+              // ownership chain: raylet has-a node manager, node manager has-a local task
+              // manager.
+              //
+              // - PID: could get from available worker
+              // - Attempt id: could pass a global attempt id generator from raylet
+              // - Cgroup application folder: could pass from raylet
+
               return PoppedWorkerHandler(worker,
                                          status,
                                          task_id,
@@ -729,6 +739,8 @@ void LocalTaskManager::RemoveFromRunningTasksIfExists(const RayTask &task) {
   auto sched_cls = task.GetTaskSpecification().GetSchedulingClass();
   auto it = info_by_sched_cls_.find(sched_cls);
   if (it != info_by_sched_cls_.end()) {
+    // TODO(hjiang): After remove the task id from `running_tasks`, corresponding cgroup
+    // will be updated.
     it->second.running_tasks.erase(task.GetTaskSpecification().TaskId());
     if (it->second.running_tasks.size() == 0) {
       info_by_sched_cls_.erase(it);

diff --git a/src/ray/raylet/local_task_manager.h b/src/ray/raylet/local_task_manager.h
@@ -297,6 +297,8 @@ class LocalTaskManager : public ILocalTaskManager {
           capacity(cap),
           next_update_time(std::numeric_limits<int64_t>::max()) {}
     /// Track the running task ids in this scheduling class.
+    ///
+    /// TODO(hjiang): Store cgroup manager along with task id as the value for map.
     absl::flat_hash_set<TaskID> running_tasks;
     /// The total number of tasks that can run from this scheduling class.
     const uint64_t capacity;

diff --git a/src/ray/raylet/main.cc b/src/ray/raylet/main.cc
@@ -200,6 +200,8 @@ int main(int argc, char *argv[]) {
   RAY_LOG(INFO) << "Setting cluster ID to: " << cluster_id;
   gflags::ShutDownCommandLineFlags();
 
+  // TODO(hjiang): Before we do any actual work, setup cgroup.
+
   // Configuration for the node manager.
   ray::raylet::NodeManagerConfig node_manager_config;
   absl::flat_hash_map<std::string, double> static_resource_conf;