Skip to content

Fix partial GPU upload cleanup after OOM - #41

Merged
dev-rinchin merged 1 commit into
sb-ai-lab:masterfrom
SergeevVladislav:fix/partial-gpu-upload-cleanup
Nov 9, 2026
Merged

dev-rinchin merged 1 commit into
sb-ai-lab:masterfrom
SergeevVladislav:fix/partial-gpu-upload-cleanup

Conversation

@SergeevVladislav

Copy link
Copy Markdown
Contributor

An exception partway through Ensemble.to_device() can leave live CuPy arrays in the trees while _on_device remains False. Both to_cpu() and pickle's __getstate__() then skip cleanup, so retaining the model also retains its GPU allocations.

Roll back failed uploads through to_cpu() and re-raise the transfer error. Make to_cpu() handle mixed CPU/GPU state independently of the flag, including a NumPy base_score and an uninitialized ensemble. Preserve CPU-only export and serialization when CuPy is unavailable.

Validation:

  • Added 25 regression cases; 22 fail on the base revision and all 25 pass with the patch. Real constrained-pool OOMs cover a partial first tree, a partial second tree, and the final base_score allocation, with compact/debug trees and default/non-default CUDA streams.
  • Verify zero live allocations in the constrained pool after rollback, preservation of the original exception, repeated cleanup, pickle recovery, and exactly unchanged predictions after retry. A fresh interpreter blocks CuPy imports to check CPU-only serialization.
  • Full suite: 444 passed, 1 skipped (Windows-specific loader test on Linux). Wheel build passed.

Tested on Linux, Python 3.12, CuPy 14.2.0, CUDA runtime 12.9, and an RTX 5090 32 GB. The bounded pools reproduce OOM without exhausting the device. Cached pool blocks are left to CuPy's normal allocator policy; the regression checks that the model releases live allocations.

@dev-rinchin dev-rinchin added this to the v0.6.0 milestone Nov 3, 2026
@dev-rinchin
dev-rinchin merged commit 4ae66cd into sb-ai-lab:master Nov 9, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants