Fix silent GPU transfer failures disabling CPU threading fallback

osipy/common/backend/array_module.py

@@ -31,6 +31,7 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -201,6 +202,11 @@ def to_gpu(array: ArrayLike) -> Any:
     - If force_cpu is True, returns NumPy array
     - If CuPy is not available, returns NumPy array
     - If input is already on GPU, returns as-is (no copy)
+    - If the GPU transfer itself fails (e.g. out of memory), a
+      ``UserWarning`` is issued and a NumPy array is returned. Callers
+      that decide GPU-vs-CPU behavior (chunk sizing, threading) based on
+      whether GPU was requested must check the type of the *returned*
+      array rather than assuming the transfer succeeded.
 
     Example
     -------
@@ -231,6 +237,11 @@ def to_gpu(array: ArrayLike) -> Any:
     # Transfer to GPU
     try:
         return cp.asarray(array)
-    except Exception:
-        # Fallback to NumPy if GPU transfer fails
+    except Exception as e:
+        warnings.warn(
+            f"GPU transfer failed ({e}); falling back to CPU. "
+            "Fitting will run on CPU for this array even if GPU was requested.",
+            UserWarning,
+            stacklevel=2,
+        )
         return to_numpy(array)

osipy/common/fitting/base.py

@@ -148,9 +148,15 @@ def fit_image(
         # Extract masked voxel data: (nt, n_voxels)
         observed_masked = data[mask].T
 
-        # Transfer to GPU once (not per-chunk)
+        # Transfer to GPU once (not per-chunk). to_gpu() falls back to
+        # NumPy (with a warning) if the transfer fails, so use_gpu is
+        # re-derived from the actual returned array rather than trusted
+        # from the pre-transfer decision above — otherwise chunk sizing
+        # and the CPU threading fallback below would silently assume GPU
+        # execution while actually running (slower) single-threaded CPU.
         if use_gpu:
             observed_masked = to_gpu(observed_masked)
+            use_gpu = hasattr(observed_masked, "__cuda_array_interface__")
 
         xp = get_array_module(observed_masked)
 

tests/unit/common/backend/test_array_module.py

@@ -120,6 +120,38 @@ def test_preserves_values(self) -> None:
         result = to_gpu(data)
         np.testing.assert_array_almost_equal(to_numpy(result), data)
 
+    def test_transfer_failure_warns_and_falls_back(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """GH-175: a failed GPU transfer must warn (not be swallowed
+        silently) and still return a usable NumPy array."""
+        import osipy.common.backend.array_module as array_module
+
+        fake_cp = type(
+            "FakeCupy",
+            (),
+            {
+                "asarray": staticmethod(
+                    lambda arr: (_ for _ in ()).throw(
+                        RuntimeError("CUDA out of memory")
+                    )
+                ),
+                "ndarray": np.ndarray,
+            },
+        )()
+        monkeypatch.setattr(array_module, "_get_cupy", lambda: fake_cp)
+
+        original_config = get_backend()
+        try:
+            set_backend(GPUConfig(force_cpu=False))
+            data = np.array([1.0, 2.0, 3.0])
+            with pytest.warns(UserWarning, match="GPU transfer failed"):
+                result = to_gpu(data)
+            assert isinstance(result, np.ndarray)
+            np.testing.assert_array_equal(result, data)
+        finally:
+            set_backend(original_config)
+
 
 class TestGpuIntegration:
     """Integration tests that run if CuPy is available."""

tests/unit/common/fitting/test_base.py

@@ -0,0 +1,116 @@
+"""Unit tests for BaseFitter.fit_image() GPU/CPU dispatch.
+
+Tests for osipy.common.fitting.base — specifically the GH-175 regression:
+fit_image() must not keep treating a chunk as GPU-bound after to_gpu()
+has silently fallen back to CPU.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, ClassVar
+
+import numpy as np
+import pytest
+
+from osipy.common.backend.config import GPUConfig, get_backend, set_backend
+from osipy.common.fitting import base as fitting_base
+from osipy.common.fitting.base import BaseFitter
+
+
+class _FakeModel:
+    """Minimal FittableModel stand-in — just enough for fit_image/create_parameter_maps."""
+
+    name = "fake_model"
+    parameters: ClassVar[list[str]] = ["a"]
+    parameter_units: ClassVar[dict[str, str]] = {"a": ""}
+    reference = ""
+
+
+class _OneParamFitter(BaseFitter):
+    """Trivial fitter: returns zeros for every chunk, records call count."""
+
+    fitting_method_name = "fake"
+    chunk_size = 4
+
+    def __init__(self) -> None:
+        self.batch_calls = 0
+
+    def fit_batch(
+        self,
+        model: Any,
+        observed_batch: np.ndarray,
+        bounds_override: dict | None = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        self.batch_calls += 1
+        n_voxels = observed_batch.shape[1]
+        params = np.zeros((1, n_voxels))
+        r2 = np.ones(n_voxels)
+        converged = np.ones(n_voxels, dtype=bool)
+        return params, r2, converged
+
+
+@pytest.fixture
+def restore_backend():
+    """Restore the global backend config after each test."""
+    original = get_backend()
+    yield
+    set_backend(original)
+
+
+class TestFitImageGpuFallback:
+    """GH-175: use_gpu must be re-derived after to_gpu() may have fallen back."""
+
+    def test_gpu_transfer_failure_reenables_cpu_threading(
+        self, monkeypatch: pytest.MonkeyPatch, restore_backend, caplog
+    ) -> None:
+        """If to_gpu() falls back to NumPy, fit_image must use the CPU
+        multi-threaded path instead of silently running single-threaded.
+        """
+        # Pretend GPU is available and requested...
+        monkeypatch.setattr(fitting_base, "is_gpu_available", lambda: True)
+        # ...but to_gpu() falls back to plain NumPy (as it does on a real
+        # transfer failure, after emitting its own warning).
+        monkeypatch.setattr(fitting_base, "to_gpu", lambda arr: np.asarray(arr))
+
+        set_backend(GPUConfig(force_cpu=False, n_workers=2))
+
+        fitter = _OneParamFitter()
+        model = _FakeModel()
+        data = np.ones((4, 4, 1, 5))  # 16 voxels, chunk_size=4 -> 4 chunks
+
+        with caplog.at_level(logging.INFO, logger=fitting_base.logger.name):
+            result = fitter.fit_image(model, data)
+
+        assert "Using 2 threads for 4 chunks" in caplog.text
+        assert "a" in result
+        assert fitter.batch_calls == 4
+
+    def test_successful_gpu_transfer_keeps_gpu_path(
+        self, monkeypatch: pytest.MonkeyPatch, restore_backend, caplog
+    ) -> None:
+        """Sanity check: when to_gpu() *doesn't* fall back (returns an
+        object that looks like a GPU array), the CPU threading path must
+        NOT be used — this guards against an overly broad fix."""
+
+        class _FakeGpuArray(np.ndarray):
+            # Marks it as "on GPU" for the hasattr() check in fit_image().
+            __cuda_array_interface__: ClassVar[dict[str, Any]] = {}
+
+        def _fake_to_gpu(arr: np.ndarray) -> _FakeGpuArray:
+            return arr.view(_FakeGpuArray)
+
+        monkeypatch.setattr(fitting_base, "is_gpu_available", lambda: True)
+        monkeypatch.setattr(fitting_base, "to_gpu", _fake_to_gpu)
+        monkeypatch.setattr(fitting_base, "get_gpu_batch_size", lambda: 0)
+
+        set_backend(GPUConfig(force_cpu=False, n_workers=2))
+
+        fitter = _OneParamFitter()
+        model = _FakeModel()
+        data = np.ones((4, 4, 1, 5))
+
+        with caplog.at_level(logging.INFO, logger=fitting_base.logger.name):
+            fitter.fit_image(model, data)
+
+        assert "Using" not in caplog.text or "threads" not in caplog.text