fix(gms): cherry-pick support SGLang 0.5.16 memory pool API (#12445) - #12492
Conversation
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
🟡 Worker startup crashes hard on SGLang builds that lack the new memory-pool entry point
The startup patch grabs the new memory-pool entry point unconditionally (ModelRunner.alloc_memory_pool at lib/gpu_memory_service/integrations/sglang/patches.py:165) without any guard, so on an engine build that does not expose it the worker dies at import time instead of continuing with a warning.
Impact: Users on a slightly older engine release get an unexplained worker crash at start instead of a degraded-but-working start.
Loss of the previous version-tolerant patch path
The old implementation wrapped ModelRunner.init_memory_pool and probed its signature with inspect, tolerating both old and new parameter names, and the only failure mode was the ImportError branch (lib/gpu_memory_service/integrations/sglang/patches.py:156-160) which logs a warning and returns. The new code accesses ModelRunner.alloc_memory_pool directly outside that try, so an AttributeError propagates out of patch_model_runner(), which is executed at module import in the scheduler child process (lib/gpu_memory_service/integrations/sglang/model_loader.py:44). Similarly, patched_alloc_memory_pool reads self.pre_model_load_memory (patches.py:175-176) with no getattr fallback, so a renamed/absent attribute raises during read-mode startup instead of being logged and skipped as the old code did. Wrapping the attribute lookups (e.g. getattr(ModelRunner, "alloc_memory_pool", None) and a hasattr check on the baseline attribute) restores graceful degradation.
Prompt for agents
In lib/gpu_memory_service/integrations/sglang/patches.py, patch_model_runner() now hard-depends on the SGLang 0.5.16 API: it reads ModelRunner.alloc_memory_pool at module import time (via model_loader.py's top-level patch_model_runner() call) and the patched wrapper reads self.pre_model_load_memory unconditionally. Previously the patch tolerated multiple SGLang signatures and only failed softly via the ImportError branch. Restore graceful degradation so that a missing method or missing baseline attribute logs a warning and leaves SGLang untouched rather than raising AttributeError inside the scheduler child process.
Was this helpful? React with 👍 or 👎 to provide feedback.
| old_value = self.pre_model_load_memory | ||
| self.pre_model_load_memory += preloaded_weights_gib | ||
| self._gms_memory_baseline_adjusted = True |
There was a problem hiding this comment.
🔍 Baseline adjustment now mutates ModelRunner state permanently
Previously the preloaded-weight correction was applied only to the argument passed into init_memory_pool, leaving ModelRunner state untouched. Now self.pre_model_load_memory is permanently increased and guarded by a per-instance _gms_memory_baseline_adjusted flag. That makes the correction survive later re-allocations of the memory pool (e.g. resume-memory-occupation flows), which appears to be the intent, but it also means any other SGLang code that reads pre_model_load_memory (logging, memory accounting, later KV re-sizing) now sees an inflated baseline. Worth confirming against SGLang 0.5.16 that no other consumer of that attribute is affected.
Was this helpful? React with 👍 or 👎 to provide feedback.
|
🎯 Code Coverage (details) 🔗 Commit SHA: 1404694 | Docs | Datadog PR Page | Give us feedback! |
2 participants
Summary
Cherry-picks merged #12445 (
290f96bc3eff3984a18a6a9f8f33b3a2d50fe2e0) ontorelease/1.4.0.Main PR
#12445
Test plan