[HiCache] Auto-size the host pool to fit available host memory (#40135)

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -10,6 +10,8 @@ from sglang.test.test_utils import CustomTestCase, enter_override, maybe_stub_sg
 
 maybe_stub_sgl_kernel()  # must precede any import that pulls in sgl_kernel
 
+import asyncio
+import gc
 import json
 import re
 import tempfile
@@ -277,6 +279,13 @@ class ServingChatTestCase(unittest.TestCase):
         # to publish one rather than hang the values off a mock manager.
         reset_context()
         self.addCleanup(reset_context)
+        # Tests drive coroutines through get_or_create_event_loop(), which
+        # creates a fresh loop per call and leaves the previous one unclosed.
+        # Finalize those loops here, between tests: if the cyclic GC collects
+        # one mid-import, its ResourceWarning imports tracemalloc while the
+        # outer import still holds the module-lock bookkeeping, which raises
+        # KeyError from importlib._bootstrap on Python < 3.12.
+        self.addCleanup(self._close_event_loops)
         publish(
             ServerArgs(
                 model_path="dummy",
@@ -314,6 +323,17 @@ class ServingChatTestCase(unittest.TestCase):
         self.fastapi_request = Mock(spec=Request)
         self.fastapi_request.headers = {}
 
+    @staticmethod
+    def _close_event_loops():
+        try:
+            loop = asyncio.get_event_loop_policy().get_event_loop()
+        except RuntimeError:
+            loop = None
+        if loop is not None and not loop.is_closed():
+            loop.close()
+        asyncio.set_event_loop(None)
+        gc.collect()
+
     @staticmethod
     def _render_tool_results_in_call_order(messages, **kwargs):
         """Block-level tool_call_id association, like the GLM chat templates."""

test/registered/unit/mem_cache/test_hicache_auto_size.py

@@ -0,0 +1,104 @@
+import unittest
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.arg_groups.overrides import resolution_result
+from sglang.srt.mem_cache import hicache_auto_size as sizing
+from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.mem_cache.pool_host import base
+from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.runtime_context import get_context, get_memory
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestHiCacheAutoSize(CustomTestCase):
+    def test_hybrid_target_draft_uses_its_sidecar_slot_capacity(self):
+        # EAGLE3 can pair a hybrid target with a plain MHA draft. Full drafts
+        # follow full target slots; SWA drafts follow the smaller SWA capacity.
+        target = Mock(
+            spec=SWAKVPool,
+            size=128,
+            full_kv_pool=Mock(size=128, get_kv_size_bytes=Mock(return_value=4096)),
+            swa_kv_pool=Mock(size=32, get_kv_size_bytes=Mock(return_value=1024)),
+        )
+        draft_mha = Mock(
+            spec=MHATokenToKVPool,
+            size=16,
+            get_kv_size_bytes=Mock(return_value=(128, 128)),
+        )
+        params = CacheInitParams(
+            disable=False,
+            req_to_token_pool=None,
+            token_to_kv_pool_allocator=Mock(get_kvcache=Mock(return_value=target)),
+            page_size=2,
+        )
+        for draft, expected_sidecar_bytes in (
+            (draft_mha, 2048),
+            (Mock(spec=BaseSWAKVPool, swa_kv_pool=draft_mha), 512),
+        ):
+            with self.subTest(
+                draft_type=type(draft).__name__, bytes=expected_sidecar_bytes
+            ):
+                plan = Mock(mode="sidecar", device_pools=(draft,))
+                self.assertEqual(
+                    sizing._estimate_hicache_bytes(params, plan),
+                    4096 + 1024 + expected_sidecar_bytes,
+                )
+
+    def test_default_ratio_fits_host_budget_and_pools_book_it(self):
+        """With only --enable-hierarchical-cache the default ratio shrinks to the
+        per-rank budget, pools book one snapshot, and an explicit ratio opts out."""
+        pool = MHATokenToKVPool(
+            size=128,
+            page_size=2,
+            dtype=torch.float16,
+            head_num=2,
+            head_dim=4,
+            layer_num=2,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+        params = CacheInitParams(
+            disable=False,
+            req_to_token_pool=None,
+            token_to_kv_pool_allocator=Mock(get_kvcache=Mock(return_value=pool)),
+            page_size=2,
+        )
+        rank_budget = 10_000
+        # Four ranks per host (e.g. TP8 over two 4-GPU nodes) share what psutil reports.
+        host_free = base.HICACHE_HOST_MEMORY_RESERVE_BYTES + 4 * rank_budget
+        with (
+            get_context().override_server_args(enable_hierarchical_cache=True),
+            patch.object(base, "ranks_per_host", return_value=4),
+            patch.object(base, "available_host_memory_bytes", return_value=host_free),
+            sizing.auto_size_hicache(params, None, enabled=True),
+        ):
+            ratio = get_memory().hicache_ratio
+            self.assertLess(ratio, 2.0)
+            host = MHATokenToKVPoolHost(
+                pool, ratio, 0, 2, "layer_first", pin_memory=False, device="cpu"
+            )
+            self.assertLessEqual(host.size * host.size_per_token, 0.8 * rank_budget)
+            with self.assertRaisesRegex(ValueError, "Not enough host memory"):
+                MHATokenToKVPoolHost(
+                    pool, ratio, 0, 2, "layer_first", pin_memory=False, device="cpu"
+                )
+        self.assertIsNone(base._host_memory_budget.get())
+
+        explicit = ServerArgs(model_path="dummy", hicache_ratio=2.0)
+        explicit.resolve_once()
+        self.assertIsNone(
+            resolution_result(explicit, "hicache_host_memory_fraction", 0.8)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_host_memory.py

@@ -0,0 +1,130 @@
+"""Exercise container and ancestor budgets using synthetic procfs/cgroup files."""
+
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+from sglang.srt.mem_cache import host_memory
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestHostMemory(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp.cleanup)
+        self.root = Path(self.temp.name)
+        self.proc = self.root / "proc"
+        (self.proc / "self").mkdir(parents=True)
+        self.mount = self.root / "cgroup mount"
+        self.mount.mkdir()
+
+    def configure(self, membership="/task/engine", mount_root="/", v1=False):
+        controllers = "memory" if v1 else ""
+        (self.proc / "self/cgroup").write_text(f"0:{controllers}:{membership}\n")
+        filesystem = "cgroup" if v1 else "cgroup2"
+        options = "rw,memory" if v1 else "rw"
+        escaped = str(self.mount).replace(" ", r"\040")
+        (self.proc / "self/mountinfo").write_text(
+            f"1 0 0:1 {mount_root} {escaped} rw - {filesystem} cgroup {options}\n"
+        )
+
+    def memory(self, path, usage, maximum="max", high="max", v1=False):
+        directory = self.mount / path
+        directory.mkdir(parents=True, exist_ok=True)
+        files = (
+            {"memory.limit_in_bytes": maximum, "memory.usage_in_bytes": usage}
+            if v1
+            else {"memory.max": maximum, "memory.high": high, "memory.current": usage}
+        )
+        for name, value in files.items():
+            (directory / name).write_text(str(value))
+
+    def test_v2_parent_and_high_limits(self):
+        self.configure()
+        self.memory("task/engine", 100)
+        for maximum, high, usage, expected in [
+            (1000, "max", 300, 700),
+            (1000, 600, 300, 300),
+            ("max", 600, 700, 0),
+            ("max", "max", 300, None),
+        ]:
+            with self.subTest(maximum=maximum, high=high):
+                self.memory("task", usage, maximum, high)
+                self.assertEqual(
+                    host_memory._cgroup_memory_headroom(self.proc), expected
+                )
+        self.memory("task", 100, 1000)
+        self.memory("task/engine", 100, 250)
+        self.assertEqual(host_memory._cgroup_memory_headroom(self.proc), 150)
+
+    def test_mount_subtree_and_cgroup_namespace(self):
+        self.memory("engine", 100, 900)
+        self.memory("", 300, 1000)
+        for membership in ["/host/task/engine", "/engine"]:
+            with self.subTest(membership=membership):
+                self.configure(membership, mount_root="/host/task")
+                self.assertEqual(host_memory._cgroup_memory_headroom(self.proc), 700)
+
+    def test_v1_parent_limit_and_unlimited_sentinel(self):
+        self.configure(v1=True)
+        self.memory("task/engine", 100, 2**63 - 4096, v1=True)
+        self.memory("task", 400, 1000, v1=True)
+        self.assertEqual(host_memory._cgroup_memory_headroom(self.proc), 600)
+
+    def test_independent_engines_have_separate_allowances(self):
+        # Both engines see the same host RAM but have different charged usage.
+        for task, usage, expected in [("a", 300, 700), ("b", 600, 400)]:
+            with self.subTest(task=task):
+                self.configure(f"/{task}/engine")
+                self.memory(f"{task}/engine", 0)
+                self.memory(task, usage, 1000)
+                with (
+                    patch.object(
+                        host_memory.psutil,
+                        "virtual_memory",
+                        return_value=Mock(available=2000),
+                    ),
+                    patch.object(
+                        host_memory,
+                        "_cgroup_memory_headroom",
+                        return_value=host_memory._cgroup_memory_headroom(self.proc),
+                    ),
+                ):
+                    self.assertEqual(
+                        host_memory.available_host_memory_bytes(), expected
+                    )
+
+    def test_host_availability_is_also_a_bound(self):
+        for cgroup, expected in [(None, 100), (200, 100), (50, 50)]:
+            with (
+                self.subTest(cgroup=cgroup),
+                patch.object(
+                    host_memory.psutil,
+                    "virtual_memory",
+                    return_value=Mock(available=100),
+                ),
+                patch.object(
+                    host_memory, "_cgroup_memory_headroom", return_value=cgroup
+                ),
+            ):
+                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 test_missing_usage_for_known_limit_fails(self):
+        self.configure()
+        self.memory("task/engine", 100, 1000)
+        (self.mount / "task/engine/memory.current").unlink()
+        with self.assertRaises(FileNotFoundError):
+            host_memory._cgroup_memory_headroom(self.proc)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -282,11 +282,10 @@ class TestHostMemoryBudget(CustomTestCase):
     def _budget_with_ranks(self, ranks):
         # Deliberate single-accessor stub: isolates the budget math from the
         # topology derivation, which the ranks_per_host case below covers.
-        fake_mem = unittest.mock.Mock(available=self._AVAILABLE)
         with (
             unittest.mock.patch.object(base, "ranks_per_host", return_value=ranks),
             unittest.mock.patch.object(
-                base.psutil, "virtual_memory", return_value=fake_mem
+                base, "available_host_memory_bytes", return_value=self._AVAILABLE
             ),
         ):
             return base.host_memory_budget_bytes()