[Bugfix][CPU] Honor cgroup memory limit when computing KV cache size (#45086)

Signed-off-by: baoloongmao <baoloongmao@tencent.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
maobaolong
2026-06-14 19:26:50 -07:00
committed by GitHub
co-authored by Li, Jiang <jiang1.li@intel.com>
parent 2725c84aae
commit b675cb7d0f
+52
View File
@@ -50,6 +50,47 @@ class MemoryNodeInfo:
available_memory: int = -1
def _read_int_file(path: str) -> int | None:
try:
with open(path) as f:
value = f.read().strip()
if not value or value == "max":
return None
return int(value)
except (OSError, ValueError):
return None
@cache
def get_cgroup_memory_limit() -> tuple[int | None, int | None]:
"""Return (limit, usage) in bytes from cgroup, or (None, None).
Supports both cgroup v2 (unified) and v1. Returns (None, None) when
not running under a constrained cgroup (e.g. bare metal, or limit
reported as `max`/an unrealistically large value).
"""
if sys.platform != "linux":
return None, None
# cgroup v2 unified hierarchy
v2_limit = _read_int_file("/sys/fs/cgroup/memory.max")
if v2_limit is not None:
v2_usage = _read_int_file("/sys/fs/cgroup/memory.current")
return v2_limit, v2_usage
# cgroup v1
v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
if v1_limit is not None:
# cgroup v1 reports a huge sentinel (close to PAGE_COUNTER_MAX)
# when unlimited. Treat absurdly large values as "no limit".
if v1_limit >= (1 << 62):
return None, None
v1_usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")
return v1_limit, v1_usage
return None, None
def get_memory_affinity(pid: int = 0) -> list[int]:
pid = os.getpid() if pid == 0 else pid
path = f"/proc/{pid}/status"
@@ -114,6 +155,17 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
free_memory + active_file_memory + inactive_file_memory + reclaimable_memory
)
# Honor cgroup memory limit (containers / k8s pods). NUMA meminfo
# reflects host-wide numbers; without this, gpu_memory_utilization
# would be applied to host RAM instead of the pod's limit. cgroup
# does not expose per-NUMA-node limits, so we just clamp the totals
# against the pod-wide limit here.
cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
if cgroup_limit is not None and cgroup_limit < total_memory:
total_memory = cgroup_limit
cgroup_available = cgroup_limit - (cgroup_usage or 0)
available_memory = max(0, min(available_memory, cgroup_available))
return MemoryNodeInfo(
total_memory=total_memory,
available_memory=available_memory,