feat(studio): run chats in parallel in the Chat tab#7455
feat(studio): run chats in parallel in the Chat tab#7455danielhanchen wants to merge 5 commits into
Conversation
New Chat used to cancel whatever the current conversation was generating. It now leaves it running, like switching to the Train or Export tab: the sidebar shows which chats are still going, and Stop is per conversation. Plain `unsloth studio` launched llama-server with one decode slot, so the admission queue serialised every chat regardless of what the UI did. Both entry points now default to the same slot count as `unsloth studio run`. A model swap still ends every running chat, since they all decode on one llama-server. /load and /unload now refuse with 409 and name those chats unless the caller passes force_cancel_active, and the UI asks first.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e87162cb0d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _raise_or_cancel_active_generations( | ||
| force = request.force_cancel_active, action = "Unloading the model" | ||
| ) |
There was a problem hiding this comment.
Recheck active generations after acquiring the unload gate
When a chat starts after this check but before inference_lifecycle_gate() is acquired, the inference middleware can register and release that gate first, after which this route unloads the model at lines 5698–5721 while the new generation is still streaming. /load explicitly repeats the check under the gate, but /unload does not; perform the same in-gate check or drain before tearing down the backend.
Useful? React with 👍 / 👎.
| @studio_router.get("/active-generations") | ||
| async def get_active_generations( |
There was a problem hiding this comment.
Hydrate frontend run state from active-generations
A repo-wide search finds no frontend caller for this new endpoint, while confirmStopRunningChatsIfNeeded consults only the in-memory Zustand map. After a page reload, or when another tab owns the generation, that map is empty, so model load/unload is sent without force_cancel_active, receives the backend 409, and never presents the confirmation dialog or restores the promised sidebar controls; consume this endpoint during hydration or handle the structured 409 with a confirm-and-retry flow.
Useful? React with 👍 / 👎.
| keepalive: true, | ||
| }).catch(() => {}); | ||
| }; | ||
| runtime.registerThreadServerCancel(threadKey, serverCancel); |
There was a problem hiding this comment.
Register cancellation for audio-only generations
For an audio-only model, the adapter returns from the branch ending at line 2554 before reaching this new server-cancel registration. Once that audio generation is backgrounded, CancelRegistrar no longer supplies a visible-thread cancel and stopChatThread has no server handle, so deleting the chat or choosing “Stop and reload” cannot stop it; the backend /audio/generate path also does not enter active_generations, allowing a forced unload to tear down the model while audio is still being produced. Register and track this branch before its early return.
Useful? React with 👍 / 👎.
| const isPinned = pinnedIdSet.has(item.id); | ||
| // A row can be generating while you look at another one, so mirror the | ||
| // Train / Export nav spinners instead of running silently. | ||
| const isGenerating = Boolean(runningThreadIds[item.id]); |
There was a problem hiding this comment.
Derive compare-row activity from its member threads
For compare entries, groupThreads sets item.id to the pair ID, but runningByThreadId is keyed by each pane's individual thread ID. Consequently this lookup is always false for a running compare conversation, so the newly added per-row spinner and aria-busy state disappear as soon as that compare is viewed from another chat; aggregate activity from the threads whose pairId matches the row instead.
Useful? React with 👍 / 👎.
| // Stop feeding the outgoing thread queued prompts, but leave its run | ||
| // alone: assistant-ui keeps every visited thread's runtime mounted, so | ||
| // it keeps streaming and the sidebar shows it running. Only an explicit | ||
| // Stop (or deleting the conversation) cancels one now. | ||
| requestPromptQueueStop({ cancelActiveRun: false }); |
There was a problem hiding this comment.
Scope transient generation state to its originating thread
Leaving the outgoing run alive also leaves it writing thread-agnostic Zustand fields such as toolStatus, generatingStatus, activeDiffusionCanvas, and contextUsage. If chat A continues after switching to chat B, A's later status and usage events are rendered in B, and A's finally can clear status or diffusion state belonging to B; the usage guard checks only the model, not the active thread. Scope these fields by thread or ignore writes from runs whose captured thread is no longer active before enabling background execution.
Useful? React with 👍 / 👎.
| _raise_or_cancel_active_generations( | ||
| force = request.force_cancel_active, action = "Loading a model" | ||
| ) |
There was a problem hiding this comment.
Preserve idempotent loads while chats are active
This gate runs before _load_model_impl reaches its exact-settings already_loaded checks at lines 4455–4528. Therefore an idempotent POST /load now returns 409 during any generation even though it would not replace the server, and retrying that no-op request with the instructed force_cancel_active=true cancels every chat before returning already_loaded; determine whether the request is actually a reload before refusing or cancelling active work.
Useful? React with 👍 / 👎.
The green "Running Python: ..." badge above the composer read a single global store value, so one chat's tool call showed above every other chat's composer, including a brand-new empty one. Its elapsed counter also restarted at 0 on every thread switch, and a run ending anywhere cleared the badge everywhere. Key the status by thread and store the moment it started, so each conversation shows only its own tool call and the counter resumes rather than restarts. Also adds a test that every conversation gets its own tool sandbox directory, which parallel tool calls depend on.
There was a problem hiding this comment.
💡 Codex Review
unsloth/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
Lines 1490 to 1496 in 4c264f3
When a tool continuation starts its next run before the previous adapter invocation finishes unwinding—a case the new backend test and the adjacent server-cancel cleanup explicitly support—both runs call setThreadRunning(threadId, true), but the first run to finish unconditionally deletes the shared boolean entry. The still-active continuation then loses its sidebar spinner, prompt-queue busy state, reload confirmation, and even per-thread Stop because stopChatThread returns early when this entry is absent. Track active run identities or a per-thread count and only remove the thread after its final run exits.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const serverCancel = () => { | ||
| const body: Record<string, string> = { cancel_id: cancelId }; | ||
| if (sandboxSessionId) body.session_id = sandboxSessionId; | ||
| const token = getAuthToken(); | ||
| void fetch(apiUrl("/api/inference/cancel"), { |
There was a problem hiding this comment.
Cancel background external-provider streams directly
When an external-provider chat is backgrounded and the user archives it or confirms “Stop and reload,” this handle only posts its cancel_id to the local cancel endpoint. The external-provider branch returns through _proxy_to_external_provider before creating any _TrackedCancel entry (studio/backend/routes/inference.py), so the POST cannot reach or close the upstream request; because the off-screen thread has no cancelRun() handle, the paid provider stream can continue despite the UI claiming it was stopped. Register an external-stream cancellation event or retain a handle that aborts that run directly.
Useful? React with 👍 / 👎.
| const running = Object.entries( | ||
| useChatRuntimeStore.getState().runningByThreadId, | ||
| ) | ||
| .filter(([, on]) => on) | ||
| .map(([threadId]) => threadId); |
There was a problem hiding this comment.
Exclude external chats from the local model-swap gate
When an external-provider conversation is generating in the background, this unqualified map includes it and blocks a local model load behind a “Stop running chats” dialog. External requests do not use the local llama-server—the backend explicitly untracks them before proxying—so choosing “Keep generating” unnecessarily prevents a safe local load, while confirming falsely says that the external run must stop. Track whether each run is local and only include generations that use the backend being replaced.
Useful? React with 👍 / 👎.
Three problems, all from the approval prompt behaving as though only one chat could ever run. Arguments were not streamed for a gated call, so the chat stayed blank for as long as the model took to write the payload, which for a large file is minutes. Nothing runs before the decision either way, and the code is what is being approved, so python and terminal now stream their card while gated. render_html stays suppressed: its card renders the payload. The status read "Running ..." with a climbing timer while the call had not started. It now reports that it is waiting for approval, then switches to running once allowed. The admission lease was held across the wait, so four unanswered prompts held all four decode slots and no other chat could start while llama-server sat idle. A parked run keeps its lease but no longer counts against capacity. Measured with four prompts left open: every gated call streamed its code, none reported running, and a fresh chat answered in 0.4s where it previously waited 290s and never did.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8cf8ab2bb2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def unpark(self) -> None: | ||
| """Take the slot back. Pairs with park(); safe to call unmatched.""" | ||
| with self._lock: | ||
| if self._parked > 0: | ||
| self._parked -= 1 |
There was a problem hiding this comment.
Reacquire capacity before unparking approval waiters
When a parked generation grants its freed slot to another chat, unpark() immediately decrements _parked without waiting for that chat to release its lease. With capacity 1, this produces _active == 2 and _parked == 0, so an approved quick tool can resume decoding concurrently with the admitted chat despite the configured one-slot VRAM limit, risking OOMs or failed streams. The parked generation must reacquire admission capacity before it resumes.
Useful? React with 👍 / 👎.
| // Passing serverCancel keeps a newer run's handle: a tool continuation | ||
| // can register the next leg before this one unwinds. | ||
| runtime.clearThreadServerCancel(threadKey, serverCancel); |
There was a problem hiding this comment.
Preserve running state when older legs unwind
When a tool continuation registers its next leg before the previous leg unwinds—the overlap this cleanup explicitly handles—the new leg has already set runningByThreadId[threadKey] to true, but the old leg unconditionally clears it immediately above this guarded server-cancel cleanup. The newer cancel handle survives while the thread is incorrectly reported idle, so its sidebar spinner disappears and stopChatThread refuses to stop it; model-swap confirmation can also be skipped and then fail with the backend 409. Protect the running flag with the same per-run identity or a reference count.
Useful? React with 👍 / 👎.
| const stopDecision = await confirmStopRunningChatsIfNeeded( | ||
| forceReload ? "Applying these settings" : "Loading a different model", | ||
| ); |
There was a problem hiding this comment.
Reconfirm active chats immediately before forcing the swap
If validation, a transformers-upgrade prompt, or the remote-code consent dialog delays the load after this confirmation, forceCancelActive remains true until the later unload/load requests. A new generation started meanwhile—especially from another tab—was not among the chats the user approved stopping, but the backend interprets the stale flag as permission to cancel every generation active at request time and silently terminates it. Perform the confirmation immediately before the destructive request or bind the authorization to the generation handles that were shown.
Useful? React with 👍 / 👎.
A gated tool call rendered two cards: the provisional one that streams the arguments, plus a second one keyed by the approval id. Only the second ever got its tool_end, so the first spun "Running" for the rest of the chat. Reuse the open part when the approval prompt arrives. The terminal card also showed nothing but a 60-char trigger label, so a long heredoc read as no progress at all. It now renders the command the same way the Python card renders its script, and neither is capped at 10k chars. Both cells moved inside the collapsible, so one chevron hides the code with the output and Copy / Download exist only while the card is open. A card parked on the prompt says so instead of counting up "Running".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 19a4fe941f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not active_generations.count(): | ||
| return 0 |
There was a problem hiding this comment.
Track direct Responses streams in the unload registry
When a client is streaming /v1/responses through the direct GGUF _responses_stream path, that path never creates a _TrackedCancel, so active_generations.count() remains zero. A concurrent non-forced /unload therefore bypasses this new safety gate and tears down llama-server while the Responses stream is still active; register the direct Responses stream for its lifetime just like the Anthropic and chat passthrough streams.
Useful? React with 👍 / 👎.
| _raise_or_cancel_active_generations( | ||
| force = request.force_cancel_active, action = "Unloading the model" | ||
| ) |
There was a problem hiding this comment.
Run safe loading-cancel checks before gating unloads
When /unload is used to cancel a model that is still loading while a chat remains active on the previous model, this unconditional check returns 409 before reaching the safe cancel_load fast paths below. The frontend's cancelLoading() then swallows that error and still reports “Stopped loading model,” leaving the backend load running; move the active-generation gate after the loading-cancel and nonresident no-op checks, since those operations do not unload the model serving the active chat.
Useful? React with 👍 / 👎.
| _PARALLEL_MIN = 1 | ||
| _PARALLEL_MAX = 64 | ||
| _PARALLEL_DEFAULT_PLAIN = 1 | ||
| _PARALLEL_DEFAULT_PLAIN = 4 |
There was a problem hiding this comment.
Apply the parallel default to programmatic server launches
This constant only changes the argparse entry point, while run_server() still defaults llama_parallel_slots to 1. The Colab launcher in studio/backend/colab.py calls run_server() without supplying that argument, so chats there remain serialized despite the new default and parallel-chat UI; make the programmatic default match this value as well.
Useful? React with 👍 / 👎.
Clicking New Chat used to cancel whatever the current conversation was generating. This makes the Chat tab behave like the rest of Studio: leaving a conversation lets it keep going, the sidebar shows which chats are still running, and Stop is per conversation.
unsloth runalready serves with--parallel 4, but plainunsloth studiolaunched llama-server with a single decode slot, so the admission queue serialised every chat regardless of what the UI did. Both entry points now default to the same slot count.Behaviour
Backend
studio/backend/state/active_generations.py: a registry of in-flight generations._TrackedCancelregisters into it, so every streaming path is covered without its own bookkeeping. Entries hold the samethreading.Eventas the existing per-run cancel registry, so cancelling reuses the per-request path instead of inventing a second teardown./loadand/unloadgate on it. The load path checks twice, the second time underinference_lifecycle_gate(), since a chat can start while the request queues.force_cancel_activeon the request cancels instead of raising. The frontend guard is bypassable from a second tab or curl; this one is not.GET /api/inference/active-generationsreturns the running conversations plus the slot count actually in use, which restores the sidebar spinners after a page reload._PARALLEL_DEFAULT_PLAINgoes 1 to 4 instudio/backend/run.pyandunsloth_cli/commands/studio.py, with a test pinning the two together. Loads that do not fit in VRAM still fall back to fewer slots.Frontend
cancelByThreadIdonly ever holds the visible thread'scancelRun(), which is not enough once background chats exist (deleting one, or a confirmed reload).call_0, so two chats mid tool call shared one store entry.Tool approvals
Three problems surfaced by running several chats at once, all from the approval prompt behaving as though only one chat could ever exist.
pythonandterminalnow stream their card while gated.render_htmlstays suppressed, since its card renders the payload.It also fixes a rendering bug that parallel chats surfaced: assistant-ui keys message parts by index, so switching conversations hands the same
MarkdownTextinstance a different message, and Streamdown keeps its parsed blocks in internal state and only extends them while streaming. A conversation could show the previously viewed one's answer until the page was reloaded.Streamdownis now keyed per message.Tool sandboxes were already per conversation and stay that way: the session id a chat sends is its thread id, so
_get_workdirputs each one in its own~/studio_sandbox/<thread id>at mode 0700. Chats inside a project deliberately share the project workspace instead.Tool cards
Streaming a gated call's arguments exposed three more problems in the card itself.
tool_startthat opens the approval prompt was keyed by the approval id, so it opened a second card. Only the second one ever receivedtool_end, leaving the first spinning "Running" for the rest of the conversation. The approval prompt now reuses the card already on screen.Both cells moved inside the collapsible, so one chevron now hides the code together with the output, and Copy / Download exist only while the card is open. A card parked on the prompt says so rather than counting up "Running", which the composer badge already did.
Verification
Playwright against
unsloth/gemma-4-E2B-it-GGUF:UD-Q4_K_XLon one GPU:count: 4, four sidebar spinners./api/inference/cancel. Chat 1 had 3062 characters when it was left and 16225 on return, still streaming./loadand/unloadreturn 409 naming the running thread, the chat keeps streaming through the refusal, and the forced load then ends it.terminalcall and a gatedpythoncall: one card per call rather than two, the whole command and script visible with their tail sentinel, the body reading "Waiting for approval" and never "Running", Copy and Download present while open and gone when collapsed and back on reopen, nothing left spinning at the end, and no "(truncated)" marker anywhere in the transcript. 8 of 8 checks on both tools.Tests:
studio/backend/tests/test_active_generations.py, 22 cases: registry lifecycle, overlapping runs on one conversation, JSON safety of the snapshot, per-thread against global cancel, an 8 thread concurrency stress, the 409 and force paths, and the parallel default agreement between the CLI and the backend.test_gated_python_call_still_streams_its_argumentsand two admission-queue cases covering a parked slot being freed andunparknever leaking capacity.studio/backend/tests/test_tool_sandbox_per_thread.py, 11 cases: a directory per conversation, stable across turns, private to its own chat, the project workspace exception, session ids that try to escape the sandbox root, and the 0700 mode. HOME is redirected to a tmp dir so nothing touches the real sandbox tree.main(test_studio_api.pywants a live keyed server,test_training_worker_flash_attn.pywants network installs).ubuntu-latest,windows-latestandmacos-14. Nothing added here is POSIX specific: it is a dict plus athreading.Lock, and cancelling is an HTTP connection close.