fix: count huge pages

python/sglang/srt/mem_cache/cp_shared_l2_pool.py

@@ -19,7 +19,7 @@ import re
 import uuid
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Iterable, Mapping
+from typing import Any, Callable, Iterable, Mapping, Optional
 
 PAYLOAD_TARGET_KV = "target_kv"
 PAYLOAD_DRAFT_KV = "draft_kv"
@@ -1439,28 +1439,80 @@ def _extract_smaps_page_sizes(smaps_text: str) -> list[int]:
     return page_sizes
 
 
-def cp_shared_l2_hugetlbfs_free_bytes(hugetlbfs_dir: str | os.PathLike[str]) -> tuple[int, int, int]:
-    """Free capacity of a hugetlbfs mount via statvfs. Returns (free_bytes, free_pages, page_bytes). On a
-    hugetlbfs mount f_frsize is the hugepage size and f_bavail is the free hugepage count, so this reflects the
-    reserved-but-unused hugepages backing the mount (NOT regular RAM, which psutil reports)."""
-    st = os.statvfs(os.fspath(hugetlbfs_dir))
-    page_bytes = int(st.f_frsize or st.f_bsize)
-    free_pages = int(st.f_bavail)
-    return free_pages * page_bytes, free_pages, page_bytes
+def _hugetlbfs_page_bytes(hugetlbfs_dir: str | os.PathLike[str]) -> int:
+    """The hugepage size of a hugetlbfs mount. statvfs.f_frsize is the mount's page size; fall back to
+    /proc/meminfo Hugepagesize (the default pool) if statvfs returns a non-hugepage value."""
+    try:
+        frsize = int(os.statvfs(os.fspath(hugetlbfs_dir)).f_frsize)
+    except OSError:
+        frsize = 0
+    if frsize >= (1 << 20):  # plausible hugepage size (>= 1 MiB)
+        return frsize
+    try:
+        with open("/proc/meminfo") as f:
+            for line in f:
+                if line.startswith("Hugepagesize:"):
+                    return int(line.split()[1]) * 1024  # "Hugepagesize:  N kB"
+    except OSError:
+        pass
+    return 2 * (1 << 20)  # last-resort default 2M
+
+
+def _hugepage_free_pages(
+    page_bytes: int, *, sysfs_root: str = "/sys/kernel/mm/hugepages"
+) -> Optional[int]:
+    """Free hugepages of the given size from sysfs (NOT statvfs.f_bavail, which is 0 for an unlimited mount
+    with no size= quota). Reads the GLOBAL pool, falling back to summing the per-NUMA-node pools. Returns
+    None if neither is readable (caller then skips the check rather than false-failing)."""
+    name = f"hugepages-{int(page_bytes) // 1024}kB"
+    try:
+        with open(os.path.join(sysfs_root, name, "free_hugepages")) as f:
+            return int(f.read().strip())
+    except OSError:
+        pass
+    import glob
+
+    total, found = 0, False
+    for p in glob.glob(f"/sys/devices/system/node/node*/hugepages/{name}/free_hugepages"):
+        try:
+            with open(p) as f:
+                total += int(f.read().strip())
+                found = True
+        except OSError:
+            pass
+    return total if found else None
+
+
+def cp_shared_l2_hugetlbfs_free_bytes(
+    hugetlbfs_dir: str | os.PathLike[str], *, sysfs_root: str = "/sys/kernel/mm/hugepages"
+) -> tuple[Optional[int], Optional[int], int]:
+    """Free capacity backing a hugetlbfs mount. The page size comes from statvfs (f_frsize); the FREE COUNT
+    comes from sysfs -- statvfs.f_bavail is unreliable (0 for an unlimited mount with no size= quota). Returns
+    (free_bytes, free_pages, page_bytes); free_bytes/free_pages are None when the free count can't be read."""
+    page_bytes = _hugetlbfs_page_bytes(hugetlbfs_dir)
+    free_pages = _hugepage_free_pages(page_bytes, sysfs_root=sysfs_root)
+    free_bytes = free_pages * page_bytes if free_pages is not None else None
+    return free_bytes, free_pages, page_bytes
 
 
 def check_cp_shared_l2_hugetlbfs_capacity(
-    hugetlbfs_dir: str | os.PathLike[str], requested_bytes: int
+    hugetlbfs_dir: str | os.PathLike[str], requested_bytes: int,
+    *, sysfs_root: str = "/sys/kernel/mm/hugepages",
 ) -> None:
-    """Fail LOUD (clean message) if the hugetlbfs mount lacks enough free hugepages for requested_bytes --
-    instead of the cryptic mmap/ftruncate ENOMEM you get otherwise. The shared-L2 slab is backed by hugepages,
-    so this is the CORRECT capacity check (the regular psutil.available check is wrong for this pool: hugepages
-    are carved out of regular RAM, so reserving them drops psutil.available)."""
-    free_bytes, free_pages, page_bytes = cp_shared_l2_hugetlbfs_free_bytes(hugetlbfs_dir)
+    """Fail LOUD (clean message) if the hugepage pool lacks enough free hugepages for requested_bytes --
+    instead of the cryptic mmap/ftruncate ENOMEM. The shared-L2 slab is backed by hugepages, so this is the
+    CORRECT capacity check (the regular psutil.available check is wrong: hugepages are carved out of regular
+    RAM). If the free count can't be determined (sysfs unreadable), the check is SKIPPED (the mmap stays the
+    backstop) -- never false-fails."""
+    free_bytes, free_pages, page_bytes = cp_shared_l2_hugetlbfs_free_bytes(
+        hugetlbfs_dir, sysfs_root=sysfs_root
+    )
+    if free_bytes is None:
+        return
     if requested_bytes > free_bytes:
         need_pages = -(-int(requested_bytes) // page_bytes)
         _failfast(
-            f"insufficient hugepages in {os.fspath(hugetlbfs_dir)!r}: need {requested_bytes / 1e9:.2f} GB "
+            f"insufficient hugepages for {os.fspath(hugetlbfs_dir)!r}: need {requested_bytes / 1e9:.2f} GB "
             f"({need_pages} x {page_bytes // (1024 * 1024)}M pages) but only {free_bytes / 1e9:.2f} GB free "
             f"({free_pages} pages). Reserve more hugepages (echo N | sudo tee "
             f".../hugepages-2048kB/nr_hugepages) or reduce --hicache-size."