Skip to content

feat(hooks): add pluggable hook modes - #79

Open
galletas1712 wants to merge 1 commit into
fzyzcjy:masterfrom
galletas1712:feat/pluggable-hook-modes
Open

feat(hooks): add pluggable hook modes#79
galletas1712 wants to merge 1 commit into
fzyzcjy:masterfrom
galletas1712:feat/pluggable-hook-modes

Conversation

@galletas1712

@galletas1712 galletas1712 commented Jul 1, 2026

Copy link
Copy Markdown

Closes #77 (pluggable memory-saver implementations) — the torch_memory_saver piece 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 existing hook_mode setter. Factories are process-global and resolved on first use; the builtin preload/torch modes are registered the same way and behave exactly as before.
  • The package-level configure_subprocess() now consults the selected mode's uses_preload flag and skips the LD_PRELOAD setup for modes that do not use the preload hook (both builtins keep it, preserving current behavior).
  • Selecting an unknown mode, re-registering a name, or using a mode never registered in the current process (e.g. a subprocess that skipped registration) fails with a descriptive error.

@galletas1712 galletas1712 changed the title feat: add pluggable hook modes feat(hooks): add pluggable hook modes Jul 1, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread torch_memory_saver/entrypoint.py
Comment thread README.md
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
@galletas1712
galletas1712 force-pushed the feat/pluggable-hook-modes branch from 96811ec to 9e3d8ea Compare July 2, 2026 07:30
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
@fzyzcjy

fzyzcjy commented Jul 10, 2026

Copy link
Copy Markdown
Owner

(replied in #77)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[RFC] Pluggable memory-saver implementation in torch_memory_saver

2 participants