torch-memory-saver 0.0.3__tar.gz → 0.0.5__tar.gz

{torch_memory_saver-0.0.3/torch_memory_saver.egg-info → torch_memory_saver-0.0.5}/PKG-INFO

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

torch_memory_saver-0.0.5/setup.py

@@ -0,0 +1,41 @@
+import os
+import glob
+import shutil
+import logging
+import setuptools
+from pathlib import Path
+from setuptools import setup
+
+logger = logging.getLogger(__name__)
+
+# copy & modify from torch/utils/cpp_extension.py
+def _find_cuda_home():
+    """Find the CUDA install path."""
+    # Guess #1
+    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+    if cuda_home is None:
+        # Guess #2
+        nvcc_path = shutil.which("nvcc")
+        if nvcc_path is not None:
+            cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
+        else:
+            # Guess #3
+            cuda_home = '/usr/local/cuda'
+    return cuda_home
+
+cuda_home = Path(_find_cuda_home())
+include_dirs = [
+    str(cuda_home.resolve() / 'targets/x86_64-linux/include'),
+]
+setup(
+    name='torch_memory_saver',
+    version='0.0.5',
+    ext_modules=[setuptools.Extension(
+        'torch_memory_saver_cpp',
+        ['csrc/torch_memory_saver.cpp'],
+        include_dirs=include_dirs,
+        libraries=['cuda']
+    )],
+    python_requires=">=3.9",
+    packages=['torch_memory_saver'],
+)

{torch_memory_saver-0.0.3 → torch_memory_saver-0.0.5}/torch_memory_saver/__init__.py

@@ -13,13 +13,14 @@ logger = logging.getLogger(__name__)
 
 class TorchMemorySaver:
     def __init__(self):
-        self._mem_pool = torch.cuda.MemPool()
+        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:
@@ -41,6 +42,10 @@ class TorchMemorySaver:
     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:

{torch_memory_saver-0.0.3 → torch_memory_saver-0.0.5/torch_memory_saver.egg-info}/PKG-INFO

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

torch_memory_saver-0.0.3/setup.py

@@ -1,19 +0,0 @@
-import logging
-
-import setuptools
-from setuptools import setup
-
-logger = logging.getLogger(__name__)
-
-setup(
-    name='torch_memory_saver',
-    version='0.0.3',
-    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'],
-)