setup cgroup v2 in C++

Signed-off-by: dentiny <dentinyhao@gmail.com>
ray-project · dentiny · Dec 23, 2024 · Dec 24, 2024 · Dec 24, 2024 · Dec 24, 2024
commit a9e8dc627537ce00c734cc7d82c4e66f815ae6aa
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,41 @@
+// 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>
+#include <string_view>
+
+namespace ray {
+
+// Context used to setup cgroupv2 for a task / actor.
+struct PhysicalModeExecutionContext {
+  // Directory for cgroup, which is appled to application process.
+  std::string_view cgroup_directory;
+  // UUID to indicate the current task / actor.
+  std::string uuid;
+  // 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,163 @@
+// 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"
+
+#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 kCgroupFilePerm = 0600;
+
+// Open a cgroup path and append write [content] into the file.
+void OpenCgroupFileAndAppend(std::string_view path, std::string content) {
+  std::ofstream out_file{path.data(), std::ios::out | std::ios::app};
+  out_file << content;
+}
+
+bool CreateNewCgroup(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.uuid.empty());
+  RAY_CHECK_NE(ctx.uuid, kDefaultCgroupUuid);
+  RAY_CHECK_GT(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.uuid);
+  int ret_code = mkdir(cgroup_folder.data(), kCgroupFilePerm);
+  if (ret_code != 0) {
+    return false;
+  }
+
+  if (ctx.max_memory > 0) {
+    const std::string procs_path = absl::StrFormat("%s/cgroup.procs", cgroup_folder);
+    OpenCgroupFileAndAppend(procs_path, absl::StrFormat("%d", ctx.pid));
+
+    const std::string max_memory_path = absl::StrFormat("%s/memory.max", cgroup_folder);
+    OpenCgroupFileAndAppend(max_memory_path, absl::StrFormat("%d", ctx.max_memory));
+  }
+
+  return true;
+}
+
+bool UpdateDefaultCgroup(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.uuid.empty());
+  RAY_CHECK_EQ(ctx.uuid, kDefaultCgroupUuid);
+  RAY_CHECK_EQ(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.uuid);
+  int ret_code = mkdir(cgroup_folder.data(), kCgroupFilePerm);
+  if (ret_code != 0) {
+    return false;
+  }
+
+  const std::string procs_path = absl::StrFormat("%s/cgroup.procs", cgroup_folder);
+  OpenCgroupFileAndAppend(procs_path, absl::StrFormat("%d", ctx.pid));
+
+  return true;
+}
+
+bool DeleteCgroup(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.uuid.empty());
+  RAY_CHECK_NE(ctx.uuid, kDefaultCgroupUuid);
+  RAY_CHECK_GT(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.uuid);
+  return rmdir(cgroup_folder.data()) == 0;
+}
+
+bool RemoveCtxFromDefaultCgroup(const PhysicalModeExecutionContext &ctx) {
+  // Sanity check.
+  RAY_CHECK(!ctx.uuid.empty());
+  RAY_CHECK_EQ(ctx.uuid, kDefaultCgroupUuid);
+  RAY_CHECK_EQ(ctx.max_memory, 0);
+
+  const std::string cgroup_folder =
+      absl::StrFormat("%s/%s", ctx.cgroup_directory, ctx.uuid);
+  int ret_code = mkdir(cgroup_folder.data(), kCgroupFilePerm);
+  if (ret_code != 0) {
+    return false;
+  }
+
+  const std::string procs_path = absl::StrFormat("%s/cgroup.procs", cgroup_folder);
+  std::ostringstream buffer;
+  {
+    std::ifstream file(procs_path.data(), std::ios::in);
+    buffer << file.rdbuf();
+  }
+  std::string content = buffer.str();  // contains all PIDs, separated by space
+
+  std::vector<std::string_view> old_pid_strings = absl::StrSplit(content, ' ');
+  std::vector<std::string_view> new_pid_strings;
+  new_pid_strings.reserve(old_pid_strings.size() - 1);
+  for (const auto &cur_pid : old_pid_strings) {
+    if (cur_pid == absl::StrFormat("%d", ctx.pid)) {
+      continue;
+    }
+    new_pid_strings.emplace_back(cur_pid);
+  }
+
+  const std::string new_pids = absl::StrJoin(new_pid_strings, " ");
+  {
+    std::ofstream out_file{procs_path.data(), std::ios::out};
+    out_file << new_pids;
+  }
+
+  return true;
+}
+
+}  // namespace
+
+bool SetupCgroupForContext(const PhysicalModeExecutionContext &ctx) {
+#ifndef __linux__
+  return false;
+#endif
+
+  // Create a new cgroup if max memory specified.
+  if (ctx.max_memory > 0) {
+    return CreateNewCgroup(ctx);
+  }
+
+  // Update default cgroup if no max resource specified.
+  return UpdateDefaultCgroup(ctx);
+}
+
+bool CleanupCgroupForContext(const PhysicalModeExecutionContext &ctx) {
+#ifndef __linux__
+  return false;
+#endif
+
+  // Delete the dedicated cgroup if max memory specified.
+  if (ctx.max_memory > 0) {
+    return DeleteCgroup(ctx);
+  }
+
+  // Update default cgroup if no max resource specified.
+  return RemoveCtxFromDefaultCgroup(ctx);
+}
+
+}  // namespace ray
diff --git a/src/ray/common/cgroup/cgroup_utils.h b/src/ray/common/cgroup/cgroup_utils.h
@@ -0,0 +1,36 @@
+// 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 <string_view>
+
+#include "ray/common/cgroup/cgroup_context.h"
+
+namespace ray {
+
+// 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.
+inline constexpr std::string_view kDefaultCgroupUuid = "default_cgroup_uuid";
+
+// Setup cgroup based on the given [ctx]. Return whether the setup succeeds or not.
+bool SetupCgroupForContext(const PhysicalModeExecutionContext &ctx);
+
+// Cleanup cgroup based on the given [ctx]. Return whether the cleanup succeds or not.
+bool CleanupCgroupForContext(const PhysicalModeExecutionContext &ctx);
+
+}  // namespace ray