From 5bb981dce1a4983341fc27ccca29b6ebe8854dfb Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 21 Aug 2026 08:48:50 +0800 Subject: [PATCH] [diffusion] chore: read the cgroup this process is actually in (#35707) --- .../memory_managers/host_memory_budget.py | 78 ++++++++++++-- .../test/unit/test_host_memory_budget.py | 102 ++++++++++++++---- 2 files changed, 151 insertions(+), 29 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py index 269bbee5e..a3f1d00f6 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py +++ b/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") -_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 diff --git a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py b/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py index 100db4d33..24182821b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py +++ b/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): - 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//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)