Merge pull request #184 from mert-kurttutan/main

Implement get_total_memory for windows

Author: noahho

.github/workflows/pull_request.yml

@@ -33,12 +33,18 @@ jobs:
           # https://github.com/actions/setup-python/issues/855
             python-version: "3.9"
             dependency-set: minimum
+          - os: windows-latest
+            python-version: "3.9"
+            dependency-set: minimum
           - os: ubuntu-latest
             python-version: "3.12"
             dependency-set: maximum
           - os: macos-latest
             python-version: "3.12"
             dependency-set: maximum
+          - os: windows-latest
+            python-version: "3.12"
+            dependency-set: maximum
     runs-on: ${{ matrix.os }}
 
     steps:
@@ -58,64 +64,20 @@ jobs:
       - name: Generate requirements file for minimum dependencies
         if: matrix.dependency-set == 'minimum'
         run: |
-          python << EOF
-          import re
-          
-          with open('pyproject.toml', 'r') as f:
-              content = f.read()
-          
-          # Find dependencies section using regex
-          deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
-          if deps_match:
-              deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
-              min_reqs = []
-              for dep in deps:
-                  match = re.match(r'([^>=<\s]+)\s*>=\s*([^,\s"\']+)', dep)
-                  if match:
-                      package, min_ver = match.groups()
-                      min_reqs.append(f"{package}=={min_ver}")
-          
-              with open('requirements.txt', 'w') as f:
-                  f.write('\n'.join(min_reqs))
-          EOF
+          python scripts/get_min_dependencies.py
 
       - name: Generate requirements file for maximum dependencies
         if: matrix.dependency-set == 'maximum'
         run: |
-          python << EOF
-          import re
-          
-          with open('pyproject.toml', 'r') as f:
-              content = f.read()
-          
-          # Find dependencies section using regex
-          deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
-          if deps_match:
-              deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
-              max_reqs = []
-              for dep in deps:
-                  # Check for maximum version constraint
-                  max_version_match = re.search(r'([^>=<\s]+).*?<\s*([^,\s"\']+)', dep)
-                  if max_version_match:
-                      # If there's a max version, use the version just below it
-                      package, max_ver = max_version_match.groups()
-                      max_reqs.append(f"{package}<{max_ver}")
-                  else:
-                      # If no max version, just use the package name
-                      package = re.match(r'([^>=<\s]+)', dep).group(1)
-                      max_reqs.append(package)
-          
-              with open('requirements.txt', 'w') as f:
-                  f.write('\n'.join(max_reqs))
-          EOF
+          python scripts/get_max_dependencies.py
 
       - name: Install dependencies
         run: |
           uv pip install --system --no-deps .
           # onnx is required for onnx export tests
           # we don't install all dev dependencies here for speed
           uv pip install --system -r requirements.txt
-          uv pip install --system pytest onnx
+          uv pip install --system pytest onnx psutil
 
       - name: Initialize submodules
         run: git submodule update --init --recursive

pyproject.toml

