diff --git a/py_boost/gpu/base.py b/py_boost/gpu/base.py index 6c3cf7c..d946a0f 100644 --- a/py_boost/gpu/base.py +++ b/py_boost/gpu/base.py @@ -4,7 +4,7 @@ try: import cupy as cp except Exception: - pass + cp = None import numpy as np from .serialization import dump, load @@ -37,30 +37,29 @@ def __init__(self): self.seed = 42 def to_device(self): - """Move trained ensemble data to current GPU device - - Returns: - - """ + """Move trained ensemble data to GPU, undoing a partial failed transfer.""" if not self._on_device: - for tree in self.models: - tree.to_device() - self.base_score = cp.asarray(self.base_score) + try: + for tree in self.models: + tree.to_device() + self.base_score = cp.asarray(self.base_score) + except Exception: + self.to_cpu() + raise self._on_device = True def to_cpu(self): - """Move trained ensemble data to CPU memory - - Returns: - - """ - if self._on_device: - for tree in self.models: + """Move all GPU arrays to CPU, including an incomplete device transfer.""" + # A failed transfer can leave GPU arrays while _on_device is still False. + # With no CuPy installed, CPU-only export/serialization needs no cleanup. + if cp is not None: + for tree in self.models or []: tree.to_cpu() - self.base_score = self.base_score.get() + if isinstance(self.base_score, cp.ndarray): + self.base_score = self.base_score.get() - self._on_device = False + self._on_device = False def __getstate__(self): """Get state dict on CPU for picking diff --git a/tests/test_device_transfer.py b/tests/test_device_transfer.py new file mode 100644 index 0000000..9081a0e --- /dev/null +++ b/tests/test_device_transfer.py @@ -0,0 +1,215 @@ +"""GPU transfer recovery must not retain live arrays after a failed upload.""" + +import pickle +import subprocess +import sys + +import numpy as np +import pytest + + +@pytest.fixture(scope='module') +def cp(): + cupy = pytest.importorskip('cupy') + try: + count = cupy.cuda.runtime.getDeviceCount() + except cupy.cuda.runtime.CUDARuntimeError as exc: + pytest.skip(f'CUDA device unavailable: {exc}') + if not count: + pytest.skip('A CUDA device is required') + return cupy + + +@pytest.fixture(scope='module', params=[False, True], ids=['compact', 'debug']) +def fitted(request, cp): + from py_boost import GradientBoosting + + rng = np.random.default_rng(12) + X = rng.normal(size=(128, 4)).astype(np.float32) + X[::11, 0] = np.nan + clean = np.nan_to_num(X) + y = np.column_stack([clean[:, 0] + .5 * clean[:, 1], clean[:, 2] - clean[:, 3]]) + model = GradientBoosting('mse', ntrees=3, max_depth=2, max_bin=16, + min_data_in_leaf=2, es=0, verbose=0, debug=request.param).fit(X, y) + expected = model.predict(X, batch_size=32) + model.to_cpu() + return X, expected, pickle.dumps(model) + + +@pytest.fixture +def model_data(fitted, cp): + X, expected, serialized = fitted + model = pickle.loads(serialized) + yield model, X, expected + # Also clean up when these regressions are run against the broken version. + for tree in model.models: + tree.to_cpu() + if isinstance(model.base_score, cp.ndarray): + model.base_score = model.base_score.get() + cp.get_default_memory_pool().free_all_blocks() + + +def gpu_array_counts(model, cp): + return [sum(isinstance(value, cp.ndarray) for value in vars(tree).values()) + for tree in model.models] + + +def assert_cpu(model, cp): + assert not model._on_device + assert isinstance(model.base_score, np.ndarray) + assert not any(gpu_array_counts(model, cp)) + + +@pytest.mark.parametrize('stage', ['first-tree', 'second-tree', 'base-score']) +@pytest.mark.parametrize('nondefault_stream', [False, True], ids=['default-stream', 'nondefault-stream']) +def test_upload_oom_rolls_back_and_allows_retry(model_data, cp, monkeypatch, stage, nondefault_stream): + model, X, expected = model_data + # Measure real pool allocation sizes instead of assuming allocator rounding. + probe = cp.cuda.MemoryPool() + sizes = [] + with cp.cuda.using_allocator(probe.malloc): + for tree in model.models: + tree.to_device() + sizes.append(probe.used_bytes()) + full_counts = gpu_array_counts(model, cp) + for tree in model.models: + tree.to_cpu() + probe.free_all_blocks() + + pool = cp.cuda.MemoryPool() + # One small allocation fits, then the next attribute fails. For the other + # stages, allow one whole tree or all trees before the failing allocation. + limit = {'first-tree': 512, 'second-tree': sizes[0] + 512, 'base-score': sizes[-1]}[stage] + pool.set_limit(size=limit) + asarray = cp.asarray + observed = {} + + def upload(array, *args, **kwargs): + try: + return asarray(array, *args, **kwargs) + except cp.cuda.memory.OutOfMemoryError as exc: + observed['counts'] = gpu_array_counts(model, cp) + observed['base_score'] = array is model.base_score + observed['error'] = exc + raise + + stream = cp.cuda.Stream(non_blocking=True) if nondefault_stream else cp.cuda.Stream.null + with monkeypatch.context() as patch: + patch.setattr(cp, 'asarray', upload) + with stream, cp.cuda.using_allocator(pool.malloc): + with pytest.raises(cp.cuda.memory.OutOfMemoryError) as caught: + model.to_device() + stream.synchronize() + + assert caught.value is observed['error'] + if stage == 'first-tree': + assert 0 < observed['counts'][0] < full_counts[0] + assert observed['counts'][1:] == [0, 0] + elif stage == 'second-tree': + assert observed['counts'][0] == full_counts[0] + assert 0 < observed['counts'][1] < full_counts[1] + assert observed['counts'][2] == 0 + else: + assert observed['counts'] == full_counts + assert observed['base_score'] + + # Check automatic rollback before any explicit to_cpu or garbage collection. + assert_cpu(model, cp) + assert pool.used_bytes() == 0 + pool.free_all_blocks() + assert pool.total_bytes() == 0 + model.to_cpu() + restored = pickle.loads(pickle.dumps(model)) + assert_cpu(restored, cp) + try: + np.testing.assert_array_equal(model.predict(X, batch_size=32), expected) + np.testing.assert_array_equal(restored.predict(X, batch_size=32), expected) + finally: + model.to_cpu() + restored.to_cpu() + + +@pytest.mark.parametrize('base_on_gpu', [False, True]) +@pytest.mark.parametrize('flag', [False, True]) +def test_to_cpu_recovers_mixed_state(model_data, cp, base_on_gpu, flag): + model, X, expected = model_data + pool = cp.cuda.MemoryPool() + with cp.cuda.using_allocator(pool.malloc): + model.models[0].values = cp.asarray(model.models[0].values) + if base_on_gpu: + model.base_score = cp.asarray(model.base_score) + model._on_device = flag + assert pool.used_bytes() > 0 + model.to_cpu() + model.to_cpu() + assert_cpu(model, cp) + assert pool.used_bytes() == 0 + pool.free_all_blocks() + np.testing.assert_array_equal(model.predict(X, batch_size=32), expected) + + +def test_pickle_recovers_partial_state(model_data, cp): + model, X, expected = model_data + pool = cp.cuda.MemoryPool() + with cp.cuda.using_allocator(pool.malloc): + model.models[0].values = cp.asarray(model.models[0].values) + model.base_score = cp.asarray(model.base_score) + assert not model._on_device + restored = pickle.loads(pickle.dumps(model)) + assert_cpu(model, cp) + assert_cpu(restored, cp) + assert pool.used_bytes() == 0 + pool.free_all_blocks() + try: + np.testing.assert_array_equal(restored.predict(X, batch_size=32), expected) + finally: + restored.to_cpu() + + +def test_upload_preserves_other_exceptions(model_data, cp, monkeypatch): + model, _, _ = model_data + failure = RuntimeError('transfer failed') + asarray = cp.asarray + + def upload(array, *args, **kwargs): + if array is model.base_score: + raise failure + return asarray(array, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(cp, 'asarray', upload) + with pytest.raises(RuntimeError) as caught: + model.to_device() + assert caught.value is failure + assert_cpu(model, cp) + + +def test_cpu_state_remains_serializable_without_cupy(): + # The CPU export/inference workflow can load an ensemble without CuPy. + # Deny imports in a fresh process so an installed CuPy cannot hide a regression. + script = r''' +import importlib.abc +import pickle +import sys +class NoCuPy(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'cupy' or fullname.startswith('cupy.'): + raise ModuleNotFoundError('CuPy is unavailable') +sys.meta_path.insert(0, NoCuPy()) +import numpy as np +from py_boost.gpu.base import Ensemble +from py_boost.gpu.tree import Tree +for trees in (None, [], [Tree(max_nodes=1, nout=1, ngroups=1)]): + model = Ensemble() + model.models = trees + if trees is not None: + model.base_score = np.array([.25], dtype=np.float32) + model.to_cpu() + restored = pickle.loads(pickle.dumps(model)) + restored.to_cpu() + assert not restored._on_device + np.testing.assert_equal(restored.base_score, model.base_score) +assert 'cupy' not in sys.modules +''' + result = subprocess.run([sys.executable, '-c', script], capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stdout + result.stderr