fix: stop active pipeline workers before unload to free VRAM#966
Open
fix: stop active pipeline workers before unload to free VRAM#966
Conversation
When /pipeline/load triggered a swap while a PipelineProcessor worker
was still producing frames, the unload path dropped pipeline_manager's
reference and called gc.collect()/empty_cache(), but the worker thread
kept the pipeline object alive through its closure and continued
allocating CUDA memory. The next load (e.g. longlive after ltx2) OOMed
with ~30 GiB still in use despite logging "CUDA cache cleared".
Add a pre-unload hook registry on PipelineManager. graph_executor
registers each processor's stop() under its node_id at creation time,
and FrameProcessor.stop() unregisters on normal teardown. The hook
fires synchronously inside _unload_pipeline_by_id_unsafe BEFORE the
pipeline reference is dropped, so the worker exits and releases its
tensors first — then gc/empty_cache can actually reclaim VRAM.
Verified: loading ltx2, running a session, then POSTing
/pipeline/load {longlive, passthrough} without a session stop now
succeeds. Log sequence is Unloading → PipelineProcessor stopped →
CUDA cache cleared, and the next session starts cleanly.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Rafal Leszko <rafal@livepeer.org>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
🚀 fal.ai Preview Deployment
Livepeer Runner
Testing Livepeer Mode |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When /pipeline/load triggered a swap while a PipelineProcessor worker was still producing frames, the unload path dropped pipeline_manager's reference and called gc.collect()/empty_cache(), but the worker thread kept the pipeline object alive through its closure and continued allocating CUDA memory. The next load (e.g. longlive after ltx2) OOMed with ~30 GiB still in use despite logging "CUDA cache cleared".
Add a pre-unload hook registry on PipelineManager. graph_executor registers each processor's stop() under its node_id at creation time, and FrameProcessor.stop() unregisters on normal teardown. The hook fires synchronously inside _unload_pipeline_by_id_unsafe BEFORE the pipeline reference is dropped, so the worker exits and releases its tensors first — then gc/empty_cache can actually reclaim VRAM.
Verified: loading ltx2, running a session, then POSTing /pipeline/load {longlive, passthrough} without a session stop now succeeds. Log sequence is Unloading → PipelineProcessor stopped → CUDA cache cleared, and the next session starts cleanly.