Add GPU lease coordinating LLM and TTS requests through llama-swap
llama-swap runs chat/vision/Chatterbox as a mutually-exclusive set on one GPU and HOLDS a request for a non-resident model until the resident model drains, then swaps. That hold burned the holder's reqwest timeout (measured: a queued TTS lost 77s behind one LLM turn; an LLM request behind a synthesis waited the entire remaining synth), so concurrent insight + read-aloud timed out instead of queueing. ai::gpu adds a fair RwLock lease acquired before each request is sent, so cross-model waits happen before the HTTP timeout starts: chat/vision share the read lease, TTS synthesis and voice-library ops (which spin Chatterbox up) take the write lease, and embeddings take none (the embed slot is in llama-swap's always-resident group). Speech jobs now flip queued->running only after acquiring the GPU, letting the client anchor its poll deadline to that transition. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,88 @@
|
||||
// GPU lease — in-process coordination for llama-swap model contention.
|
||||
//
|
||||
// llama-swap runs the heavyweight models (chat / vision / Chatterbox TTS) as
|
||||
// a mutually-exclusive set on one GPU (matrix DSL `(q27 | … | tts) & e`): a
|
||||
// request for a non-resident model is HELD by llama-swap until the resident
|
||||
// model's in-flight requests drain, then the models swap. That hold counts
|
||||
// against the *holder's* reqwest timeout — measured live: a queued TTS burned
|
||||
// 77s of its budget behind a single LLM turn, and an LLM request behind a
|
||||
// running synthesis waited the entire remaining synth. Uncoordinated
|
||||
// cross-model traffic therefore times out instead of queueing.
|
||||
//
|
||||
// The lease moves that wait into this process, BEFORE the HTTP request is
|
||||
// sent and before its timeout starts:
|
||||
// - chat/vision requests (the LLM-side slots) share the READ lease;
|
||||
// - TTS synthesis and voice-library ops (anything that spins Chatterbox up
|
||||
// and evicts the LLM) take the WRITE lease;
|
||||
// - embeddings take NO lease: the `embed` slot is in llama-swap's
|
||||
// always-resident group (the `& e` term) and never participates in a swap,
|
||||
// so leasing it would only stall searches behind a queued synthesis.
|
||||
//
|
||||
// tokio's RwLock is fair (FIFO, write-preferring): a queued TTS gets the GPU
|
||||
// right after the current LLM request drains, and later LLM requests queue
|
||||
// behind it — bounded waits in both directions, no starvation, no timeout
|
||||
// budget burned while waiting.
|
||||
//
|
||||
// RULES: hold a lease for exactly one HTTP request (for streaming, the
|
||||
// stream's lifetime) and NEVER acquire one while already holding one — once a
|
||||
// writer is queued, new read acquisitions block, so nested acquisition can
|
||||
// deadlock.
|
||||
|
||||
use std::sync::LazyLock;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
|
||||
|
||||
static GPU_LEASE: LazyLock<RwLock<()>> = LazyLock::new(|| RwLock::new(()));
|
||||
|
||||
/// Waits longer than this are logged — they mean a cross-model swap was
|
||||
/// avoided and quantify what the request *would* have burned of its timeout.
|
||||
const SLOW_WAIT_LOG_SECS: f64 = 2.0;
|
||||
|
||||
/// Shared lease for LLM-side requests (chat / vision slots).
|
||||
pub async fn llm_lease() -> RwLockReadGuard<'static, ()> {
|
||||
let started = Instant::now();
|
||||
let guard = GPU_LEASE.read().await;
|
||||
log_slow_wait("llm", started);
|
||||
guard
|
||||
}
|
||||
|
||||
/// Exclusive lease for TTS-side requests (speech synthesis + voice-library
|
||||
/// ops that spin up Chatterbox).
|
||||
pub async fn tts_lease() -> RwLockWriteGuard<'static, ()> {
|
||||
let started = Instant::now();
|
||||
let guard = GPU_LEASE.write().await;
|
||||
log_slow_wait("tts", started);
|
||||
guard
|
||||
}
|
||||
|
||||
fn log_slow_wait(kind: &str, started: Instant) {
|
||||
let waited = started.elapsed().as_secs_f64();
|
||||
if waited > SLOW_WAIT_LOG_SECS {
|
||||
log::info!("GPU lease ({kind}): waited {waited:.1}s for the other model class to drain");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// One sequential test, not several: the lease is a single global, so
|
||||
// parallel tests interleaving reads and writes on it can hit the very
|
||||
// nested-acquisition deadlock the module comment warns about.
|
||||
#[tokio::test]
|
||||
async fn write_lease_excludes_readers_then_reads_share() {
|
||||
let w = tts_lease().await;
|
||||
// A reader must not acquire while the writer is held.
|
||||
let pending = tokio::spawn(async { drop(llm_lease().await) });
|
||||
tokio::task::yield_now().await;
|
||||
assert!(!pending.is_finished());
|
||||
drop(w);
|
||||
pending.await.expect("reader acquires after writer drops");
|
||||
|
||||
// With no writer queued, read leases are shared.
|
||||
let a = llm_lease().await;
|
||||
let b = llm_lease().await;
|
||||
drop(a);
|
||||
drop(b);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user