torch-memory-saver 0.0.1__tar.gz → 0.0.2__tar.gz

{torch_memory_saver-0.0.1/torch_memory_saver.egg-info → torch_memory_saver-0.0.2}/PKG-INFO

@@ -1,5 +1,5 @@
 Metadata-Version: 2.1
 Name: torch_memory_saver
-Version: 0.0.1
+Version: 0.0.2
 Requires-Python: >=3.9
 License-File: LICENSE

torch_memory_saver-0.0.2/README.md

@@ -0,0 +1,29 @@
+# torch_memory_saver
+
+Allow torch tensor memory to be released and resumed later.
+
+API:
+
+```python
+memory_saver = TorchMemorySaver()
+
+# 1. For tensors that wants to be paused, create them within `region`
+with memory_saver.region():
+    x = torch.full((1_000_000_000,), 100, dtype=torch.uint8, device='cuda')
+
+# 2. After `pause`, CUDA memory is released for those tensors.
+# For example, check `nvidia-smi`'s memory usage to verify.
+memory_saver.pause()
+
+# 3. After `resume`, CUDA memory is re-occupied for those tensors.
+memory_saver.resume()
+```
+
+Please refer to https://github.com/sgl-project/sglang/issues/2542#issuecomment-2563641647 for details.
+
+TODO:
+
+- [x] Implementation
+- [x] Publish to pypi
+- [ ] More tests and infra
+- [ ] Documentation

torch_memory_saver-0.0.2/setup.py

@@ -0,0 +1,19 @@
+import logging
+
+import setuptools
+from setuptools import setup
+
+logger = logging.getLogger(__name__)
+
+setup(
+    name='torch_memory_saver',
+    version='0.0.2',
+    ext_modules=[setuptools.Extension(
+        'torch_memory_saver_cpp',
+        ['csrc/torch_memory_saver.cpp'],
+        extra_compile_args=['-I/usr/local/cuda/include'],
+        extra_link_args=['-lcuda'],
+    )],
+    python_requires=">=3.9",
+    packages=['torch_memory_saver'],
+)

{torch_memory_saver-0.0.1 → torch_memory_saver-0.0.2}/torch_memory_saver/__init__.py

@@ -2,6 +2,7 @@ import ctypes
 import logging
 import os
 from contextlib import contextmanager
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -18,31 +19,54 @@ class TorchMemorySaver:
 
     @contextmanager
     def region(self):
-        with torch.cuda.use_mem_pool(self._mem_pool):
-            _global_info.cdll.tms_region_enter()
-            try:
-                yield
-            finally:
-                _global_info.cdll.tms_region_leave()
+        if _global_info.binary_info.enabled:
+            with torch.cuda.use_mem_pool(self._mem_pool):
+                _global_info.binary_info.cdll.tms_region_enter()
+                try:
+                    yield
+                finally:
+                    _global_info.binary_info.cdll.tms_region_leave()
+        else:
+            yield
 
     def pause(self):
-        _global_info.cdll.tms_pause()
+        if _global_info.binary_info.enabled:
+            _global_info.binary_info.cdll.tms_pause()
 
     def resume(self):
-        _global_info.cdll.tms_resume()
+        if _global_info.binary_info.enabled:
+            _global_info.binary_info.cdll.tms_resume()
+
+
+@dataclass
+class _BinaryInfo:
+    cdll: Optional[ctypes.CDLL]
+
+    @property
+    def enabled(self):
+        return self.cdll is not None
+
+    @staticmethod
+    def compute():
+        env_ld_preload = os.environ.get('LD_PRELOAD', '')
+        if 'torch_memory_saver' in env_ld_preload:
+            return _BinaryInfo(cdll=ctypes.CDLL(env_ld_preload))
+        else:
+            logger.warning(
+                f'TorchMemorySaver is disabled for the current process because invalid LD_PRELOAD="{env_ld_preload}" (process_id={os.getpid()})')
+            return _BinaryInfo(cdll=None)
 
 
 class _GlobalInfo:
     def __init__(self):
-        self._cdll: Optional[ctypes.CDLL] = None
+        self._binary_info: Optional[_BinaryInfo] = None
         self._last_id = 0
 
     @property
-    def cdll(self):
-        if self._cdll is None:
-            self._cdll = _compute_cdll()
-            logger.debug(f'Use cdll={self._cdll}')
-        return self._cdll
+    def binary_info(self):
+        if self._binary_info is None:
+            self._binary_info = _BinaryInfo.compute()
+        return self._binary_info
 
     def next_id(self):
         self._last_id += 1
@@ -52,12 +76,6 @@ class _GlobalInfo:
 _global_info = _GlobalInfo()
 
 
-def _compute_cdll():
-    env_ld_preload = os.environ.get('LD_PRELOAD', '')
-    assert 'torch_memory_saver' in env_ld_preload, f'Please specify correct LD_PRELOAD (currently: {env_ld_preload})'
-    return ctypes.CDLL(env_ld_preload)
-
-
 def get_binary_path():
     dir_package = Path(__file__).parent
     candidates = [

{torch_memory_saver-0.0.1 → torch_memory_saver-0.0.2/torch_memory_saver.egg-info}/PKG-INFO

@@ -1,5 +1,5 @@
 Metadata-Version: 2.1
 Name: torch_memory_saver
-Version: 0.0.1
+Version: 0.0.2
 Requires-Python: >=3.9
 License-File: LICENSE

torch_memory_saver-0.0.1/README.md

@@ -1,11 +0,0 @@
-# torch_memory_saver
-
-Allow torch tensor memory to be released and resumed later 
-
-Please refer to https://github.com/sgl-project/sglang/issues/2542#issuecomment-2563641647 for details.
-
-TODO:
-
-- [x] Implementation
-- [ ] More tests and infra
-- [ ] Publish to pypi

torch_memory_saver-0.0.1/setup.py

@@ -1,19 +0,0 @@
-from setuptools import setup
-from torch.utils import cpp_extension
-
-ext_module = cpp_extension.CppExtension(
-    'torch_memory_saver_cpp',
-    ['csrc/torch_memory_saver.cpp'],
-    extra_compile_args=['-I/usr/local/cuda/include'],
-    extra_link_args=['-lcuda'],
-)
-
-setup(
-    name='torch_memory_saver',
-    version='0.0.1',
-    # https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-c-extension
-    ext_modules=[ext_module],
-    cmdclass={'build_ext': cpp_extension.BuildExtension},
-    python_requires=">=3.9",
-    packages=['torch_memory_saver'],
-)