diff --git a/osipy/common/backend/array_module.py b/osipy/common/backend/array_module.py index a289bd1..72f4271 100644 --- a/osipy/common/backend/array_module.py +++ b/osipy/common/backend/array_module.py @@ -31,10 +31,13 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING, Any import numpy as np +from osipy.common.exceptions import GPUTransferError + # NumPy 2.0 renamed trapz -> trapezoid. Ensure compatibility with NumPy <2.0. if not hasattr(np, "trapezoid"): np.trapezoid = np.trapz # type: ignore[attr-defined] # noqa: NPY201 @@ -42,6 +45,8 @@ if TYPE_CHECKING: from numpy.typing import ArrayLike, NDArray +logger = logging.getLogger(__name__) + # Cache the CuPy module to avoid repeated imports _cupy_module: Any = None _cupy_available: bool | None = None @@ -201,6 +206,21 @@ 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 a GPU is available, ``force_cpu`` is not set, and the transfer + itself fails (e.g. out of memory), the error is logged and + re-raised as a ``GPUTransferError`` rather than silently falling + back to CPU. A silent fallback here would let callers that decide + GPU-vs-CPU behavior (chunk sizing, threading) keep assuming GPU + execution when it never happened. Callers that want a graceful + CPU fallback on transfer failure (e.g. batch processing that + retries a batch on CPU after a GPU out-of-memory error) must catch + this exception themselves. + + Raises + ------ + GPUTransferError + If a GPU is available, ``force_cpu`` is not set, and the + transfer to GPU fails for any reason. Example ------- @@ -231,6 +251,11 @@ def to_gpu(array: ArrayLike) -> Any: # Transfer to GPU try: return cp.asarray(array) - except Exception: - # Fallback to NumPy if GPU transfer fails - return to_numpy(array) + except Exception as e: + msg = ( + f"GPU transfer failed ({e}). GPU was requested and detected, " + "so refusing to silently fall back to CPU; pass force_cpu=True " + "to run on CPU instead." + ) + logger.error(msg) + raise GPUTransferError(msg) from e diff --git a/osipy/common/backend/batch.py b/osipy/common/backend/batch.py index 09b0631..8ca9644 100644 --- a/osipy/common/backend/batch.py +++ b/osipy/common/backend/batch.py @@ -37,7 +37,7 @@ get_gpu_memory_info, is_gpu_available, ) -from osipy.common.exceptions import DataValidationError +from osipy.common.exceptions import DataValidationError, GPUTransferError if TYPE_CHECKING: from collections.abc import Callable @@ -238,8 +238,33 @@ def map( batches_processed += 1 + except GPUTransferError as e: + # to_gpu() raises this instead of silently falling back + # to CPU (see array_module.to_gpu). BatchProcessor + # deliberately wants to keep degrading gracefully to CPU + # when GPU memory fills mid-run, so handle it here. + if not self.auto_fallback: + raise + + warnings.warn( + f"GPU transfer failed at batch {batches_processed}: {e}. " + "Falling back to CPU processing.", + UserWarning, + stacklevel=2, + ) + fallback_occurred = True + use_gpu = False + + # Clear GPU memory + self._clear_gpu_memory() + + # Retry on CPU + results.append(func(batch)) + batches_processed += 1 + except Exception as e: - # Check if this is a GPU memory error + # Check if this is a GPU memory error raised by `func` + # itself (as opposed to the transfer, handled above). error_str = str(e).lower() is_memory_error = any( phrase in error_str diff --git a/osipy/common/exceptions.py b/osipy/common/exceptions.py index 1eacd19..188121b 100644 --- a/osipy/common/exceptions.py +++ b/osipy/common/exceptions.py @@ -89,6 +89,26 @@ class AIFError(OsipyError): pass +class GPUTransferError(OsipyError): + """Raised when transferring an array to GPU fails. + + Raised by ``to_gpu()`` when a GPU is available, ``force_cpu`` is not + set, and the transfer itself fails for any reason (e.g. out of + memory). Callers that want a graceful CPU fallback instead of + surfacing this error (e.g. batch processing retrying on CPU after a + GPU out-of-memory error) must catch it explicitly. + + Examples + -------- + >>> raise GPUTransferError("GPU transfer failed (CUDA out of memory)") + Traceback (most recent call last): + ... + osipy.common.exceptions.GPUTransferError: GPU transfer failed (CUDA out of memory) + """ + + pass + + class IOError(OsipyError): """Raised when file I/O operations fail. diff --git a/osipy/common/fitting/base.py b/osipy/common/fitting/base.py index 6b923f2..0db84b3 100644 --- a/osipy/common/fitting/base.py +++ b/osipy/common/fitting/base.py @@ -148,7 +148,9 @@ 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() raises instead of + # silently falling back to NumPy on transfer failure, so use_gpu + # stays accurate for the chunk sizing and threading decisions below. if use_gpu: observed_masked = to_gpu(observed_masked) diff --git a/tests/unit/common/backend/test_array_module.py b/tests/unit/common/backend/test_array_module.py index 5a6399e..1234e5b 100644 --- a/tests/unit/common/backend/test_array_module.py +++ b/tests/unit/common/backend/test_array_module.py @@ -20,6 +20,7 @@ get_backend, set_backend, ) +from osipy.common.exceptions import GPUTransferError class TestGetArrayModule: @@ -120,6 +121,41 @@ def test_preserves_values(self) -> None: result = to_gpu(data) np.testing.assert_array_almost_equal(to_numpy(result), data) + def test_transfer_failure_raises_instead_of_falling_back( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """GH-175: a failed GPU transfer must not be swallowed silently. + + With a GPU detected and force_cpu unset, a transfer failure must + raise GPUTransferError rather than silently returning a NumPy + array, since callers rely on use_gpu staying accurate for + chunk-sizing and threading decisions. + """ + 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.raises(GPUTransferError, match="GPU transfer failed"): + to_gpu(data) + finally: + set_backend(original_config) + class TestGpuIntegration: """Integration tests that run if CuPy is available.""" diff --git a/tests/unit/common/backend/test_batch.py b/tests/unit/common/backend/test_batch.py index fd6ed7c..e5f8ead 100644 --- a/tests/unit/common/backend/test_batch.py +++ b/tests/unit/common/backend/test_batch.py @@ -8,9 +8,10 @@ import numpy as np import pytest +from osipy.common.backend import batch as batch_module from osipy.common.backend.batch import BatchProcessor, BatchResult, batch_apply from osipy.common.backend.config import GPUConfig, set_backend -from osipy.common.exceptions import DataValidationError +from osipy.common.exceptions import DataValidationError, GPUTransferError class TestBatchResult: @@ -136,6 +137,52 @@ def test_uses_cpu_when_forced(self) -> None: set_backend(original_config) +class TestGpuTransferFallback: + """GH-175/GH-176: BatchProcessor must keep its deliberate CPU fallback + when a GPU transfer fails, even though to_gpu() itself now raises + instead of silently degrading.""" + + def test_falls_back_to_cpu_on_gpu_transfer_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A GPUTransferError from to_gpu() should be caught and the + batch retried on CPU when auto_fallback is True (the default).""" + monkeypatch.setattr(batch_module, "is_gpu_available", lambda: True) + + def _raise(arr: np.ndarray) -> np.ndarray: + raise GPUTransferError("GPU transfer failed (CUDA out of memory)") + + monkeypatch.setattr(batch_module, "to_gpu", _raise) + + processor = BatchProcessor(batch_size=50, use_gpu=True) + data = np.arange(100).astype(np.float64) + + with pytest.warns(UserWarning, match="GPU transfer failed"): + result = processor.map(data, lambda x: x**2) + + np.testing.assert_array_almost_equal(result.data, data**2) + assert result.fallback_occurred is True + assert result.used_gpu is False + + def test_reraises_gpu_transfer_error_when_auto_fallback_disabled( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """With auto_fallback disabled, the GPUTransferError must propagate + instead of being swallowed.""" + monkeypatch.setattr(batch_module, "is_gpu_available", lambda: True) + + def _raise(arr: np.ndarray) -> np.ndarray: + raise GPUTransferError("GPU transfer failed (CUDA out of memory)") + + monkeypatch.setattr(batch_module, "to_gpu", _raise) + + processor = BatchProcessor(batch_size=50, use_gpu=True, auto_fallback=False) + data = np.arange(100).astype(np.float64) + + with pytest.raises(GPUTransferError, match="GPU transfer failed"): + processor.map(data, lambda x: x**2) + + class TestBatchApply: """Tests for batch_apply convenience function.""" diff --git a/tests/unit/common/fitting/__init__.py b/tests/unit/common/fitting/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/common/fitting/test_base.py b/tests/unit/common/fitting/test_base.py new file mode 100644 index 0000000..f33bc9d --- /dev/null +++ b/tests/unit/common/fitting/test_base.py @@ -0,0 +1,118 @@ +"""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 when the GPU +transfer actually failed. to_gpu() now raises GPUTransferError instead of +silently falling back to NumPy, so fit_image() should let that error +propagate rather than continue on a wrong assumption about the device. +""" + +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.exceptions import GPUTransferError +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: fit_image must not keep assuming GPU after a failed transfer.""" + + def test_gpu_transfer_failure_propagates( + self, monkeypatch: pytest.MonkeyPatch, restore_backend + ) -> None: + """If to_gpu() fails, fit_image must raise rather than silently + continue on CPU while still believing it is on GPU. + """ + monkeypatch.setattr(fitting_base, "is_gpu_available", lambda: True) + + def _raise(arr: np.ndarray) -> np.ndarray: + raise GPUTransferError("GPU transfer failed (CUDA out of memory)") + + monkeypatch.setattr(fitting_base, "to_gpu", _raise) + + 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 pytest.raises(GPUTransferError, match="GPU transfer failed"): + fitter.fit_image(model, data) + + assert fitter.batch_calls == 0 + + def test_successful_gpu_transfer_keeps_gpu_path( + self, monkeypatch: pytest.MonkeyPatch, restore_backend, caplog + ) -> None: + """Sanity check: when to_gpu() succeeds (returns an object that + looks like a GPU array), the CPU threading path must NOT be + used.""" + + 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