Create the CUDA context only when a GPU tensor is actually used

RELEASES.md

@@ -8,6 +8,7 @@
 
 #### Closed issues
 
+- Build the CUDA generator and the CUDA entries of `TorchBackend.__type_list__` lazily, so that using POT with CPU-only torch tensors no longer initialises a CUDA context and claims device memory (PR #847, Issue #612)
 - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842)
 - Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816)
 - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840)

ot/backend.py

@@ -2066,30 +2066,25 @@ class TorchBackend(Backend):
 
     __name__ = "torch"
     __type__ = torch_type
-    __type_list__ = None
+    # __type_list__ is a property below: its CUDA entries are built lazily.
 
     rng_ = None
 
     def __init__(self):
         self.rng_ = torch.Generator("cpu")
         self.rng_.seed()
 
-        self.__type_list__ = [
+        # The CUDA generator and the CUDA entries of the type list are built the
+        # first time something asks for them. Building either here initialises a
+        # CUDA context, which claims device memory and wakes the GPU even when
+        # the computation stays entirely on the CPU.
+        self._type_list_cpu = [
             torch.tensor(1, dtype=torch.float32),
             torch.tensor(1, dtype=torch.float64),
         ]
-
-        if torch.cuda.is_available():
-            self.rng_cuda_ = torch.Generator("cuda")
-            self.rng_cuda_.seed()
-            self.__type_list__.append(
-                torch.tensor(1, dtype=torch.float32, device="cuda")
-            )
-            self.__type_list__.append(
-                torch.tensor(1, dtype=torch.float64, device="cuda")
-            )
-        else:
-            self.rng_cuda_ = torch.Generator("cpu")
+        self._type_list_cuda = None
+        self._rng_cuda = None
+        self._cuda_seed = None
 
         from torch.autograd import Function
         from torch.autograd.function import once_differentiable
@@ -2133,6 +2128,35 @@ def backward(ctx, g):
         self.ValFunction = ValFunction
         self.MatrixSqrtFunction = MatrixSqrtFunction
 
+    @property
+    def __type_list__(self):
+        if self._type_list_cuda is None:
+            if torch.cuda.is_available():
+                self._type_list_cuda = [
+                    torch.tensor(1, dtype=torch.float32, device="cuda"),
+                    torch.tensor(1, dtype=torch.float64, device="cuda"),
+                ]
+            else:
+                self._type_list_cuda = []
+        return self._type_list_cpu + self._type_list_cuda
+
+    @property
+    def rng_cuda_(self):
+        if self._rng_cuda is None:
+            if torch.cuda.is_available():
+                self._rng_cuda = torch.Generator("cuda")
+                if self._cuda_seed is None:
+                    self._rng_cuda.seed()
+                else:
+                    self._rng_cuda.manual_seed(self._cuda_seed)
+            else:
+                self._rng_cuda = torch.Generator("cpu")
+        return self._rng_cuda
+
+    @rng_cuda_.setter
+    def rng_cuda_(self, generator):
+        self._rng_cuda = generator
+
     def _to_numpy(self, a):
         if isinstance(a, float) or isinstance(a, int) or isinstance(a, np.ndarray):
             return np.array(a)
@@ -2412,7 +2436,11 @@ def seed(self, seed=None):
             pass
         elif isinstance(seed, int):
             self.rng_.manual_seed(seed)
-            self.rng_cuda_.manual_seed(seed)
+            # Remember the seed rather than forcing the CUDA generator into
+            # existence; it is applied when that generator is first needed.
+            self._cuda_seed = seed
+            if self._rng_cuda is not None:
+                self._rng_cuda.manual_seed(seed)
         elif isinstance(seed, torch.Generator):
             if self.device_type(seed) == "GPU":
                 self.rng_cuda_ = seed

test/test_backend.py

@@ -942,3 +942,31 @@ def test_torch_optimizer_after_tensorflow_import():
         f"interpreter died with returncode {result.returncode}: "
         f"{result.stderr.decode(errors='replace')[-2000:]}"
     )
+
+
+@pytest.mark.skipif(
+    not torch or not torch.cuda.is_available(),
+    reason="Requires torch with CUDA available",
+)
+def test_no_cuda_context_for_cpu_only_work():
+    """Non-regression test for issue #612.
+
+    Building the torch backend used to create a CUDA generator and the CUDA
+    entries of the type list straight away. That initialises a CUDA context, so
+    device memory is claimed and the GPU wakes up for a computation that stays
+    entirely on the CPU. The check needs an interpreter that has not touched
+    CUDA yet, so it runs in a subprocess.
+    """
+    code = (
+        "import torch\n"
+        "import ot\n"
+        "x = torch.randn(64, 2)\n"
+        "ot.dist(x, x)\n"
+        "assert not torch.cuda.is_initialized(), 'a CUDA context was created'\n"
+        "assert torch.cuda.memory_allocated() == 0, 'device memory was claimed'\n"
+    )
+    result = subprocess.run([sys.executable, "-c", code], capture_output=True)
+    assert result.returncode == 0, (
+        f"interpreter exited with returncode {result.returncode}: "
+        f"{result.stderr.decode(errors='replace')[-2000:]}"
+    )