From 386c1c4def3b21e8f201bfd93f022f92e8a05aa4 Mon Sep 17 00:00:00 2001 From: iris-ai <4+iris-ai@noreply.localhost> Date: Mon, 21 Sep 2026 03:16:45 -0400 Subject: [PATCH] Pause a llama turn at once, and let Stop take the model with it A turn waits on three things that look nowhere at all: a permission question, a tool call with a minute to run, and the completion itself, which says nothing while the prompt is read -- tens of seconds on a long conversation. Setting a flag left the turn exactly where it was until whichever it was came back. The wait is now what ends, not the work. Each of those runs on a thread of its own and the interrupt answers the wait; the abandoned thread finishes into a channel nobody is reading. 43ms to end a turn in every state, measured against a real model -- including mid prompt-processing, which used to be a twenty-second wait. That makes cancellation a token per turn rather than a flag on the session: the abandoned thread wakes up some time later, and a flag the next turn had reset would let it write into a conversation it is no longer part of. The open thinking block moves to Shared for the same reason -- the thread that knows one is open is no longer the thread that ends the turn. Stop, meanwhile, did nothing at all to a llama session: it signals the session's recorded process and process::stop refuses a Shared one, which is the whole point of that record -- so the session sat at idle. A Shared record routes to the driver now, because what stopping means for a session that borrows the machine's process is the driver's to say. It ends the turn, says exited itself, and gives up its claim on the model; each live session claims the model it is on, and the model is unloaded when the last claim goes. A model another session is using stays where it is. --- AGENTS.md | 38 ++- PLAN.md | 24 ++ TODO.md | 27 ++ server/src/session/llama/mod.rs | 400 ++++++++++++++++++++--------- server/src/session/llama/router.rs | 33 +++ server/src/session/mod.rs | 34 ++- 6 files changed, 419 insertions(+), 137 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5b593d0..fd3eee8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -511,14 +511,36 @@ written, and the fold uses that same predicate to decide a reply is settled. reply, and the transcript keeps whatever arrived before it. - **A cancel flag is only as prompt as the next place somebody looks.** A - llama turn waits on two things that look nowhere at all: a permission - question, and a tool call `llama-server` is running -- a shell command there - runs to its own timeout, up to a minute. Setting `cancel` left the turn - exactly where it was until that came back, so Pause did nothing on screen - for as long as the command took. `Shared::abandon_turn` sets the flag *and* - releases both waits, the tool call by answering its channel with - `tools::UNFINISHED` -- the call is left running over there, because nothing - in that protocol takes one back, and its answer is dropped. + llama turn waits on three things that look nowhere at all: a permission + question, a tool call `llama-server` is running (a shell command there runs + to its own timeout, up to a minute), and the completion itself, which says + nothing for as long as the prompt takes to read -- tens of seconds on a long + conversation. Setting `cancel` left the turn exactly where it was until + whichever it was came back, so Pause did nothing on screen for all of it. + **The wait is what ends, not the work**: `awaiting` runs each of those on a + thread of its own and `Shared::abandon_turn` answers the wait, so the turn + ends in milliseconds (measured 43ms in every state) and the abandoned thread + finishes into a channel nobody is reading. Nothing here can stop a shell + command or a model mid-reply, and pretending otherwise is what the old code + did. + **Which is why cancellation is a token per turn** (`Cancel`), not a flag on + the session: the abandoned thread wakes up some time later, and a flag the + next turn had reset would let it write into a conversation it is no longer + part of. Its own token stays set for ever, so it says nothing -- and the + open thinking block is closed by `Shared::abandon_turn` rather than by that + thread, since the one that knows is not the one that ends the turn. + +- **Stop ends a llama session and takes the model with it, if nobody else + wants it** (2026-09-21). Until then Stop did *nothing at all* to one: + `stop_session` signals the session's recorded process, and `process::stop` + refuses a `Shared` record -- so the session sat at `idle` with no sign + anything had happened. A `Shared` record now routes to `Driver::stop`, + because what stopping means for a session that borrows somebody else's + process is the driver's to say. The llama driver ends the turn, says + `exited` itself (nothing else will -- there is no process of its own to + die), and asks `Router::release`: each live session claims the model it is + on, and the model comes out of memory only when the last claim goes. A model + another session is using stays. - **A path is stored as it was typed, and `~` is expanded where it is used.** `~/repos/x` and `/home/someone/repos/x` are a path and a snapshot of where it diff --git a/PLAN.md b/PLAN.md index d9eaf20..a6d4c69 100644 --- a/PLAN.md +++ b/PLAN.md @@ -477,6 +477,30 @@ deliberate and easy to undo by accident: only a turn whose model was changed under it. The conversation is read *before* the message is announced, which is what makes "everything before this message" true rather than a race against the pump. +- **An interrupt ends the wait, not the work** (2026-09-21, `awaiting`, + `Cancel`). Everything a llama turn waits on is on the far side of something + that cannot be told to stop -- a permission nobody has answered, a shell + command with a minute to run, a completion that says nothing until the + prompt has been read. Each now runs on a thread of its own and the + interrupt answers the *wait*; the abandoned thread finishes into a channel + nobody is reading. Measured at 43ms to end a turn in every state, against + up to a minute before. Two things fall out of it. Cancellation is a **token + per turn** rather than a flag on the session, because the abandoned thread + wakes later and a flag the next turn had reset would let it speak into a + conversation it is no longer part of. And the open thinking block is closed + by the interrupt, not by that thread, since the thread that knows is not + the one that ends the turn. +- **Stop takes the model with it when nobody else wants it** (2026-09-21). + Stop did nothing to a llama session at all: `stop_session` signals the + session's own recorded process and refuses a `Shared` one, which is the + whole point of that record -- so the session sat at `idle`. A `Shared` + record routes to `Driver::stop` now, because what stopping means for a + session that borrows the machine's process is the driver's to say. Each + live session claims the model it is on (`Router::claim`) and the model is + unloaded when the last claim goes, which is the honest reading of "give me + my GPU back": a model another session is on is nobody's to take. Rejected: + unloading unconditionally, which is the thing the shared router exists to + prevent. - **A wait that can be measured says how far along it is** (2026-09-21, `GET /sessions/{id}/progress`, `Driver::progress`). Both of this driver's waits have a real number behind them and neither used to reach diff --git a/TODO.md b/TODO.md index ddb6aa9..39e2907 100644 --- a/TODO.md +++ b/TODO.md @@ -22,6 +22,33 @@ one in place when it turns out to need a decision. ## Session settings +Asked for by Bryan on 2026-09-21, in one run while other work was in flight. +The first two are one change; the rest can land separately. + +- [ ] **A screen, not a modal.** The settings dialog has outgrown one: it + scrolls inside itself and covers the session it is about. + +- [ ] **Two tabs on that screen, the way the main screen has three.** One is + the session's own settings; the other is *the same provider screen* + reached from the machines tab (`ProviderScreen`), for this session's + provider -- two ways in, one screen, no second copy of the truth. + +- [ ] **Compact fields everywhere.** The label goes above the box rather than + floating inside it, and the padding around the value comes down. The + value's own text size does not change: what is costing a row its height + is the framing, not the text. + +- [ ] **A system prompt per session.** For llama.cpp it is a `system` message + on each request, so it is one entry in `DriverKind::params` and no app + change; whether the CLI drivers get one (`--append-system-prompt`) is a + separate question. + +- [ ] **Stop unloads the model where nothing else is using it.** Today + `Driver::stop` deliberately leaves it in memory, because the server is + the machine's and another session may be on the same model. The answer + is a claim per live session on the router, and an unload when the last + one goes -- not an unconditional unload. + - [ ] Autocompact belongs in session settings; empty disables it, which is the default. Iris chose "hand it to the driver" — only where a driver has auto-compaction of its own. **That option was offered on a false premise diff --git a/server/src/session/llama/mod.rs b/server/src/session/llama/mod.rs index 2a7cb45..be0b688 100644 --- a/server/src/session/llama/mod.rs +++ b/server/src/session/llama/mod.rs @@ -280,6 +280,7 @@ impl Call { } /// What one completion produced. +#[derive(Default)] struct Reply { text: String, calls: Vec, @@ -370,6 +371,20 @@ struct Respawn { router: Arc, } +/// One turn's cancellation, and what decides whether a thread still belongs to +/// the turn that is running now. +/// +/// A token per turn rather than one flag on the session, because an interrupt +/// abandons the *wait* and leaves the work running: the abandoned thread wakes +/// some time later, and a flag the next turn had reset would let it write into +/// a conversation it is no longer part of. Its own token stays set for ever, +/// so it cannot. +type Cancel = Arc; + +/// How one abandoned wait is answered. Held by [`Shared::abandonable`] and +/// called by [`Shared::abandon_turn`]; see [`awaiting`]. +type Canceller = Box; + /// Everything the driver's own threads need, which is nearly all of it: a turn /// runs on a thread of its own and outlives any borrow of the driver. struct Shared { @@ -389,11 +404,17 @@ struct Shared { /// `~`. `None` leaves the directory to `llama-server`, which is the honest /// answer rather than a guess at one. cwd: Option, - /// Set by [`Shared::abandon_turn`]; the streaming loop and the tool loop - /// both check it, leaving what was produced in the transcript. A flag is - /// only as prompt as the next place somebody looks, which is why the two - /// waits that look nowhere are released by hand there as well. - cancel: AtomicBool, + /// The running turn's cancellation -- see [`Cancel`]. Replaced at the + /// start of every turn; [`Shared::abandon_turn`] sets whichever one is + /// current. + cancel: Mutex, + /// The open thinking block, if the model is in one: when it started. + /// + /// Here rather than in [`generate`]'s own frame because a turn can be + /// let go of while a block is open, and then the thread that knows is not + /// the thread that ends the turn -- see [`Shared::done_thinking`]. Without + /// it the card on the phone spins for ever. + thinking_since: Mutex>, /// Whether a turn is running, and what is waiting behind it. /// /// One lock over both, because the two decide each other: with a flag and @@ -426,9 +447,10 @@ struct Shared { /// is only ever about the wait the session is in now -- see /// [`SessionStatus::Reading`]. reading: Mutex>, - /// Tool calls a turn is blocked on, by the call's id, so that an - /// interrupt can let go of one rather than sit out whatever it is doing. - running_tools: Mutex>>, + /// Waits an interrupt can end, by the id of what is being waited on -- + /// see [`awaiting`]. Each entry answers its own wait; nothing here stops + /// the work, which goes on wherever it is until it notices. + abandonable: Mutex>, /// Which of the machine's tools this session offers its model. Live like /// the sampling settings and for the same reason -- it is applied to the /// next request rather than to anything that was started. @@ -516,7 +538,8 @@ impl LlamaDriver { session_dir: session_dir.to_path_buf(), sampling: Mutex::new(sampling), cwd, - cancel: AtomicBool::new(false), + cancel: Mutex::new(Arc::new(AtomicBool::new(false))), + thinking_since: Mutex::new(None), turns: Mutex::new(Turns::default()), serving: Mutex::new(Serving::Loading(meta.model.clone().unwrap_or_default())), settled: Condvar::new(), @@ -533,7 +556,7 @@ impl LlamaDriver { allowed: Mutex::new(allowances(transcript)), asked: Mutex::new(HashMap::new()), reading: Mutex::new(None), - running_tools: Mutex::new(HashMap::new()), + abandonable: Mutex::new(HashMap::new()), tools_wanted: Mutex::new(Chosen::from(meta.params.get(TOOLS).map(String::as_str))), watching: AtomicBool::new(false), thinking: Mutex::new(chosen_thinking(&meta.params)), @@ -656,6 +679,10 @@ impl LlamaDriver { shared.emit(Event::Images { images: shared.images(), }); + // This session is on this model now, which is what + // decides whether stopping it can take the model out of + // memory -- see `Router::release`. + router.claim(&session, &model); Serving::Ready { serves, tools: Arc::new(tools), @@ -844,37 +871,80 @@ impl Shared { } } - /// Lets go of every tool call this session is waiting on, telling the - /// model what a call that never came back was. - /// - /// The call itself is not stopped: `llama-server` does not take one back - /// once it has started it, so the thread that asked is left to collect an - /// answer nobody reads. What ends here is the *waiting*, which is what a - /// turn is actually made of. - fn abandon_tools(&self) { - for (_, finished) in std::mem::take( - &mut *self - .running_tools + /// The running turn's cancellation, for a thread that is about to act on + /// behalf of it. + fn cancel(&self) -> Cancel { + Arc::clone( + &self + .cancel .lock() .unwrap_or_else(std::sync::PoisonError::into_inner), - ) { - let _ = finished.send(tools::UNFINISHED.to_string()); + ) + } + + /// Opens a turn: a token of its own, which is what makes everything the + /// last turn left running unable to speak into this one. + fn open_turn(&self) -> Cancel { + let fresh: Cancel = Arc::new(AtomicBool::new(false)); + *self + .cancel + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::clone(&fresh); + fresh + } + + /// Closes the open thinking block, saying how long it was open for. + /// + /// Exactly once, whoever gets here first: the stream's next fragment, the + /// end of the reply, or the reader ending the turn from another thread. + fn done_thinking(&self) { + let started = self + .thinking_since + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + if let Some(started) = started { + self.emit(Event::ThinkingDone { + ms: started.elapsed().as_millis() as u64, + }); } } - /// Ends the running turn: the flag, and everything blocked on something + /// Lets go of everything this turn is waiting on, telling whatever was + /// waiting what a thing that never came back is. + /// + /// None of it is *stopped*: a tool call runs to `llama-server`'s own + /// timeout, a completion generates until its connection is dropped, and + /// neither takes an instruction to stop. So each abandoned thread is left + /// to finish and find nobody there -- and it finds its turn's token set, + /// which is what keeps it silent. + fn abandon_waits(&self) { + for (_, abandon) in std::mem::take( + &mut *self + .abandonable + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) { + abandon(); + } + } + + /// Ends the running turn: the token, and everything blocked on something /// that will not read it. /// - /// Both of those waits are on another thread with no deadline a reader - /// would sit through -- a permission nobody is going to answer, and a tool - /// call running a shell command, which is up to a minute of `llama-server` - /// not being asked anything. Setting the flag alone left the turn exactly - /// where it was until that came back, which on screen was a Pause button - /// that did nothing at all. + /// Every one of those waits is on another thread with no deadline a reader + /// would sit through -- a permission nobody is going to answer, a shell + /// command that has a minute to run, a model that has said nothing yet and + /// will not for another twenty seconds. Setting the flag alone left the + /// turn exactly where it was until that came back, which on screen was a + /// Pause button that did nothing at all. fn abandon_turn(&self) { - self.cancel.store(true, Ordering::SeqCst); + self.cancel().store(true, Ordering::SeqCst); self.abandon_questions(); - self.abandon_tools(); + self.abandon_waits(); + // The block the model had open is over, and the thread that was + // watching it is not coming back to say so. + self.done_thinking(); } fn emit(&self, event: Event) { @@ -1103,7 +1173,7 @@ impl LlamaDriver { ) { let shared = Arc::clone(shared); std::thread::spawn(move || { - shared.cancel.store(false, Ordering::SeqCst); + let cancel = shared.open_turn(); // Everything before this message is read here, before the // message is announced: the record is written behind us, so a // conversation read after the announcement can already contain @@ -1142,7 +1212,7 @@ impl LlamaDriver { text, image_parts(&shared.session_dir, &images), )); - if let Err(err) = converse(&shared, &serves, &tools, messages) { + if let Err(err) = converse(&shared, &serves, &tools, messages, &cancel) { shared.emit(Event::Error { message: format!("{err:#}"), }); @@ -1180,15 +1250,34 @@ fn converse( serves: &Serves, tools: &Arc, mut messages: Vec, + cancel: &Cancel, ) -> Result<()> { for _ in 0..MAX_STEPS { - if shared.cancel.load(Ordering::SeqCst) { + if cancel.load(Ordering::SeqCst) { return Ok(()); } // Read per call rather than per turn, so a sampling change made while // a long turn is running reaches the rest of it. let sampling = shared.sampling.lock().unwrap().clone(); - let reply = generate(serves, &messages, tools, &sampling, shared)?; + // The conversation is handed *to* the thread that generates and + // handed back with the reply, rather than borrowed: the request can + // be abandoned, and what a thread with no borrow to honour leaves + // behind is nothing at all. It is also the whole prompt, which is not + // a thing to copy once a turn. + let asked = std::mem::take(&mut messages); + let generated = awaiting(shared, GENERATING, cancel, { + let shared = Arc::clone(shared); + let serves = serves.clone(); + let tools = Arc::clone(tools); + let cancel = Arc::clone(cancel); + move || generate(&serves, asked, &tools, &sampling, &shared, &cancel) + }); + // Let go of mid-request: the reader ended the turn, and what arrived + // before that is already in the transcript. + let Some((asked, reply)) = generated.transpose()? else { + return Ok(()); + }; + messages = asked; let calls = reply.calls; messages.push(Message { tool_calls: calls.iter().map(Call::wire).collect(), @@ -1198,10 +1287,10 @@ fn converse( return Ok(()); } for call in &calls { - let output = run_call(shared, tools, call); + let output = run_call(shared, tools, call, cancel); messages.push(Message::result_of(&call.id, output)); } - take_steers(shared, &mut messages); + take_steers(shared, &mut messages, cancel); } // Said rather than left as a turn that simply stopped: a reply that ends // here and one that ends because the model was finished look identical on @@ -1235,11 +1324,11 @@ fn converse( /// the next turn, which rebuilds the conversation out of the transcript -- /// which is what the note says anyway, since by then the turn it was typed /// into is over. -fn take_steers(shared: &Arc, messages: &mut Vec) { +fn take_steers(shared: &Arc, messages: &mut Vec, cancel: &Cancel) { // An interrupted turn takes nothing: this request is not going out, so a // message taken here would be recorded as read and then never answered. // Left in the queue it opens the next turn instead. - if shared.cancel.load(Ordering::SeqCst) { + if cancel.load(Ordering::SeqCst) { return; } let waiting = std::mem::take(&mut shared.turns.lock().unwrap().waiting); @@ -1270,7 +1359,7 @@ fn take_steers(shared: &Arc, messages: &mut Vec) { /// refused or interrupted. A call with no result is a conversation a chat /// template cannot render -- see [`tools::UNFINISHED`] -- so there is no path /// out of here that leaves one. -fn run_call(shared: &Arc, tools: &Arc, call: &Call) -> String { +fn run_call(shared: &Arc, tools: &Arc, call: &Call, cancel: &Cancel) -> String { let arguments = call.arguments(); shared.emit(Event::ToolStart { id: call.id.clone(), @@ -1284,12 +1373,12 @@ fn run_call(shared: &Arc, tools: &Arc, call: &Call) -> String { "There is no tool called {}. Use one of the tools you were given.", call.name ) - } else if shared.cancel.load(Ordering::SeqCst) { + } else if cancel.load(Ordering::SeqCst) { tools::UNFINISHED.to_string() } else if !permitted(shared, call) { tools::REFUSED.to_string() } else { - run_tool(shared, tools, call, arguments) + run_tool(shared, tools, call, arguments, cancel) }; shared.emit(Event::ToolEnd { id: call.id.clone(), @@ -1298,56 +1387,89 @@ fn run_call(shared: &Arc, tools: &Arc, call: &Call) -> String { output } -/// Runs one tool on a thread of its own, and waits for it in a way an -/// interrupt can reach. +/// Runs one tool where an interrupt can let go of it. /// -/// The wait is a channel with two writers -- the call's own thread and -/// [`Shared::abandon_tools`] -- and whichever speaks first is what the model -/// is told. A call is the one part of a turn this server can neither hurry -/// nor take back: `llama-server` runs it to its own timeout, up to a minute -/// for a shell command, and nothing in the protocol says "stop that one". So -/// an interrupt leaves it running over there and drops its answer, which is -/// what makes Pause end the turn when it is pressed rather than when the -/// command it happened to be running is over. -fn run_tool(shared: &Arc, tools: &Arc, call: &Call, arguments: Value) -> String { - let (finished, waiting) = std::sync::mpsc::channel(); - shared - .running_tools - .lock() - .unwrap() - .insert(call.id.clone(), finished.clone()); - // Checked after registering, not before: the flag is set ahead of the - // drain, so a call that registered after that drain would be waiting on a - // wake-up already given out. Either the insert beat the drain and this - // will be released, or it did not and the flag is already true. - if shared.cancel.load(Ordering::SeqCst) { - shared.running_tools.lock().unwrap().remove(&call.id); - return tools::UNFINISHED.to_string(); +/// A call is the one part of a turn this server can neither hurry nor take +/// back: `llama-server` runs it to its own timeout, up to a minute for a +/// shell command, and nothing in the protocol says "stop that one". So an +/// interrupt leaves it running over there and drops its answer, which is what +/// makes Pause end the turn when it is pressed rather than when the command +/// it happened to be running is over. +fn run_tool( + shared: &Arc, + tools: &Arc, + call: &Call, + arguments: Value, + cancel: &Cancel, +) -> String { + let tools = Arc::clone(tools); + let name = call.name.clone(); + let cwd = shared.cwd.clone(); + // Read here rather than in the thread: the filter is live, and this is the + // call it applies to. + let chosen = shared.tools_wanted.lock().unwrap().clone(); + awaiting(shared, &call.id, cancel, move || { + match tools.execute(&name, &arguments, cwd.as_deref(), &chosen) { + Ok(output) => output, + // Reaching the tool failed, which is this server's problem and not + // the model's work going wrong -- but the model is still what has + // to carry on, so it is told in the result rather than only in the + // log. + Err(err) => format!("This tool could not be run: {err:#}"), + } + }) + .unwrap_or_else(|| tools::UNFINISHED.to_string()) +} + +/// The id [`awaiting`] registers a completion under. One name is enough: a +/// session runs one turn at a time, and an abandoned turn's entry is taken by +/// the abandoning. +const GENERATING: &str = "generating"; + +/// Runs `work` on a thread of its own and waits for it in a way an interrupt +/// can reach, answering `None` where one arrived first. +/// +/// What this buys is that **the wait is what ends, not the work**: nothing +/// here can stop a shell command or a model mid-reply, and pretending +/// otherwise would mean holding the turn open until whatever it was doing +/// finished -- a minute for a command, twenty seconds for a long prompt, all +/// of it with a Pause button that had visibly done nothing. The abandoned +/// thread finishes into a channel nobody is reading, finds its turn's +/// [`Cancel`] set, and says nothing. +fn awaiting( + shared: &Arc, + id: &str, + cancel: &Cancel, + work: impl FnOnce() -> T + Send + 'static, +) -> Option { + let (done, waiting) = std::sync::mpsc::channel(); + let interrupt = done.clone(); + shared.abandonable.lock().unwrap().insert( + id.to_string(), + Box::new(move || { + let _ = interrupt.send(None); + }), + ); + // Registered before the token is read, not after: an interrupt sets it and + // *then* takes the entries, so anything registering after that would be + // waiting on a wake-up already given out. Either this got in before, or + // the token is set and this answers itself with the same entry. + if cancel.load(Ordering::SeqCst) { + if let Some(abandon) = shared.abandonable.lock().unwrap().remove(id) { + abandon(); + } + // Whichever of this and the interrupt took the entry has called it, so + // there is an answer waiting either way. + return waiting.recv().ok().flatten(); } - { - let tools = Arc::clone(tools); - let name = call.name.clone(); - let cwd = shared.cwd.clone(); - // Read here rather than in the thread, which is where it was read - // before: the filter is live, and this is the call it applies to. - let chosen = shared.tools_wanted.lock().unwrap().clone(); - std::thread::spawn(move || { - let output = match tools.execute(&name, &arguments, cwd.as_deref(), &chosen) { - Ok(output) => output, - // Reaching the tool failed, which is this server's problem and - // not the model's work going wrong -- but the model is still - // what has to carry on, so it is told in the result rather - // than only in the log. - Err(err) => format!("This tool could not be run: {err:#}"), - }; - let _ = finished.send(output); - }); - } - let output = waiting - .recv() - .unwrap_or_else(|_| tools::UNFINISHED.to_string()); - shared.running_tools.lock().unwrap().remove(&call.id); - output + std::thread::spawn(move || { + let _ = done.send(Some(work())); + }); + // `Err` is the working thread having panicked, which says the same thing + // an interrupt does: nothing is coming. + let answer = waiting.recv().ok().flatten(); + shared.abandonable.lock().unwrap().remove(id); + answer } /// Whether this call may go ahead, asking whoever is reading if it has to. @@ -1605,8 +1727,9 @@ impl Driver for LlamaDriver { }); return; } - self.shared.cancel.store(true, Ordering::SeqCst); - self.shared.cancel.store(false, Ordering::SeqCst); + // A token of its own for whatever comes next, so nothing left over + // from the last turn is holding the one this model will answer. + self.shared.open_turn(); match self.start(model) { // Reported when it is true and not before: `start` has put the // session into `Loading`, and the model it is loading is this one. @@ -1659,15 +1782,44 @@ impl Driver for LlamaDriver { self.shared.abandon_turn(); } - /// Ends this session's claim on the model and nothing else. + /// Ends this session's turn and its claim on the model, and takes the + /// model out of memory if this session was the last thing using it. /// - /// The server holding it is the machine's, shared with every other session - /// on it, so stopping one session unloads nothing -- see - /// [`process::Detail::Shared`], which is what makes that structural rather - /// than a rule to remember. A model is taken out of memory from the - /// machine's provider settings, where what it costs everybody is visible. + /// The server itself stays: it is the machine's, shared with every other + /// session on it, and `process::stop` refuses to signal a + /// [`process::Detail::Shared`] record -- which is what makes that + /// structural rather than a rule to remember. What Stop can honestly do + /// is give back the memory *this* session was the reason for, and a model + /// somebody else is on is not that; see [`Router::release`], which + /// answers with the model only when nobody else holds it. + /// + /// On a thread because it is an HTTP round trip to the serving machine, + /// and stopping a session should not wait on one: the session is stopped + /// either way, and a model that could not be unloaded is a message in the + /// log rather than a failure to report. fn stop(&self) { self.shared.abandon_turn(); + // Said here because nothing else will: a session whose process is the + // machine's shared server has no process of its own to die and be + // noticed. Without it the session sits at `idle` after Stop, with no + // Start button and no sign anything happened. + self.shared.emit(Event::Status { + state: SessionStatus::Exited, + }); + // Anything waiting on this model is waiting on a session that has + // stopped, and nothing else will wake it. + self.shared + .settle(Serving::Failed("this session was stopped.".to_string())); + if let Some(model) = self.respawn.router.release(&self.respawn.meta.id) { + let router = Arc::clone(&self.respawn.router); + std::thread::spawn(move || { + if let Err(err) = router.unload(&model) { + tracing::warn!("could not unload {model}: {err:#}"); + } else { + tracing::info!("unloaded {model}, which nothing else was using"); + } + }); + } process::clear(&self.shared.session_dir); } } @@ -2168,11 +2320,12 @@ fn thinking_kwargs(shared: &Shared) -> Option { /// how long it thought for. fn generate( serves: &Serves, - messages: &[Message], + messages: Vec, tools: &Tools, sampling: &serde_json::Map, shared: &Shared, -) -> Result { + cancel: &Cancel, +) -> Result<(Vec, Reply)> { let mut body = json!({ // Which model, because one `llama-server` is serving every model this // machine has loaded and this is how a request says which it means. @@ -2238,11 +2391,6 @@ fn generate( // worst direction for a figure somebody is watching to see how much room // is left. let mut context = None; - // The open thinking block: when it started, so its duration is measured - // where the deltas actually arrive rather than worked out later from two - // timestamps. `None` between blocks -- a turn can think, speak, call a - // tool and think again. - let mut thinking: Option = None; // Whether the model has produced anything at all yet; until it has, this // turn is still reading its prompt. let mut speaking = false; @@ -2259,17 +2407,8 @@ fn generate( // server's own account of prompt processing, and a prompt it had cached is // a small number rather than a missing one. let mut prefill = None; - // Closes the open block, which is anything the model says that is not more - // working: the first word of the reply, or a tool call. - let done_thinking = |thinking: &mut Option| { - if let Some(started) = thinking.take() { - shared.emit(Event::ThinkingDone { - ms: started.elapsed().as_millis() as u64, - }); - } - }; for line in std::io::BufRead::lines(reader) { - if shared.cancel.load(Ordering::SeqCst) { + if cancel.load(Ordering::SeqCst) { break; } let line = line.context("reading the generation stream")?; @@ -2289,7 +2428,7 @@ fn generate( // `refusal` never sees it. Skipped as an unrecognised chunk, it ended // the turn with nothing said at all. if let Some(why) = chunk.pointer("/error/message").and_then(Value::as_str) { - done_thinking(&mut thinking); + shared.done_thinking(); bail!("llama-server stopped generating: {why}"); } // One figure answering both questions, which for this dialect it @@ -2335,7 +2474,11 @@ fn generate( if let Some(fragment) = delta.get("reasoning_content").and_then(Value::as_str) && !fragment.is_empty() { - thinking.get_or_insert_with(std::time::Instant::now); + shared + .thinking_since + .lock() + .unwrap() + .get_or_insert_with(std::time::Instant::now); shared.emit(Event::Thinking { delta: fragment.to_string(), }); @@ -2343,7 +2486,7 @@ fn generate( if let Some(fragment) = delta.get("content").and_then(Value::as_str) && !fragment.is_empty() { - done_thinking(&mut thinking); + shared.done_thinking(); text.push_str(fragment); shared.emit(Event::AssistantText { delta: fragment.to_string(), @@ -2355,7 +2498,7 @@ fn generate( .into_iter() .flatten() { - done_thinking(&mut thinking); + shared.done_thinking(); absorb(&mut calls, fragment); } } @@ -2364,15 +2507,22 @@ fn generate( // ever having spoken. *shared.reading.lock().unwrap() = None; // A block left open by the end of the stream -- a model that thought and - // then said nothing, or a turn the reader cancelled -- is still a block - // that ended. Without this its card spins for ever. - done_thinking(&mut thinking); + // then said nothing -- is still a block that ended. Without this its card + // spins for ever. The interrupted case is closed by the interrupt itself, + // which is the thread that ends that turn. + shared.done_thinking(); // Said before anything else is reported, because what arrived is a partial // reply rather than a reply. A cancelled turn leaves the same stream // unfinished and is not a failure: it is what the reader asked for. - if !finished && !shared.cancel.load(Ordering::SeqCst) { + if !finished && !cancel.load(Ordering::SeqCst) { bail!("llama-server stopped sending this reply before it was finished"); } + // Nothing more is said into a turn somebody ended. This thread is here + // only to notice that the stream stopped, and what it has left to report + // is about a turn that is over -- see `awaiting`. + if cancel.load(Ordering::SeqCst) { + return Ok((messages, Reply::default())); + } if tokens > 0 { shared.emit(Event::UsageDelta { tokens, @@ -2385,7 +2535,7 @@ fn generate( // is cut mid-fragment, and running it would mean inventing what was asked // for. calls.retain(|call| !call.name.is_empty()); - Ok(Reply { text, calls }) + Ok((messages, Reply { text, calls })) } /// Whether a delta carries anything the model produced, of any kind. diff --git a/server/src/session/llama/router.rs b/server/src/session/llama/router.rs index fa851d6..2ca052a 100644 --- a/server/src/session/llama/router.rs +++ b/server/src/session/llama/router.rs @@ -146,6 +146,7 @@ impl Routers { gate: Mutex::new(()), loading: Arc::new(Mutex::new(HashMap::new())), watching: Arc::new(AtomicBool::new(false)), + claims: Mutex::new(HashMap::new()), }) }); *router.spec.lock().unwrap_or_else(|e| e.into_inner()) = spec; @@ -180,6 +181,16 @@ pub struct Router { /// is enough, and every load asks -- an adopted router has none until /// somebody wants a model from it. watching: Arc, + /// Which model each live session on this router is holding, so that the + /// last session to let go of one is the one that can take it out of + /// memory -- see [`release`](Self::release). + /// + /// One model per session because that is what a session has: claiming a + /// new one is how the old claim ends, which is what a model change is. + /// Only sessions *this* backend has running are in here, which is the + /// right set: a stopped session is not using a model, whatever it is + /// configured with. + claims: Mutex>, } /// Where one model's load has got to, as the router last said. @@ -361,6 +372,28 @@ impl Router { .unwrap_or_default() } + /// Records that `session` is on `model`, which is also how its claim on + /// whatever it was on before ends. + pub fn claim(&self, session: &str, model: &str) { + self.claims + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(session.to_string(), model.to_string()); + } + + /// Lets go of `session`'s model, answering with it when no other session + /// is on it -- which is to say, when taking it out of memory would cost + /// nobody anything. + /// + /// `None` for a session that had no claim, and for one whose model + /// somebody else is still using. Both of those are the same instruction: + /// leave it where it is. + pub fn release(&self, session: &str) -> Option { + let mut claims = self.claims.lock().unwrap_or_else(|e| e.into_inner()); + let model = claims.remove(session)?; + (!claims.values().any(|held| *held == model)).then_some(model) + } + /// Takes one model out of memory, leaving the router and every other model /// alone. pub fn unload(&self, key: &str) -> Result<()> { diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs index 36715ed..49260af 100644 --- a/server/src/session/mod.rs +++ b/server/src/session/mod.rs @@ -2040,10 +2040,19 @@ impl SessionManager { /// ahead of the measurement, and wrong for the grace period a process /// that ignores SIGTERM keeps running. /// - /// Deliberately not routed through the driver: the record is the - /// session's rather than any dialect's, so asking it here stops a - /// session whose driver is in no state to be asked, and adds no method a - /// new driver could implement wrongly. + /// Deliberately not routed through the driver *where the process is the + /// session's own*: the record is the session's rather than any dialect's, + /// so asking it here stops a session whose driver is in no state to be + /// asked, and adds no method a new driver could implement wrongly. + /// + /// A [`process::Detail::Shared`] record is the exception, and it is one + /// this rule could not have anticipated: the process is the machine's + /// llama-server, shared with every other session on it, so `process::stop` + /// refuses to signal it -- and Stop therefore did nothing at all, leaving + /// the session sitting at `idle`. What stopping means for a session that + /// borrows somebody else's process is the driver's to say, and only it + /// can say it: it is the half that knows whether the model is still + /// wanted by anybody else. pub fn stop_session(&self, id: &str) -> Result<()> { if !self .inner @@ -2072,6 +2081,23 @@ impl SessionManager { }; tracing::info!("stopping session {id} (pid {})", record.pid); process::mark_stopping(&self.data_dir.join(id))?; + if let process::Detail::Shared { .. } = record.detail { + let driver = self + .inner + .read() + .unwrap() + .live + .get(id) + .and_then(|session| session.driver()); + match driver { + Some(driver) => driver.stop(), + // A borrowed process and no driver to ask: the record is all + // there is, and clearing it is what says this session is not + // using that server any more. + None => process::clear(&self.data_dir.join(id)), + } + return Ok(()); + } process::stop(&record, process::STOP_GRACE); Ok(()) }