[Simulator] Fix meta host memory budgets on constrained runners (#40440)

.github/workflows/_pr-test-simulator-cpu.yml

@@ -88,6 +88,7 @@ jobs:
           MALLOC_ARENA_MAX: "2"
           MAX_JOBS: "1"
           PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1"
+          PYTEST_ADDOPTS: "-s"
           TORCH_EXTENSIONS_DIR: ${{ runner.temp }}/torch-extensions
         run: |
           python3 -m pytest -q tools/sglang-simulator/test/test_simulation_sglang_runner.py

tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_pool_host.py

@@ -104,37 +104,12 @@ def _install_meta_allocators() -> None:
 _SIMULATED_AVAILABLE_HOST_MEMORY_BYTES = 1 << 60
 
 
-class _PsutilProxy:
-    def __init__(self, psutil_module):
-        self._psutil_module = psutil_module
-
-    def virtual_memory(self):
-        snapshot = self._psutil_module.virtual_memory()
-        return snapshot._replace(
-            available=max(
-                snapshot.available,
-                _SIMULATED_AVAILABLE_HOST_MEMORY_BYTES,
-            )
-        )
-
-    def __getattr__(self, name):
-        return getattr(self._psutil_module, name)
-
-
 def _call_with_meta_host_memory(original_init, self, *args, **kwargs):
     """Bypass physical host-payload checks while meta allocation is active."""
-    init_globals = getattr(original_init, "__globals__", None)
+    from sglang.srt.mem_cache.pool_host.base import host_memory_budget_scope
-    psutil_module = init_globals.get("psutil") if init_globals is not None else None
-    if psutil_module is None:
-        return original_init(self, *args, **kwargs)
 
-    proxy = _PsutilProxy(psutil_module)
+    with host_memory_budget_scope(_SIMULATED_AVAILABLE_HOST_MEMORY_BYTES):
-    init_globals["psutil"] = proxy
-    try:
         return original_init(self, *args, **kwargs)
-    finally:
-        if init_globals.get("psutil") is proxy:
-            init_globals["psutil"] = psutil_module
 
 
 @lru_cache(maxsize=256)