fix(gms): support SGLang 0.5.16 memory pool API

lib/gpu_memory_service/integrations/sglang/patches.py

@@ -10,7 +10,6 @@
 
 from __future__ import annotations
 
-import inspect
 import logging
 from contextlib import contextmanager
 from typing import Optional
@@ -163,84 +162,33 @@ def patch_model_runner() -> None:
     if hasattr(ModelRunner, "_gms_patched"):
         return
 
-    original_init_memory_pool = ModelRunner.init_memory_pool
-    memory_arg_name = next(
-        (
-            name
-            for name in inspect.signature(original_init_memory_pool).parameters
-            if name != "self"
-        ),
-        None,
-    )
+    original_alloc_memory_pool = ModelRunner.alloc_memory_pool
 
-    def patched_init_memory_pool(self, *args, **kwargs):
-        """Patch memory baseline for SGLang old/new init_memory_pool signatures."""
+    def patched_alloc_memory_pool(self, *args, **kwargs):
         impl = get_gms_memory_saver_impl()
-        preloaded_weights_gib = 0.0
-        if impl is not None:
+        if (
+            impl is not None
+            and impl.preloaded_weights_bytes > 0
+            and not self.__dict__.get("_gms_memory_baseline_adjusted", False)
+        ):
             preloaded_weights_gib = impl.preloaded_weights_bytes / (1 << 30)
+            old_value = self.pre_model_load_memory
+            self.pre_model_load_memory += preloaded_weights_gib
+            self._gms_memory_baseline_adjusted = True
+            logger.info(
+                "[GMS] Adjusted pre_model_load_memory for preloaded weights: "
+                "%.2f GiB + %.2f GiB = %.2f GiB",
+                old_value,
+                preloaded_weights_gib,
+                self.pre_model_load_memory,
+            )
 
-        if preloaded_weights_gib > 0 and memory_arg_name in (
-            "pre_model_load_memory",
-            "total_gpu_memory",
-        ):
-            if args:
-                old_value = args[0]
-                new_value = (
-                    old_value + preloaded_weights_gib
-                    if isinstance(old_value, (int, float))
-                    else old_value
-                )
-                args = (new_value,) + args[1:]
-            elif memory_arg_name in kwargs:
-                old_value = kwargs[memory_arg_name]
-                new_value = (
-                    old_value + preloaded_weights_gib
-                    if isinstance(old_value, (int, float))
-                    else old_value
-                )
-                kwargs = dict(kwargs)
-                kwargs[memory_arg_name] = new_value
-            else:
-                old_value = None
-                new_value = None
-
-            if isinstance(old_value, (int, float)) and isinstance(
-                new_value, (int, float)
-            ):
-                logger.info(
-                    "[GMS] Adjusted %s for preloaded weights: "
-                    "%.2f GiB + %.2f GiB = %.2f GiB",
-                    memory_arg_name,
-                    old_value,
-                    preloaded_weights_gib,
-                    new_value,
-                )
-            else:
-                logger.info(
-                    "[GMS] Could not adjust %s for preloaded weights; value=%r",
-                    memory_arg_name,
-                    old_value,
-                )
-        elif impl is not None and impl.imported_weights_bytes > 0:
-            if preloaded_weights_gib > 0:
-                logger.info(
-                    "[GMS] Leaving %s unchanged; unsupported SGLang "
-                    "init_memory_pool signature for preloaded weights",
-                    memory_arg_name,
-                )
-            else:
-                logger.info(
-                    "[GMS] Leaving %s unchanged; weights were loaded by this process",
-                    memory_arg_name,
-                )
-
-        return original_init_memory_pool(self, *args, **kwargs)
-
-    ModelRunner.init_memory_pool = patched_init_memory_pool
+        return original_alloc_memory_pool(self, *args, **kwargs)
+
+    ModelRunner.alloc_memory_pool = patched_alloc_memory_pool
     ModelRunner._gms_patched = True
     _model_runner_patched = True
-    logger.info("[GMS] Patched ModelRunner.init_memory_pool")
+    logger.info("[GMS] Patched ModelRunner.alloc_memory_pool")
 
 
 def patch_static_state_for_gms() -> None:

lib/gpu_memory_service/tests/test_sglang_patches.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""CPU-only tests for SGLang ModelRunner memory accounting patches."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+from _deps import HAS_GMS, HAS_TORCH
+
+if not HAS_GMS:
+    pytest.skip(
+        "gpu_memory_service package is not available in this test image",
+        allow_module_level=True,
+    )
+
+if not HAS_TORCH:
+    pytest.skip("torch is required", allow_module_level=True)
+
+from gpu_memory_service.integrations.sglang import patches
+
+pytestmark = [
+    pytest.mark.pre_merge,
+    pytest.mark.unit,
+    pytest.mark.sglang,
+    pytest.mark.core,
+    pytest.mark.gpu_0,
+]
+
+
+def _patch_model_runner(monkeypatch, model_runner, preloaded_weights_bytes):
+    module_name = "sglang.srt.model_executor.model_runner"
+    module = ModuleType(module_name)
+    module.ModelRunner = model_runner
+    monkeypatch.setitem(sys.modules, module_name, module)
+    monkeypatch.setattr(patches, "_model_runner_patched", False)
+    monkeypatch.setattr(
+        patches,
+        "get_gms_memory_saver_impl",
+        lambda: SimpleNamespace(preloaded_weights_bytes=preloaded_weights_bytes),
+    )
+    patches.patch_model_runner()
+
+
+def test_patch_model_runner_adjusts_persistent_baseline_once(monkeypatch):
+    class ModelRunner:
+        def alloc_memory_pool(self, memory_pool_config=None):
+            self.calls.append((self.pre_model_load_memory, memory_pool_config))
+            return memory_pool_config
+
+    _patch_model_runner(monkeypatch, ModelRunner, 2 << 30)
+    patched_method = ModelRunner.alloc_memory_pool
+    patches.patch_model_runner()
+    monkeypatch.setattr(patches, "_model_runner_patched", False)
+    patches.patch_model_runner()
+
+    runner = ModelRunner()
+    runner.pre_model_load_memory = 10.0
+    runner.calls = []
+    positional_config = object()
+    keyword_config = object()
+
+    assert runner.alloc_memory_pool(positional_config) is positional_config
+    assert runner.alloc_memory_pool(memory_pool_config=keyword_config) is keyword_config
+    assert ModelRunner.alloc_memory_pool is patched_method
+    assert runner.pre_model_load_memory == 12.0
+    assert runner.calls == [(12.0, positional_config), (12.0, keyword_config)]
+
+
+def test_patch_model_runner_leaves_baseline_unchanged_without_preload(monkeypatch):
+    class ModelRunner:
+        def alloc_memory_pool(self):
+            return self.pre_model_load_memory
+
+    _patch_model_runner(monkeypatch, ModelRunner, 0)
+    runner = ModelRunner()
+    runner.pre_model_load_memory = 10.0
+
+    assert runner.alloc_memory_pool() == 10.0
+    assert runner.pre_model_load_memory == 10.0