[diffusion] chore: read the cgroup this process is actually in (#35707)
This commit is contained in:
+67
-11
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user