A worker that has run a jax node keeps holding the GPU after the run: XLA
preallocates most of the VRAM at import and never releases it, so the next
process OOMs on preallocation while a warm idle worker sits on the card.
Pools are kept warm on purpose — a library reads its environment at import,
so a warm worker cannot be re-told — but the end of a run is a point where
the memory should go back, and the environments carrying a GPU assignment
are exactly the pools that ran on one. Idle ones go now, busy ones when
they return.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019Hra4ndWMCLU5F3KjUuVAc
stroblmeandClaude Opus 5
53b49e5f68