feat(hooks): add pluggable hook modes - #79
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a registry for hook modes (register_hook_mode) allowing custom lazy factories to be registered, and updates TorchMemorySaver to resolve its implementation lazily. The feedback suggests several improvements: implementing __getstate__ and __setstate__ on TorchMemorySaver to prevent pickling errors after initialization, simplifying the README example imports for consistency, and raising descriptive RuntimeErrors instead of cryptic KeyErrors in _ensure_initialized and _uses_preload if the selected hook mode is not registered in the current process.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
96811ec to
9e3d8ea
Compare
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
9e3d8ea to
cd7b4b6
Compare
|
(replied in #77) |
Closes #77 (pluggable memory-saver implementations) — the
torch_memory_saverpiece of the GPU Memory Service integration proposed in sgl-project/sglang#27310.register_hook_mode(name, factory, uses_preload=...)registers a lazily constructed custom implementation, selected through the existinghook_modesetter. Factories are process-global and resolved on first use; the builtinpreload/torchmodes are registered the same way and behave exactly as before.configure_subprocess()now consults the selected mode'suses_preloadflag and skips theLD_PRELOADsetup for modes that do not use the preload hook (both builtins keep it, preserving current behavior).