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

python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py

@@ -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")
+# (mount root, limit file, usage file), v2 before v1.
-_CGROUP_V1 = (
+_CGROUP_MOUNTS = (
-    "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+    ("/sys/fs/cgroup", "memory.max", "memory.current"),
-    "/sys/fs/cgroup/memory/memory.usage_in_bytes",
+    ("/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 _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
+        # 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."""
+    """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)
+    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
-        usage = _read_int(usage_path) or 0
+            if tightest is not None and limit >= tightest[0]:
-        return limit, usage
+                continue
+            tightest = (limit, _read_int(os.path.join(directory, usage_name)) or 0)
+        if tightest is not None:
+            return tightest
     return None
 
 

python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py

@@ -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):
+def _write_cgroup(directory, files, values):
-        path = tmp_path / name
+    """Lay out one cgroup directory's limit and usage files."""
-        path.write_text(str(value))
+    directory.mkdir(parents=True, exist_ok=True)
-        return str(path)
+    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 = (
+def _point_at(monkeypatch, tmp_path, *, v2=None, v1=None, own_path="", nested=None):
-        (write("memory.max", v2[0]), write("memory.current", v2[1]))
+    """Redirect the cgroup lookups at directories under tmp_path.
-        if v2
+
-        else (missing, missing)
+    `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 = (
+    proc = tmp_path / "proc_self_cgroup"
-        (write("limit_in_bytes", v1[0]), write("usage_in_bytes", v1[1]))
+    proc.write_text(f"11:memory:{own_path}\n" if own_path else "")
-        if v1
+    monkeypatch.setattr(host_memory_budget, "_PROC_SELF_CGROUP", str(proc))
-        else (missing, missing)
-    )
-    monkeypatch.setattr(host_memory_budget, "_CGROUP_V2", v2_paths)
-    monkeypatch.setattr(host_memory_budget, "_CGROUP_V1", v1_paths)
 
 
 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)