Make sure gpu4pyscf allocator is never used

src/skala/conftest.py

@@ -13,6 +13,7 @@
 
 This file lives in ``src/skala/`` rather than ``src/skala/gpu4pyscf/`` so
 pytest can load it without first importing ``skala.gpu4pyscf``.
+Similar problems arise for torch_allocator and cupy.
 """
 
 import torch
@@ -21,3 +22,4 @@
 
 if not torch.cuda.is_available():
     collect_ignore_glob.append("gpu4pyscf/*.py")
+    collect_ignore_glob.append("utils/torch_allocator.py")

src/skala/gpu4pyscf/__init__.py

@@ -35,11 +35,6 @@
 
 from skala.functional import ExcFunctionalBase, load_functional
 from skala.gpu4pyscf import dft
-from skala.gpu4pyscf.torch_allocator import use_torch_mempool_in_cupy
-
-# Install this last because GPU4PySCF configures its own CuPy allocator during import.
-# Subsequent CuPy allocations use PyTorch's caching allocator
-use_torch_mempool_in_cupy()
 
 
 def SkalaKS(

src/skala/pyscf/backend.py

@@ -14,14 +14,14 @@
 GPU_EXCEPTION: BaseException | None = None
 
 __all__ = [
+    "KS",
     "Array",
     "Grid",
-    "KS",
-    "dft_gpu",
     "check_gpu_imports_were_successful",
+    "dft_gpu",
     "from_numpy_or_cupy",
-    "to_numpy",
     "to_cupy",
+    "to_numpy",
 ]
 
 
@@ -40,6 +40,21 @@
         import cupy
         from gpu4pyscf import dft as dft_gpu
 
+        from skala.utils.torch_allocator import use_torch_mempool_in_cupy
+
+        # Install this last because GPU4PySCF configures its own CuPy allocator during import.
+        # Subsequent CuPy allocations use PyTorch's caching allocator
+        if not torch.cuda.is_available():
+            raise ImportError(
+                "CUDA is not available; cannot configure CuPy to use the PyTorch allocator."
+            )
+        try:
+            use_torch_mempool_in_cupy()
+        except RuntimeError as e:
+            raise ImportError(
+                "Failed to configure CuPy to use the PyTorch allocator."
+            ) from e
+
         Array = TypeVar("Array", np.ndarray, cupy.ndarray)
         Grid: TypeAlias = dft.Grids | dft_gpu.Grids
         KS: TypeAlias = dft.rks.RKS | dft.uks.UKS | dft_gpu.rks.RKS | dft_gpu.uks.UKS

src/skala/gpu4pyscf/torch_allocator.py → src/skala/utils/torch_allocator.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-"""Use torch allocator in CuPy.
+"""Use the PyTorch allocator in CuPy.
 
 Adapted from pytorch_pfn_extras.
 https://github.com/pfnet/pytorch-pfn-extras/blob/master/pytorch_pfn_extras/cuda/_allocator.py

tests/test_gpu4pyscf_gradients.py

@@ -22,14 +22,15 @@
 from test_pyscf_gradients import FULL_GRAD_REF
 
 from skala.functional.base import ExcFunctionalBase
-from skala.gpu4pyscf import SkalaKS, torch_allocator
+from skala.gpu4pyscf import SkalaKS
 from skala.gpu4pyscf.gradients import (
     SkalaRKSGradient,
     SkalaUKSGradient,
     nuc_grad_from_veff,
     veff_and_expl_nuc_grad,
 )
 from skala.pyscf.features import generate_features
+from skala.utils import torch_allocator
 
 
 def test_torch_allocator_is_active_after_import() -> None: