torch-memory-saver 0.0.6__cp39-abi3-manylinux2014_x86_64.whl

torch_memory_saver/__init__.py

@@ -0,0 +1,115 @@
+import ctypes
+import logging
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class TorchMemorySaver:
+    def __init__(self):
+        self._mem_pool = None
+        self._id = _global_info.next_id()
+        assert self._id == 1, 'Only support one single instance yet (multi-instance will be implemented later)'
+
+    @contextmanager
+    def region(self):
+        if _global_info.binary_info.enabled:
+            self._ensure_mem_pool()
+            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):
+        if _global_info.binary_info.enabled:
+            _global_info.binary_info.cdll.tms_pause()
+
+    def resume(self):
+        if _global_info.binary_info.enabled:
+            _global_info.binary_info.cdll.tms_resume()
+
+    @property
+    def enabled(self):
+        return _global_info.binary_info.enabled
+
+    def _ensure_mem_pool(self):
+        if self._mem_pool is None:
+            self._mem_pool = torch.cuda.MemPool()
+
+
+@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._binary_info: Optional[_BinaryInfo] = None
+        self._last_id = 0
+
+    @property
+    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
+        return self._last_id
+
+
+_global_info = _GlobalInfo()
+
+
+def get_binary_path():
+    dir_package = Path(__file__).parent
+    candidates = [
+        p
+        for d in [dir_package, dir_package.parent]
+        for p in d.glob('torch_memory_saver_cpp.*.so')
+    ]
+    assert len(candidates) == 1, f'{candidates=}'
+    return candidates[0]
+
+
+@contextmanager
+def configure_subprocess():
+    with change_env('LD_PRELOAD', str(get_binary_path())):
+        yield
+
+
+@contextmanager
+def change_env(key: str, value: str):
+    old_value = os.environ.get(key, '')
+    os.environ[key] = value
+    logger.debug(f'change_env set key={key} value={value}')
+    try:
+        yield
+    finally:
+        assert os.environ[key] == value
+        os.environ[key] = old_value
+        logger.debug(f'change_env restore key={key} value={old_value}')

torch_memory_saver-0.0.6.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 fzyzcjy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_memory_saver-0.0.6.dist-info/METADATA

@@ -0,0 +1,6 @@
+Metadata-Version: 2.1
+Name: torch_memory_saver
+Version: 0.0.6
+Requires-Python: >=3.9
+License-File: LICENSE
+

torch_memory_saver-0.0.6.dist-info/RECORD

@@ -0,0 +1,7 @@
+torch_memory_saver_cpp.abi3.so,sha256=OCweTnvdmyg5zhUIMJjfH9NW0lYjtcNwqzS9-89cCvQ,315896
+torch_memory_saver/__init__.py,sha256=B3AXwxxJeUbNFKdrfaGzXvl3vTcgPOf2UjaFVtGCZ68,3072
+torch_memory_saver-0.0.6.dist-info/LICENSE,sha256=i806R5xShJFB4k9yNQJ2GYCcSBlu1frTx2vH_nWdWE8,1064
+torch_memory_saver-0.0.6.dist-info/METADATA,sha256=P21LYFkCJHFwaMAkxZBoiQRkhvIQmOBAKbHHoQdQiEI,108
+torch_memory_saver-0.0.6.dist-info/WHEEL,sha256=HUPiMa7ZA9BvJ9gdJRYwZIjK2rWbCcrqYvJ4Onw0owE,102
+torch_memory_saver-0.0.6.dist-info/top_level.txt,sha256=uJ27-bVSKHxdcfHRcakvEr_KQxnUlMia6v19fHbfHxA,42
+torch_memory_saver-0.0.6.dist-info/RECORD,,

torch_memory_saver-0.0.6.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.0.0)
+Root-Is-Purelib: false
+Tag: cp39-abi3-linux_x86_64
+

torch_memory_saver-0.0.6.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+torch_memory_saver
+torch_memory_saver_cpp