From 14046946474b66607af654a9ee3d9293465ea4f7 Mon Sep 17 00:00:00 2001 From: Schwinn Saereesitthipitak <17022745+galletas1712@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:14:19 -0700 Subject: [PATCH] fix(gms): support SGLang 0.5.16 memory pool API (#12445) Signed-off-by: Schwinn Saereesitthipitak --- .../tests/test_sglang_gms_memory_saver.py | 121 ++++++++++++++++++ .../integrations/sglang/memory_saver.py | 39 +++++- .../integrations/sglang/patches.py | 94 +++----------- .../tests/test_sglang_patches.py | 81 ++++++++++++ 4 files changed, 258 insertions(+), 77 deletions(-) create mode 100644 lib/gpu_memory_service/tests/test_sglang_patches.py diff --git a/components/src/dynamo/sglang/tests/test_sglang_gms_memory_saver.py b/components/src/dynamo/sglang/tests/test_sglang_gms_memory_saver.py index ed1cce83d821..7eff00acbfea 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_gms_memory_saver.py +++ b/components/src/dynamo/sglang/tests/test_sglang_gms_memory_saver.py @@ -4,6 +4,8 @@ from __future__ import annotations from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -23,7 +25,9 @@ pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0, + pytest.mark.profiled_vram_gib(0), pytest.mark.sglang, + pytest.mark.core, ] @@ -161,3 +165,120 @@ def test_region_requires_rw_allocator(build_impl, tag): with pytest.raises(RuntimeError, match=rf"requires '{tag}' to be RW"): with impl.region(tag, enable_cpu_backup=False): pass + + +def test_write_publication_follows_outermost_region_and_switches_weights_ro( + build_impl, monkeypatch +): + impl, weights, _, _ = build_impl() + model = torch.nn.Module() + impl.preloaded_weights_bytes = 456 + finalized_models = [] + events = [] + + @contextmanager + def traced_pool(tag, device): + events.append(f"pool_enter:{tag}") + try: + yield + finally: + events.append(f"pool_exit:{tag}") + + def finalize(allocator, pending_model): + events.append("finalize") + finalized_models.append(pending_model) + allocator.granted_lock_type = GrantedLockType.RO + return SimpleNamespace(committed_bytes=123) + + monkeypatch.setattr(gms_memory_saver, "gms_use_mem_pool", traced_pool) + monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize) + + with impl.region("weights", enable_cpu_backup=False): + with impl.region("kv_cache", enable_cpu_backup=False): + impl.finalize_write_mode(model) + events.append("defer") + + assert events == [ + "pool_enter:weights", + "pool_enter:kv_cache", + "defer", + "pool_exit:kv_cache", + "pool_exit:weights", + "finalize", + ] + assert finalized_models == [model] + assert weights.granted_lock_type == GrantedLockType.RO + assert impl.imported_weights_bytes == 123 + assert impl.preloaded_weights_bytes == 0 + + with impl.region("weights", enable_cpu_backup=False): + impl.finalize_write_mode(torch.nn.Module()) + + assert events[-1] == "finalize" + assert finalized_models == [model] + + +@pytest.mark.parametrize("failure_source", ["body", "pool_exit"]) +def test_nested_region_failure_discards_pending_publication( + build_impl, monkeypatch, failure_source +): + impl, weights, _, _ = build_impl() + stale_model = torch.nn.Module() + fresh_model = torch.nn.Module() + finalize = Mock(return_value=SimpleNamespace(committed_bytes=1)) + monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize) + + @contextmanager + def maybe_failing_pool(tag, device): + yield + if failure_source == "pool_exit" and tag == "kv_cache": + raise ValueError("pool exit failed") + + monkeypatch.setattr(gms_memory_saver, "gms_use_mem_pool", maybe_failing_pool) + + failure_message = failure_source.replace("_", " ") + with impl.region("weights", enable_cpu_backup=False): + impl.finalize_write_mode(stale_model) + with pytest.raises(ValueError, match=f"{failure_message} failed"): + with impl.region("kv_cache", enable_cpu_backup=False): + if failure_source == "body": + raise ValueError("body failed") + + finalize.assert_not_called() + + with impl.region("weights", enable_cpu_backup=False): + impl.finalize_write_mode(fresh_model) + finalize.assert_called_once_with(weights, fresh_model) + + +def test_failed_finalization_is_cleared_and_not_retried(build_impl, monkeypatch): + impl, weights, _, _ = build_impl() + model = torch.nn.Module() + finalize = Mock(side_effect=ValueError("finalization failed")) + monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize) + + with pytest.raises(ValueError, match="finalization failed"): + with impl.region("weights", enable_cpu_backup=False): + impl.finalize_write_mode(model) + + with impl.region("weights", enable_cpu_backup=False): + pass + finalize.assert_called_once_with(weights, model) + + +def test_invalid_or_duplicate_publication_preserves_first_model( + build_impl, monkeypatch +): + impl, weights, _, _ = build_impl() + first_model = torch.nn.Module() + finalize = Mock(return_value=SimpleNamespace(committed_bytes=1)) + monkeypatch.setattr(gms_memory_saver, "finalize_gms_write", finalize) + + with impl.region("weights", enable_cpu_backup=False): + impl.finalize_write_mode(first_model) + with pytest.raises(TypeError, match="must not be None"): + impl.finalize_write_mode(None) + with pytest.raises(RuntimeError, match="publication is already pending"): + impl.finalize_write_mode(torch.nn.Module()) + + finalize.assert_called_once_with(weights, first_model) diff --git a/lib/gpu_memory_service/integrations/sglang/memory_saver.py b/lib/gpu_memory_service/integrations/sglang/memory_saver.py index f0831292604b..570c09df407a 100644 --- a/lib/gpu_memory_service/integrations/sglang/memory_saver.py +++ b/lib/gpu_memory_service/integrations/sglang/memory_saver.py @@ -73,6 +73,8 @@ def __init__( self.imported_weights_bytes = 0 self.preloaded_weights_bytes = 0 self.ro_connect_timeout_ms = ro_connect_timeout_ms + self._active_region_depth = 0 + self._pending_write_model: Optional[torch.nn.Module] = None requested_mode = mode or RequestedLockType.RW_OR_RO self.allocators = { tag: get_or_create_gms_client_memory_manager( @@ -95,7 +97,7 @@ def __init__( @contextmanager def region(self, tag: str, enable_cpu_backup: bool): - """Mark allocation region with tag.""" + """Use the tag's RW pool and publish pending weights after clean exit.""" if enable_cpu_backup: raise ValueError( "SGLang with GMS does not support CPU backup for allocations." @@ -134,8 +136,19 @@ def region(self, tag: str, enable_cpu_backup: bool): f"SGLang with GMS requires {tag!r} to be RW for allocations; got {mode}" ) - with gms_use_mem_pool(tag, self._device): - yield + self._active_region_depth += 1 + clean_exit = False + try: + with gms_use_mem_pool(tag, self._device): + yield + clean_exit = True + finally: + self._active_region_depth -= 1 + if not clean_exit: + self._pending_write_model = None + + if self._active_region_depth == 0: + self._finalize_pending_write() @contextmanager def cuda_graph( @@ -184,9 +197,27 @@ def resume(self, tag: Optional[str] = None) -> None: self.allocators[target_tag].remap_all_vas() def finalize_write_mode(self, model: torch.nn.Module) -> None: - """Finalize write mode: register tensors, commit, and switch to read.""" + """Publish write-mode weights after all managed GMS regions exit.""" + if model is None: + raise TypeError("GMS weight publication model must not be None") + if self.allocators["weights"].granted_lock_type != GrantedLockType.RW: # Read-only import mode never republishes weights. + self._pending_write_model = None + return + + if self._pending_write_model is not None: + raise RuntimeError("GMS weight publication is already pending") + + self._pending_write_model = model + if self._active_region_depth == 0: + self._finalize_pending_write() + + def _finalize_pending_write(self) -> None: + """Consume and publish pending weights outside managed GMS regions.""" + model = self._pending_write_model + self._pending_write_model = None + if model is None: return stats = finalize_gms_write(self.allocators["weights"], model) diff --git a/lib/gpu_memory_service/integrations/sglang/patches.py b/lib/gpu_memory_service/integrations/sglang/patches.py index f364cca83d40..824a0787c3cc 100644 --- a/lib/gpu_memory_service/integrations/sglang/patches.py +++ b/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: diff --git a/lib/gpu_memory_service/tests/test_sglang_patches.py b/lib/gpu_memory_service/tests/test_sglang_patches.py new file mode 100644 index 000000000000..07c521c80051 --- /dev/null +++ b/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