[diffusion] chore: read the cgroup this process is actually in (#35707)

This commit is contained in:
Mick
2026-08-21 08:48:50 +08:00
committed by GitHub
parent f825d72936
commit 5bb981dce1
2 changed files with 151 additions and 29 deletions
@@ -13,6 +13,8 @@ container at 1117.2 GiB, a 900 GiB over-report. Serving runs in containers, so
the cap is read directly from whichever cgroup version is mounted.
"""
import os
import psutil
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -21,11 +23,12 @@ logger = init_logger(__name__)
GIB_BYTES = 1024**3
_CGROUP_V2 = ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory.current")
_CGROUP_V1 = (
"/sys/fs/cgroup/memory/memory.limit_in_bytes",
"/sys/fs/cgroup/memory/memory.usage_in_bytes",
# (mount root, limit file, usage file), v2 before v1.
_CGROUP_MOUNTS = (
("/sys/fs/cgroup", "memory.max", "memory.current"),
("/sys/fs/cgroup/memory", "memory.limit_in_bytes", "memory.usage_in_bytes"),
)
_PROC_SELF_CGROUP = "/proc/self/cgroup"
# An unlimited v1 cgroup reports a sentinel near 2**63 rather than omitting the
# file, so treat anything implausibly large as "no cap".
@@ -53,14 +56,67 @@ def _read_int(path: str) -> int | None:
return None
def cgroup_memory_limit_bytes() -> tuple[int, int] | None:
"""This process's (cap, usage) under its cgroup, or None when uncapped."""
for limit_path, usage_path in (_CGROUP_V2, _CGROUP_V1):
limit = _read_int(limit_path)
if limit is None or limit >= _UNLIMITED_ABOVE:
def _own_cgroup_path() -> str:
"""The cgroup path /proc reports for this process, or "" if it reports none."""
try:
with open(_PROC_SELF_CGROUP) as handle:
lines = handle.read().splitlines()
except OSError:
return ""
for line in lines:
fields = line.split(":", 2)
if len(fields) != 3:
continue
usage = _read_int(usage_path) or 0
return limit, usage
# v2 leaves the controller field empty; v1 lists memory among its own
if not fields[1] or "memory" in fields[1].split(","):
return fields[2]
return ""
def _cgroup_dirs(mount: str) -> list[str]:
"""This process's cgroup directory and its ancestors up to `mount`.
The path /proc reports is relative to the host's cgroup root, while the
mount seen inside a container is already the container's own cgroup -- so
the two do not simply concatenate. Measured in a Docker container: /proc
says /docker/8e10..., and the deepest directory that exists under the mount
is the mount itself. Trying progressively shorter suffixes finds the leaf
whichever way the container was set up.
"""
if not os.path.isdir(mount):
return []
parts = [part for part in _own_cgroup_path().split("/") if part]
leaf = mount
for start in range(len(parts)):
candidate = os.path.join(mount, *parts[start:])
if os.path.isdir(candidate):
leaf = candidate
break
dirs = [leaf]
while dirs[-1] != mount:
dirs.append(os.path.dirname(dirs[-1]))
return dirs
def cgroup_memory_limit_bytes() -> tuple[int, int] | None:
"""This process's (cap, usage) under its cgroup, or None when uncapped.
The tightest cap in the chain wins. A nested cgroup -- a systemd scope with
MemoryMax, a container started with --cgroup-parent -- holds this process
below whatever the mount root allows, and planning against the root would
commit memory the process cannot have.
"""
for mount, limit_name, usage_name in _CGROUP_MOUNTS:
tightest = None
for directory in _cgroup_dirs(mount):
limit = _read_int(os.path.join(directory, limit_name))
if limit is None or limit >= _UNLIMITED_ABOVE:
continue
if tightest is not None and limit >= tightest[0]:
continue
tightest = (limit, _read_int(os.path.join(directory, usage_name)) or 0)
if tightest is not None:
return tightest
return None
@@ -13,28 +13,43 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget i
pin_benefit_bytes,
)
_V2_FILES = ("memory.max", "memory.current")
_V1_FILES = ("memory.limit_in_bytes", "memory.usage_in_bytes")
def _point_at(monkeypatch, tmp_path, *, v2=None, v1=None):
"""Redirect the cgroup lookups at files under tmp_path."""
def write(name, value):
path = tmp_path / name
path.write_text(str(value))
return str(path)
def _write_cgroup(directory, files, values):
"""Lay out one cgroup directory's limit and usage files."""
directory.mkdir(parents=True, exist_ok=True)
if values is not None:
for name, value in zip(files, values):
(directory / name).write_text(str(value))
return directory
missing = str(tmp_path / "absent")
v2_paths = (
(write("memory.max", v2[0]), write("memory.current", v2[1]))
if v2
else (missing, missing)
def _point_at(monkeypatch, tmp_path, *, v2=None, v1=None, own_path="", nested=None):
"""Redirect the cgroup lookups at directories under tmp_path.
`own_path` is what /proc would report, and `nested` is the (limit, usage) of
the cgroup at that path -- together they stand in for a process held below
its mount root.
"""
roots = {}
for key, files, values in (("v2", _V2_FILES, v2), ("v1", _V1_FILES, v1)):
roots[key] = _write_cgroup(tmp_path / key, files, values)
if nested is not None:
leaf = roots["v1"]
for part in [part for part in own_path.split("/") if part]:
leaf = leaf / part
_write_cgroup(leaf, _V1_FILES, nested)
monkeypatch.setattr(
host_memory_budget,
"_CGROUP_MOUNTS",
((str(roots["v2"]),) + _V2_FILES, (str(roots["v1"]),) + _V1_FILES),
)
v1_paths = (
(write("limit_in_bytes", v1[0]), write("usage_in_bytes", v1[1]))
if v1
else (missing, missing)
)
monkeypatch.setattr(host_memory_budget, "_CGROUP_V2", v2_paths)
monkeypatch.setattr(host_memory_budget, "_CGROUP_V1", v1_paths)
proc = tmp_path / "proc_self_cgroup"
proc.write_text(f"11:memory:{own_path}\n" if own_path else "")
monkeypatch.setattr(host_memory_budget, "_PROC_SELF_CGROUP", str(proc))
class TestCgroupLimit:
@@ -83,6 +98,57 @@ class TestCgroupLimit:
assert host_memory_available_bytes() == 12 * GIB_BYTES
class TestNestedCgroup:
def test_a_tighter_nested_cap_wins_over_the_root(self, monkeypatch, tmp_path):
# a systemd scope with MemoryMax, or --cgroup-parent: planning against
# the root would commit memory this process cannot have
_point_at(
monkeypatch,
tmp_path,
v1=(1117 * GIB_BYTES, 0),
own_path="/h3",
nested=(32 * GIB_BYTES, 0),
)
assert cgroup_memory_limit_bytes() == (32 * GIB_BYTES, 0)
def test_a_looser_nested_cap_loses_to_the_root(self, monkeypatch, tmp_path):
_point_at(
monkeypatch,
tmp_path,
v1=(32 * GIB_BYTES, 0),
own_path="/wide",
nested=(900 * GIB_BYTES, 0),
)
assert cgroup_memory_limit_bytes() == (32 * GIB_BYTES, 0)
def test_a_container_path_that_does_not_exist_falls_back_to_the_mount(
self, monkeypatch, tmp_path
):
# the measured Docker case: /proc says /docker/8e10..., and the mount is
# already that cgroup, so nothing joins
_point_at(
monkeypatch,
tmp_path,
v1=(1117 * GIB_BYTES, 488 * GIB_BYTES),
own_path="/docker/8e1010720ccd",
)
assert cgroup_memory_limit_bytes() == (1117 * GIB_BYTES, 488 * GIB_BYTES)
def test_a_suffix_of_the_reported_path_is_found_under_the_mount(
self, monkeypatch, tmp_path
):
# /proc says /docker/<id>/h3 while the mount is the container's own
# cgroup, so only the trailing "h3" resolves
_point_at(
monkeypatch,
tmp_path,
v1=(1117 * GIB_BYTES, 0),
own_path="/docker/8e1010720ccd/h3",
)
_write_cgroup(tmp_path / "v1" / "h3", _V1_FILES, (32 * GIB_BYTES, 0))
assert cgroup_memory_limit_bytes() == (32 * GIB_BYTES, 0)
class TestHostPinBudget:
def test_a_component_that_fits_is_granted(self):
budget = HostPinBudget(available_bytes=40 * GIB_BYTES)