[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:
co-authored by
Li, Jiang <jiang1.li@intel.com>
parent
2725c84aae
commit
b675cb7d0f
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user