@@ -62,6 +62,7 @@ dev = [
   # Test
   "pytest",
   "onnx", # required for onnx export tests
+  "psutil", # required for testing internal memory tool on windows
   # Docs
   "mkdocs",
   "mkdocs-material",

scripts/get_max_dependencies.py

@@ -0,0 +1,28 @@
+import re
+
+def main() -> None:
+    with open('pyproject.toml', 'r') as f:
+        content = f.read()
+
+    # Find dependencies section using regex
+    deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
+    if deps_match:
+        deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
+        max_reqs = []
+        for dep in deps:
+            # Check for maximum version constraint
+            max_version_match = re.search(r'([^>=<\s]+).*?<\s*([^,\s"\']+)', dep)
+            if max_version_match:
+                # If there's a max version, use the version just below it
+                package, max_ver = max_version_match.groups()
+                max_reqs.append(f"{package}<{max_ver}")
+            else:
+                # If no max version, just use the package name
+                package = re.match(r'([^>=<\s]+)', dep).group(1)
+                max_reqs.append(package)
+
+        with open('requirements.txt', 'w') as f:
+            f.write('\n'.join(max_reqs))
+
+if __name__ == '__main__':
+    main()

scripts/get_min_dependencies.py

@@ -0,0 +1,22 @@
+import re
+
+def main() -> None:
+    with open('pyproject.toml', 'r') as f:
+        content = f.read()
+
+    # Find dependencies section using regex
+    deps_match = re.search(r'dependencies\s*=\s*\[(.*?)\]', content, re.DOTALL)
+    if deps_match:
+        deps = [d.strip(' "\'') for d in deps_match.group(1).strip().split('\n') if d.strip()]
+        min_reqs = []
+        for dep in deps:
+            match = re.match(r'([^>=<\s]+)\s*>=\s*([^,\s"\']+)', dep)
+            if match:
+                package, min_ver = match.groups()
+                min_reqs.append(f"{package}=={min_ver}")
+
+        with open('requirements.txt', 'w') as f:
+            f.write('\n'.join(min_reqs))
+
+if __name__ == '__main__':
+    main()

src/tabpfn/model/memory.py

@@ -252,12 +252,23 @@ def get_max_free_memory(
                     os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 1e9
                 )
             except AttributeError:
-                # TODO: `os.sysconf` does not exist on windows.
-                free_memory = cls.convert_units(
-                    default_gb_cpu_if_failed_to_calculate,
-                    "gb",
-                    "b",
-                )
+                from tabpfn.utils import get_total_memory_windows
+
+                if os.name == "nt":
+                    free_memory = get_total_memory_windows()
+                else:
+                    warnings.warn(
+                        "Could not get system memory, defaulting to"
+                        f" {default_gb_cpu_if_failed_to_calculate} GB",
+                        RuntimeWarning,
+                        stacklevel=2,
+                    )
+                    free_memory = cls.convert_units(
+                        default_gb_cpu_if_failed_to_calculate,
+                        "gb",
+                        "b",
+                    )
+
             except ValueError:
                 warnings.warn(
                     "Could not get system memory, defaulting to"

src/tabpfn/utils.py

@@ -4,8 +4,10 @@
 
 from __future__ import annotations
 
+import ctypes
 import os
 import sys
+import typing
 import warnings
 from collections.abc import Sequence
 from pathlib import Path
@@ -829,3 +831,39 @@ def _transform_borders_one(
         )
 
     return logit_cancel_mask, descending_borders, borders_t
+
+
+# Terminology: Use memory to referent physical memory, swap for swap memory
+def get_total_memory_windows() -> float:
+    """Get the total memory of the system for windows OS, using windows API.
+
+    Returns:
+        The total memory of the system in GB.
+    """
+
+    # ref: https://github.com/microsoft/windows-rs/blob/c9177f7a65c764c237a9aebbd3803de683bedaab/crates/tests/bindgen/src/fn_return_void_sys.rs#L12
+    # ref: https://learn.microsoft.com/en-us/windows/win32/api/sysinfoapi/ns-sysinfoapi-memorystatusex
+    # this class is needed to load the memory status with GlobalMemoryStatusEx function
+    # using win32 API, for more details see microsoft docs link above
+    class _MEMORYSTATUSEX(ctypes.Structure):
+        _fields_: typing.ClassVar = [
+            ("dwLength", ctypes.c_ulong),
+            ("dwMemoryLoad", ctypes.c_ulong),
+            ("ullTotalPhys", ctypes.c_ulonglong),
+            ("ullAvailPhys", ctypes.c_ulonglong),
+            ("ullTotalPageFile", ctypes.c_ulonglong),
+            ("ullAvailPageFile", ctypes.c_ulonglong),
+            ("ullTotalVirtual", ctypes.c_ulonglong),
+            ("ullAvailVirtual", ctypes.c_ulonglong),
+            ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
+        ]
+
+    # Initialize the structure
+    mem_status = _MEMORYSTATUSEX()
+    # need to initialize lenght of structure, see microsft docs above
+    mem_status.dwLength = ctypes.sizeof(_MEMORYSTATUSEX)
+
+    k32_lib = ctypes.windll.LoadLibrary("kernel32.dll")
+    k32_lib.GlobalMemoryStatusEx(ctypes.byref(mem_status))
+
+    return mem_status.ullTotalPhys / 1e9  # Convert bytes to GB

tests/test_classifier_interface.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+import os
 from itertools import product
 from typing import Callable, Literal
 
@@ -248,6 +249,8 @@ def forward(
 
 @pytest.mark.filterwarnings("ignore::torch.jit.TracerWarning")
 def test_onnx_exportable_cpu(X_y: tuple[np.ndarray, np.ndarray]) -> None:
+    if os.name == "nt":
+        pytest.skip("onnx export is not tested on windows")
     X, y = X_y
     with torch.no_grad():
         classifier = TabPFNClassifier(n_estimators=1, device="cpu", random_state=42)

tests/test_regressor_interface.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+import os
 from itertools import product
 from typing import Callable, Literal
 
@@ -246,6 +247,8 @@ def forward(
 # WARNING: unstable for scipy<1.11.0
 @pytest.mark.filterwarnings("ignore::torch.jit.TracerWarning")
 def test_onnx_exportable_cpu(X_y: tuple[np.ndarray, np.ndarray]) -> None:
+    if os.name == "nt":
+        pytest.skip("onnx export is not tested on windows")
     X, y = X_y
     with torch.no_grad():
         regressor = TabPFNRegressor(n_estimators=1, device="cpu", random_state=43)

tests/test_utils.py

@@ -0,0 +1,41 @@
+# use get_total_memory and compare it against result from psutils
+# run it only if the it is windows os.name == "nt"
+from __future__ import annotations
+
+import os
+
+
+def test_internal_windows_total_memory():
+    if os.name == "nt":
+        import psutil
+
+        from tabpfn.utils import get_total_memory_windows
+
+        utils_result = get_total_memory_windows()
+        psutil_result = psutil.virtual_memory().total / 1e9
+        assert utils_result == psutil_result
+
+
+def test_internal_windows_total_memory_multithreaded():
+    # collect results from multiple threads
+    if os.name == "nt":
+        import threading
+
+        import psutil
+
+        from tabpfn.utils import get_total_memory_windows
+
+        results = []
+
+        def get_memory():
+            results.append(get_total_memory_windows())
+
+        threads = []
+        for _ in range(10):
+            t = threading.Thread(target=get_memory)
+            threads.append(t)
+            t.start()
+        for t in threads:
+            t.join()
+        psutil_result = psutil.virtual_memory().total / 1e9
+        assert all(result == psutil_result for result in results)