LD_PRELOAD Script - #8
Conversation
- Package data contained in torch_memory_saver - tms python sets `LD_PRELOAD`
There was a problem hiding this comment.
nit: what about naming tms to torch_memory_saver
| import torch_memory_saver, importlib.machinery, pathlib, os | ||
| package_dir = pathlib.Path(torch_memory_saver.__file__).parent | ||
| ext_suffix = importlib.machinery.EXTENSION_SUFFIXES[0] | ||
| lib_name = f'torch_memory_saver_cpp{ext_suffix}' | ||
| print(str(package_dir / lib_name)) |
There was a problem hiding this comment.
nit: what about we provide such a function in torch_memory_saver python code, and here we only need to do as simple as $(python -c "import torch_memory_saver; torch_memory_saver.that_function()"
| ") | ||
|
|
||
| # Detect if script is being sourced | ||
| (return 0 2>/dev/null) && SOURCED=1 || SOURCED=0 |
There was a problem hiding this comment.
nit: curious whether this will accurately detect whether it is being sourced, or will there false positive / negative?
| return _BinaryInfo(cdll=ctypes.CDLL(env_ld_preload)) | ||
| paths = env_ld_preload.split(':') | ||
| lib_path = next((path for path in paths if 'torch_memory_saver' in path), None) | ||
| if lib_path: |
There was a problem hiding this comment.
nit: looks like we may have a bug for the else case (or, if the else branch will never occur, we may change if to assert)
| assert os.environ[key] == value | ||
| os.environ[key] = old_value | ||
| logger.debug(f'change_env restore key={key} value={old_value}') | ||
| _global_info = _GlobalInfo() No newline at end of file |
There was a problem hiding this comment.
well, it would be great to keep the old configure_subprocess etc, because it is how SGLang uses it. https://github.com/sgl-project/sglang/blob/8aab7fdb21e86b87a3141951ce498246fd4aae73/python/sglang/srt/entrypoints/engine.py#L522
There was a problem hiding this comment.
Will look at the approach later and modify PR, not by system for a few days
|
|
||
| ## Installation | ||
|
|
||
| For the library to work properly, it needs to be preloaded using the `LD_PRELOAD` environment variable. |
There was a problem hiding this comment.
nit: would be great to mention the SGLang approach, which indeed does not require any operations
|
@EduardDurech Btw, I am curious how/where are you using torch_memory_saver? |
|
@fzyzcjy right now standalone API import os, torch
from torch_memory_saver import TorchMemorySaver
memory_saver = TorchMemorySaver()
print(os.system('nvidia-smi'))
with memory_saver.region():
x = torch.full((1_000_000_000,), 100, dtype=torch.uint8, device='cuda')
print(os.system('nvidia-smi'))
memory_saver.pause()
print(os.system('nvidia-smi'))
memory_saver.resume()
print(os.system('nvidia-smi'))Testing, will use with veRL-SGLang, did you guys implement |
Yes, if you are using memory saver inside SGLang, then the LD_PRELOAD is automatically configured and you do not need to do anything to make it work |
|
Okay, do you know where in the code it finds the .so? I couldn't find invocation of torch_memory_saver, maybe it would be good to be compatible as standalone and as dependency, but currently the PR moves the .so to packages/torch_memory_saver instead of packages |
Just follow the configure_subprocess function you will see |
2 participants
LD_PRELOAD