Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

feat(studio): run chats in parallel in the Chat tab#7455

Open
danielhanchen wants to merge 5 commits into
mainunslothai/unsloth:mainfrom
feat/parallel-chatsunslothai/unsloth:feat/parallel-chatsCopy head branch name to clipboard
Open

feat(studio): run chats in parallel in the Chat tab#7455
danielhanchen wants to merge 5 commits into
mainunslothai/unsloth:mainfrom
feat/parallel-chatsunslothai/unsloth:feat/parallel-chatsCopy head branch name to clipboard

Conversation

@danielhanchen

@danielhanchen danielhanchen commented Jul 26, 2026

Copy link
Copy Markdown
Member

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 run already serves with --parallel 4, but plain unsloth studio launched 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

  • New Chat leaves the previous conversation generating. Its runtime stays mounted, it keeps streaming into its own thread, and its sidebar row spins.
  • Every row in Recents gets its own spinner. Off the Chat route the nav item spins and reads "Return to N chats", the same way Train and Export already work.
  • Stop, tool calls, uploads, nudging and self healing are isolated per conversation. Stopping one closes only that upstream stream. Siblings keep decoding and llama-server is never signalled.
  • Changing the model, or any setting that needs Apply, still ends every running chat, because they all decode on one llama-server. The UI now says so first and names the chats, and the backend refuses the swap with 409 unless the caller opts in.

Backend

  • New studio/backend/state/active_generations.py: a registry of in-flight generations. _TrackedCancel registers into it, so every streaming path is covered without its own bookkeeping. Entries hold the same threading.Event as the existing per-run cancel registry, so cancelling reuses the per-request path instead of inventing a second teardown.
  • /load and /unload gate on it. The load path checks twice, the second time under inference_lifecycle_gate(), since a chat can start while the request queues. force_cancel_active on the request cancels instead of raising. The frontend guard is bypassable from a second tab or curl; this one is not.
  • New GET /api/inference/active-generations returns the running conversations plus the slot count actually in use, which restores the sidebar spinners after a page reload.
  • _PARALLEL_DEFAULT_PLAIN goes 1 to 4 in studio/backend/run.py and unsloth_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

  • New Chat and thread switches stop the prompt queue without cancelling the prompt already dispatched.
  • A per-run server cancel handle, so Stop reaches a conversation generating off screen. cancelByThreadId only ever holds the visible thread's cancelRun(), which is not enough once background chats exist (deleting one, or a confirmed reload).
  • Tool output store keys are scoped per conversation. Local GGUF names every tool call call_0, so two chats mid tool call shared one store entry.
  • The composer's "Running Python: ..." badge is scoped per conversation too. It read one global string, so a chat mid tool call put its badge above every other composer, including a brand new empty chat. Its elapsed counter also restarted on every thread switch, and a run ending anywhere cleared the badge everywhere. The status is now keyed by thread and carries the moment it started, so the counter resumes instead of restarting.
  • Confirmation dialog before a model swap, listing the chats it would stop, with "Keep generating" and "Stop and reload".

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.

  • A gated call did not stream its arguments, 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, since its card renders the payload.
  • The status read "Running ..." with a climbing timer while the call had not started. It now says 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.

It also fixes a rendering bug that parallel chats surfaced: assistant-ui keys message parts by index, so switching conversations hands the same MarkdownText instance 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. Streamdown is 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_workdir puts 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.

  • The card was rendered twice. The provisional card that streams the arguments is keyed by the backend tool call id, but the tool_start that opens the approval prompt was keyed by the approval id, so it opened a second card. Only the second one ever received tool_end, leaving the first spinning "Running" for the rest of the conversation. The approval prompt now reuses the card already on screen.
  • The terminal card had no body at all, only a 60-character trigger label, so a long heredoc showed a truncated first line and nothing else for as long as it took to write. It now renders the command the way the Python card renders its script.
  • Neither cell is capped any more. The Python script used to stop at 10k characters with a "... (truncated)" marker. Both render in full; above 20k characters the block stays plain monospace rather than paying for shiki, and it stays plain while the model is still writing so a long payload does not re-tokenize on every fragment.

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_XL on one GPU:

  • Four conversations streaming at once, backend reporting count: 4, four sidebar spinners.
  • New Chat issued no /api/inference/cancel. Chat 1 had 3062 characters when it was left and 16225 on return, still streaming.
  • Stop on chat 1 took the active count from 4 to 3 with the other three still streaming, and the llama-server PID was 355324 before and after.
  • Reload guard: /load and /unload return 409 naming the running thread, the chat keeps streaming through the refusal, and the forced load then ends it.
  • Two concurrent tool-calling chats end with the right answer in each, live, after switching between them and after a reload, 3 runs out of 3.
  • Tool badge: shown in the chat that owns the run, absent in a new empty chat opened next to it, and still counting on return (1s before leaving, 12s after).
  • Sandboxes: four concurrent tool calls landed in four directories named after their threads, and none could see a sibling's file. Repeated through the UI with two conversations, reading the paths out of the tools' own stdout.
  • Approvals: with four prompts left open, every gated call streamed its code, none reported "Running", and a fresh chat answered in 0.4s. Before the change the same request waited 290s and never answered while the GPU was idle.
  • Tool cards, through the UI on a gated terminal call and a gated python call: 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_arguments and two admission-queue cases covering a parked slot being freed and unpark never 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.
  • Full studio backend suite passes, with the same pre-existing failures as main (test_studio_api.py wants a live keyed server, test_training_worker_flash_attn.py wants network installs).
  • The registry tests also run green on ubuntu-latest, windows-latest and macos-14. Nothing added here is POSIX specific: it is a dict plus a threading.Lock, and cancelling is an HTTP connection close.

danielhanchen and others added 2 commits July 26, 2026 05:04
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.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +5651 to +5653
_raise_or_cancel_active_generations(
force = request.force_cancel_active, action = "Unloading the model"
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +4312 to +4313
@studio_router.get("/active-generations")
async def get_active_generations(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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]);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +1265 to +1269
// 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 });

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment thread studio/backend/routes/inference.py Outdated
Comment on lines +4360 to +4362
_raise_or_cancel_active_generations(
force = request.force_cancel_active, action = "Loading a model"
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

setThreadRunning: (threadId, running) =>
set((state) => {
const next = { ...state.runningByThreadId };
if (running) {
next[threadId] = true;
} else {
delete next[threadId];

P2 Badge Reference-count overlapping runs before clearing thread activity

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

Comment on lines +2797 to +2801
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"), {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +29 to +33
const running = Object.entries(
useChatRuntimeStore.getState().runningByThreadId,
)
.filter(([, on]) => on)
.map(([threadId]) => threadId);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +292 to +296
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +4191 to +4193
// Passing serverCancel keeps a newer run's handle: a tool continuation
// can register the next leg before this one unwinds.
runtime.clearThreadServerCancel(threadKey, serverCancel);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +537 to +539
const stopDecision = await confirmStopRunningChatsIfNeeded(
forceReload ? "Applying these settings" : "Loading a different model",
);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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".

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +4300 to +4301
if not active_generations.count():
return 0

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +5664 to +5666
_raise_or_cancel_active_generations(
force = request.force_cancel_active, action = "Unloading the model"
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment thread studio/backend/run.py
_PARALLEL_MIN = 1
_PARALLEL_MAX = 64
_PARALLEL_DEFAULT_PLAIN = 1
_PARALLEL_DEFAULT_PLAIN = 4

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

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

Morty Proxy This is a proxified and sanitized view of the page, visit original site.