From 489a69413d2b6ca9311d14ffdc1f455787af5e42 Mon Sep 17 00:00:00 2001 From: raghotham <853234+raghotham@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:53:50 -0700 Subject: [PATCH] [HiCache] Fall back to host memory when no cgroup fs is mounted HiCache sizes its host pool from available_host_memory_bytes(), which bounds psutil's available memory by the process's cgroup headroom. When the process belongs to a cgroup but no cgroup or cgroup2 filesystem is mounted (for example a container that shares the host cgroup namespace without mounting cgroupfs), discovery raised "Cannot locate the process memory cgroup in mounted cgroup filesystems" and the engine failed to start. - If no cgroup filesystem is mounted at all, no limit is visible: log a one-time warning and size against host available memory. - A mounted cgroup filesystem that does not contain the process still raises, since it may hide a real limit; the error now names the override below. - New SGLANG_HICACHE_HOST_MEMORY_BYTES sets the host memory budget directly and skips cgroup and psutil discovery. The existing reserve and per-local-rank split still apply. Values <= 0 are rejected. Tests exercise available_host_memory_bytes() over a synthetic procfs: no cgroupfs (v1 and v2), the explicit budget where discovery would raise, and a mounted but unmatched cgroup. --- python/sglang/srt/environ.py | 4 ++ python/sglang/srt/mem_cache/host_memory.py | 37 +++++++++++++- .../unit/mem_cache/test_host_memory.py | 51 +++++++++++++++++-- 3 files changed, 85 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5a9a98b184ad..e85801ff165d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -763,6 +763,10 @@ class Envs: # HiCache host<->device transfers use the TMA staging kernel when the GPU # (sm_90+), row size and page size allow; set to 0 to force the register kernel. SGLANG_HICACHE_TMA_TRANSFER = EnvBool(True) + # Host memory in bytes available to HiCache pools on this host. When set, + # it replaces cgroup and psutil discovery; the 10 GiB reserve and the + # per-local-rank split still apply. + SGLANG_HICACHE_HOST_MEMORY_BYTES = EnvInt(None) # Base token count for each MLA/DSA dedup broadcast chunk. SGLANG_MLA_DEDUP_CHUNK_TOKENS = EnvInt(2048) SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) diff --git a/python/sglang/srt/mem_cache/host_memory.py b/python/sglang/srt/mem_cache/host_memory.py index 8cd8b4132b89..5032c3e9c161 100644 --- a/python/sglang/srt/mem_cache/host_memory.py +++ b/python/sglang/srt/mem_cache/host_memory.py @@ -2,14 +2,29 @@ from __future__ import annotations +import functools import logging import re from pathlib import Path, PurePosixPath import psutil +from sglang.srt.environ import envs + logger = logging.getLogger(__name__) +_BUDGET_ENV = "SGLANG_HICACHE_HOST_MEMORY_BYTES" + + +@functools.cache +def _warn_no_cgroup_mount() -> None: + logger.warning( + "The process belongs to a cgroup but no cgroup filesystem is mounted; " + "sizing HiCache against host available memory. If the process runs " + "under a memory limit, set %s to its host memory budget.", + _BUDGET_ENV, + ) + def _unescape_mount_path(value: str) -> str: return re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), value) @@ -32,9 +47,11 @@ def _cgroup_memory_headroom(proc_root: Path = Path("/proc")) -> int | None: headroom = None resolved = False + cgroup_mounted = False for line in mounts.splitlines(): before, after = line.split(" - ", 1) filesystem, _, options = after.split()[:3] + cgroup_mounted |= filesystem in ("cgroup", "cgroup2") if filesystem not in memberships: continue if filesystem == "cgroup" and "memory" not in options.split(","): @@ -84,14 +101,30 @@ def _cgroup_memory_headroom(proc_root: Path = Path("/proc")) -> int | None: break directory = directory.parent if memberships and not resolved: + if not cgroup_mounted: + # E.g. a container sharing the host cgroup namespace without + # cgroupfs: no limit is visible, so size against host memory. + _warn_no_cgroup_mount() + return None + # A mounted but unmatched hierarchy may hide a real limit. raise RuntimeError( - "Cannot locate the process memory cgroup in mounted cgroup filesystems" + "Cannot locate the process memory cgroup in mounted cgroup " + f"filesystems; set {_BUDGET_ENV} to the host memory budget to skip " + "cgroup discovery" ) return headroom def available_host_memory_bytes() -> int: - """Conservative allocatable RAM; charged file cache is not assumed reclaimable.""" + """Conservative allocatable RAM; charged file cache is not assumed reclaimable. + + An explicit SGLANG_HICACHE_HOST_MEMORY_BYTES replaces both probes. + """ + budget = envs.SGLANG_HICACHE_HOST_MEMORY_BYTES.get() + if budget is not None: + if budget <= 0: + raise ValueError(f"{_BUDGET_ENV} must be positive, got {budget}") + return budget available = psutil.virtual_memory().available cgroup_headroom = _cgroup_memory_headroom() if cgroup_headroom is not None: diff --git a/test/registered/unit/mem_cache/test_host_memory.py b/test/registered/unit/mem_cache/test_host_memory.py index 7e8a9034f245..e3e4b022f036 100644 --- a/test/registered/unit/mem_cache/test_host_memory.py +++ b/test/registered/unit/mem_cache/test_host_memory.py @@ -1,13 +1,17 @@ """Exercise container and ancestor budgets using synthetic procfs/cgroup files.""" +import functools import tempfile import unittest from pathlib import Path from unittest.mock import Mock, patch +from sglang.srt.environ import envs from sglang.srt.mem_cache import host_memory from sglang.test.ci.ci_register import register_cpu_ci +_cgroup_memory_headroom = host_memory._cgroup_memory_headroom + register_cpu_ci(est_time=1, suite="base-a-test-cpu") @@ -112,11 +116,48 @@ def test_host_availability_is_also_a_bound(self): ): self.assertEqual(host_memory.available_host_memory_bytes(), expected) - def test_unmounted_memory_cgroup_fails(self): - self.configure() - (self.proc / "self/mountinfo").write_text("") - with self.assertRaisesRegex(RuntimeError, "Cannot locate"): - host_memory._cgroup_memory_headroom(self.proc) + def available(self, host_available=5000): + """Public sizing entry point over the synthetic procfs.""" + with ( + patch.object( + host_memory.psutil, + "virtual_memory", + return_value=Mock(available=host_available), + ), + patch.object( + host_memory, + "_cgroup_memory_headroom", + functools.partial(_cgroup_memory_headroom, self.proc), + ), + ): + return host_memory.available_host_memory_bytes() + + def without_cgroupfs(self, v1=False): + # A container sharing the host cgroup namespace, no cgroupfs mounted. + self.configure("/system.slice/engine.scope", v1=v1) + (self.proc / "self/mountinfo").write_text("1 0 0:1 / /proc rw - proc proc rw\n") + + def test_no_cgroup_filesystem_sizes_against_host(self): + for v1 in [False, True]: + with self.subTest(v1=v1): + self.without_cgroupfs(v1=v1) + self.assertEqual(self.available(host_available=5000), 5000) + + def test_explicit_budget_replaces_discovery(self): + # Mounted but unresolvable: discovery alone would raise. + self.configure("/system.slice/engine.scope") + with envs.SGLANG_HICACHE_HOST_MEMORY_BYTES.override(1234): + self.assertEqual(self.available(host_available=5000), 1234) + with envs.SGLANG_HICACHE_HOST_MEMORY_BYTES.override(0): + with self.assertRaisesRegex(ValueError, "must be positive"): + self.available() + + def test_unresolved_mounted_cgroup_fails(self): + self.configure("/system.slice/engine.scope") + with self.assertRaisesRegex( + RuntimeError, "Cannot locate.*SGLANG_HICACHE_HOST_MEMORY_BYTES" + ): + self.available() def test_missing_usage_for_known_limit_fails(self): self.configure()