Skip to content

fix(agent): free the LLM's VRAM on the normal path after a workflow runs - #289

Open
kevin9327 wants to merge 1 commit into
lightningpixel:devfrom
kevin9327:fix/agent-workflow-vram-unload
Open

fix(agent): free the LLM's VRAM on the normal path after a workflow runs#289
kevin9327 wants to merge 1 commit into
lightningpixel:devfrom
kevin9327:fix/agent-workflow-vram-unload

Conversation

@kevin9327

Copy link
Copy Markdown
Contributor

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-call for loop:

for _ in range(10):
    ...
    if not tool_calls:
        return AgentChatResponse(...)   # normal exit — returns from *inside* the loop
    for tc in tool_calls:
        ... execute_tool ...

has_workflow = any(a.tool == "run_workflow" for a in actions_done)   # only reached if the loop runs out
if has_workflow:
    await client.post(.../api/generate, json={"keep_alive": 0})       # "so the workflow has full GPU memory"

The normal flow is: round 1 the model calls run_workflow; round 2 it produces a final answer with no tool calls, so the function returns from inside the loop at the if not tool_calls branch. 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

  • Added api/tests/test_agent_workflow_unload.py (unittest + httpx.MockTransport, scripting the Ollama /api/chat turns and recording /api/generate calls). test_llm_unloaded_on_the_normal_return_after_a_workflow fails before (no unload on the early-return path) and passes after; test_no_unload_when_no_workflow_was_dispatched guards against unloading when nothing ran.
  • python -m unittest discover -s tests in api/ (venv with fastapi + python-multipart + httpx): all pass, 0 failures.

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>
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.

1 participant