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

torch_memory_saver/__init__.py

@@ -14,29 +14,34 @@ 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):
+    def region(self, tag: str = "default"):
+        """Context manager for memory saving with optional tag"""
         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_set_current_tag(tag.encode('utf-8'))
                 _global_info.binary_info.cdll.tms_region_enter()
                 try:
                     yield
                 finally:
+                    _global_info.binary_info.cdll.tms_set_current_tag(b"default")
                     _global_info.binary_info.cdll.tms_region_leave()
         else:
             yield
 
-    def pause(self):
+    def pause(self, tag: Optional[str] = None):
+        """Pause memory for specific tag or all memory if tag is None"""
         if _global_info.binary_info.enabled:
-            _global_info.binary_info.cdll.tms_pause()
+            tag_bytes = tag.encode('utf-8') if tag else None
+            _global_info.binary_info.cdll.tms_pause(tag_bytes)
 
-    def resume(self):
+    def resume(self, tag: Optional[str] = None):
+        """Resume memory for specific tag or all memory if tag is None"""
         if _global_info.binary_info.enabled:
-            _global_info.binary_info.cdll.tms_resume()
+            tag_bytes = tag.encode('utf-8') if tag else None
+            _global_info.binary_info.cdll.tms_resume(tag_bytes)
 
     @property
     def enabled(self):
@@ -46,7 +51,6 @@ class TorchMemorySaver:
         if self._mem_pool is None:
             self._mem_pool = torch.cuda.MemPool()
 
-
 @dataclass
 class _BinaryInfo:
     cdll: Optional[ctypes.CDLL]
@@ -55,21 +59,39 @@ class _BinaryInfo:
     def enabled(self):
         return self.cdll is not None
 
+    @staticmethod
+    def _setup_function_signatures(cdll):
+        """Define function signatures for the C library"""
+        cdll.tms_region_enter.argtypes = []
+        cdll.tms_region_leave.argtypes = []
+        cdll.tms_set_current_tag.argtypes = [ctypes.c_char_p]
+        cdll.tms_pause.argtypes = [ctypes.c_char_p]
+        cdll.tms_resume.argtypes = [ctypes.c_char_p]
+
     @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))
+            try:
+                cdll = ctypes.CDLL(env_ld_preload)
+                _BinaryInfo._setup_function_signatures(cdll)
+                return _BinaryInfo(cdll=cdll)
+            except OSError as e:
+                logger.error(f'Failed to load CDLL from {env_ld_preload}: {e}')
+                return _BinaryInfo(cdll=None)
         else:
-            logger.warning(
-                f'TorchMemorySaver is disabled for the current process because invalid LD_PRELOAD="{env_ld_preload}" (process_id={os.getpid()})')
+            print(
+                f'TorchMemorySaver is disabled for the current process because invalid LD_PRELOAD. '
+                f'You can use configure_subprocess() utility, '
+                f'or directly specify `LD_PRELOAD=/path/to/torch_memory_saver_cpp.some-postfix.so python your_script.py. '
+                f'(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):
@@ -77,13 +99,11 @@ class _GlobalInfo:
             self._binary_info = _BinaryInfo.compute()
         return self._binary_info
 
-    def next_id(self):
-        self._last_id += 1
-        return self._last_id
-
 
 _global_info = _GlobalInfo()
 
+# Global singleton instance
+torch_memory_saver = TorchMemorySaver()
 
 def get_binary_path():
     dir_package = Path(__file__).parent
@@ -92,7 +112,7 @@ def get_binary_path():
         for d in [dir_package, dir_package.parent]
         for p in d.glob('torch_memory_saver_cpp.*.so')
     ]
-    assert len(candidates) == 1, f'{candidates=}'
+    assert len(candidates) == 1, f'Expected exactly one torch_memory_saver_cpp library, found: {candidates}'
     return candidates[0]
 
 

{torch_memory_saver-0.0.6.dist-info → torch_memory_saver-0.0.8.dist-info}/METADATA

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

torch_memory_saver-0.0.8.dist-info/RECORD

@@ -0,0 +1,7 @@
+torch_memory_saver_cpp.abi3.so,sha256=kI54XLX1E_R-PgmfwrZoM2-emKOWQXRoabNEMZS0dCE,391672
+torch_memory_saver/__init__.py,sha256=MfMSSGSNhP7xBo4jq6RPug6svnMVPtmua8-l0yWlkTg,4403
+torch_memory_saver-0.0.8.dist-info/LICENSE,sha256=i806R5xShJFB4k9yNQJ2GYCcSBlu1frTx2vH_nWdWE8,1064
+torch_memory_saver-0.0.8.dist-info/METADATA,sha256=tUjtugtoDTFWeIct5kTvWRbGUUTlJlG85h_zSClzUdA,108
+torch_memory_saver-0.0.8.dist-info/WHEEL,sha256=HUPiMa7ZA9BvJ9gdJRYwZIjK2rWbCcrqYvJ4Onw0owE,102
+torch_memory_saver-0.0.8.dist-info/top_level.txt,sha256=uJ27-bVSKHxdcfHRcakvEr_KQxnUlMia6v19fHbfHxA,42
+torch_memory_saver-0.0.8.dist-info/RECORD,,

torch_memory_saver-0.0.6.dist-info/RECORD

@@ -1,7 +0,0 @@
-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,,