fix: count huge pages

2026-06-23 10:29:52 +00:00
parent 6114c3e3e5
commit 0025b0e819
2 changed files with 109 additions and 27 deletions
--- a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py
+++ b/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."
--- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py
+++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py
@@ -2838,23 +2838,53 @@ class TestCpSharedL2EightRankReserveDeterminism(unittest.TestCase):

 class TestCpSharedL2HugetlbfsCapacity(unittest.TestCase):
    """The hugetlbfs capacity check fails LOUD on insufficient hugepages (clean message vs a cryptic mmap
-    ENOMEM). statvfs works on any mount, so the logic is exercised here against a tmpdir; the real hugetlbfs
-    semantics (f_frsize = hugepage size, f_bavail = free hugepages) validate on the live server."""
+    ENOMEM). The FREE COUNT comes from sysfs (`hugepages-NkB/free_hugepages`), NOT statvfs.f_bavail -- the
+    latter is 0 for an unlimited mount with no size= quota (the false "0 GB free" the live server hit). Here we
+    inject a fake sysfs pool via `sysfs_root=`; the real sysfs validates on the server."""

    def test_free_bytes_and_capacity_check(self):
        import os
        import tempfile

-        free_bytes_fn = _cp_shared_l2_pool.cp_shared_l2_hugetlbfs_free_bytes
-        check_fn = _cp_shared_l2_pool.check_cp_shared_l2_hugetlbfs_capacity
-        with tempfile.TemporaryDirectory() as d:
-            free, pages, page_bytes = free_bytes_fn(d)
-            self.assertGreater(free, 0)
-            self.assertGreater(page_bytes, 0)
-            self.assertEqual(free, pages * page_bytes)
-            check_fn(d, 4096)  # tiny request fits -> no raise
+        pool = _cp_shared_l2_pool
+        free_bytes_fn = pool.cp_shared_l2_hugetlbfs_free_bytes
+        check_fn = pool.check_cp_shared_l2_hugetlbfs_capacity
+        with tempfile.TemporaryDirectory() as d, tempfile.TemporaryDirectory() as sysfs:
+            # Build a fake sysfs pool whose dir name matches this mount's hugepage size so the lookup resolves
+            # (the page size still comes from statvfs/meminfo; only the free COUNT is read from sysfs).
+            page_bytes = pool._hugetlbfs_page_bytes(d)
+            name = f"hugepages-{page_bytes // 1024}kB"
+            os.makedirs(os.path.join(sysfs, name))
+            with open(os.path.join(sysfs, name, "free_hugepages"), "w") as f:
+                f.write("100\n")
+
+            free, pages, pb = free_bytes_fn(d, sysfs_root=sysfs)
+            self.assertEqual(pages, 100)
+            self.assertEqual(pb, page_bytes)
+            self.assertEqual(free, 100 * page_bytes)
+
+            check_fn(d, page_bytes, sysfs_root=sysfs)  # 1 page fits in 100 -> no raise
            with self.assertRaises(ValueError):  # exceeds free capacity -> fail loud
-                check_fn(d, free + 10**12)
+                check_fn(d, free + page_bytes, sysfs_root=sysfs)
+
+    def test_skips_check_when_free_count_unreadable(self):
+        # When the free count can't be read (unreadable sysfs root + a hugepage size that won't match any real
+        # per-NUMA pool), the helper returns None and the check is SKIPPED -- never false-fails (the bug the
+        # live server hit, where statvfs.f_bavail==0 wrongly reported "0 GB free").
+        import os
+        import tempfile
+
+        pool = _cp_shared_l2_pool
+        with tempfile.TemporaryDirectory() as empty_sysfs, tempfile.TemporaryDirectory() as d:
+            # 7 MiB is not a real hugepage size, so the per-NUMA glob fallback finds nothing either.
+            self.assertIsNone(pool._hugepage_free_pages(7 << 20, sysfs_root=empty_sysfs))
+            free, pages, pb = pool.cp_shared_l2_hugetlbfs_free_bytes(d, sysfs_root=empty_sysfs)
+            # Page size is still resolvable; free count is None iff this mount's size isn't in the empty root.
+            self.assertGreater(pb, 0)
+            if free is None:
+                self.assertIsNone(pages)
+                # An impossibly large request must NOT raise when the free count is unknown.
+                pool.check_cp_shared_l2_hugetlbfs_capacity(d, 10**18, sysfs_root=empty_sysfs)


 if __name__ == "__main__":