fix(agent): free the LLM's VRAM on the normal path after a workflow runs - #289
Open
kevin9327 wants to merge 1 commit into
Open
fix(agent): free the LLM's VRAM on the normal path after a workflow runs#289kevin9327 wants to merge 1 commit into
kevin9327 wants to merge 1 commit into
Conversation
agent_chat's VRAM-unload block sat *after* the tool-call `for` loop, so it only ran when all 10 iterations completed without an early return — i.e. the rare "Reached maximum tool iterations" path. In the normal flow (round 1 dispatches run_workflow, round 2 returns a final answer with no tool calls) the function returns from inside the loop, so the unload the comment promises — "so the workflow has full GPU memory" — never happened. On a single-GPU machine the LLM kept holding VRAM while the workflow tried to generate. Extract the unload into a helper and call it on both exits (the normal return and the loop-exhausted return) whenever a workflow was dispatched. It stays best-effort and only fires once the agent is done reasoning, so intermediate tool rounds still have the model loaded. Adds api/tests/test_agent_workflow_unload.py: asserts keep_alive:0 is sent on the normal post-workflow return (fails before, passes after) and not sent when no workflow ran. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
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.
What
The in-app assistant's LLM-VRAM unload never runs in the normal case. When the assistant dispatches a workflow it's supposed to free the Ollama model's VRAM so the workflow gets the full GPU — but that only happens on the rare "reached max iterations" path, so on a single-GPU machine the LLM keeps holding VRAM while the workflow tries to generate.
Why it triggers
In
agent_chat, the unload block sits after the tool-callforloop:The normal flow is: round 1 the model calls
run_workflow; round 2 it produces a final answer with no tool calls, so the functionreturns from inside the loop at theif not tool_callsbranch. The block after the loop is only reached when all 10 iterations complete without a final answer — the"Reached maximum tool iterations."case. So the documented VRAM unload is effectively dead code in real use.Fix
Extract the unload into
_unload_llm_after_workflow(...)and call it on both exits — the normal early return and the loop-exhausted return — when a workflow was dispatched. It stays best-effort (never fails the chat) and fires only once the agent has finished reasoning, so intermediate tool rounds still have the model loaded.Verification
api/tests/test_agent_workflow_unload.py(unittest +httpx.MockTransport, scripting the Ollama/api/chatturns and recording/api/generatecalls).test_llm_unloaded_on_the_normal_return_after_a_workflowfails before (no unload on the early-return path) and passes after;test_no_unload_when_no_workflow_was_dispatchedguards against unloading when nothing ran.python -m unittest discover -s testsinapi/(venv withfastapi+python-multipart+httpx): all pass, 0 failures.