Implement batched server parity and observability

src/runtime.rs

@@ -1,4 +1,5 @@
 use crate::engine::{ChatTurn, GenerationOutput, Generator};
+use crate::metrics::{Metrics, WorkSource};
 use crate::settings::{EngineSettings, TurnSettings};
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -10,6 +11,7 @@ use std::time::{Duration, Instant};
 #[derive(Clone)]
 pub(crate) struct GenerationService {
     commands: Sender<Command>,
+    metrics: Arc<Metrics>,
 }
 
 pub(crate) struct ActiveGeneration {
@@ -47,13 +49,14 @@ struct Command {
 }
 
 impl GenerationService {
-    pub(crate) fn spawn() -> Result<Self, String> {
+    pub(crate) fn spawn(metrics: Arc<Metrics>) -> Result<Self, String> {
         let (commands, receiver) = mpsc::channel::<Command>();
+        let worker_metrics = Arc::clone(&metrics);
         thread::Builder::new()
             .name("model-runtime".into())
-            .spawn(move || run(receiver))
+            .spawn(move || run(receiver, worker_metrics))
             .map_err(|error| format!("Could not start the model runtime: {error}"))?;
-        Ok(Self { commands })
+        Ok(Self { commands, metrics })
     }
 
     pub(crate) fn generate(
@@ -64,9 +67,15 @@ impl GenerationService {
         checkpoint: CheckpointTarget,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
+        let source = match &checkpoint {
+            CheckpointTarget::Local(_) => WorkSource::LocalChat,
+            CheckpointTarget::Transient(_) => WorkSource::Http,
+        };
         let cancel = Arc::new(AtomicBool::new(false));
         let (events, receiver) = mpsc::channel();
-        self.commands
+        self.metrics.request_queued(source);
+        if self
+            .commands
             .send(Command {
                 engine,
                 turn,
@@ -76,7 +85,11 @@ impl GenerationService {
                 cancel: Arc::clone(&cancel),
                 events,
             })
-            .map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?;
+            .is_err()
+        {
+            self.metrics.request_rejected();
+            return Err("The model runtime stopped unexpectedly.".to_owned());
+        }
         Ok(ActiveGeneration {
             events: receiver,
             cancel,
@@ -84,34 +97,70 @@ impl GenerationService {
     }
 }
 
-fn run(commands: Receiver<Command>) {
+fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
     let mut loaded = None::<(EngineSettings, Generator)>;
     let mut last_used = Instant::now();
     let mut idle_timeout = Duration::from_secs(15 * 60);
     loop {
         match commands.recv_timeout(Duration::from_secs(1)) {
             Ok(command) => {
+                let request_started = Instant::now();
+                let source = match &command.checkpoint {
+                    CheckpointTarget::Local(_) => WorkSource::LocalChat,
+                    CheckpointTarget::Transient(_) => WorkSource::Http,
+                };
+                metrics.request_started(source);
                 idle_timeout = command.idle_timeout;
                 if loaded
                     .as_ref()
                     .is_none_or(|(settings, _)| settings != &command.engine)
                 {
+                    if loaded.take().is_some() {
+                        metrics.unloaded();
+                    }
                     let _ = command.events.send(GenerationEvent::Loading);
-                    loaded = match Generator::open(&command.engine) {
-                        Ok(generator) => Some((command.engine.clone(), generator)),
+                    metrics.loading();
+                    let load_started = Instant::now();
+                    loaded = match Generator::open(&command.engine, Arc::clone(&metrics)) {
+                        Ok(generator) => {
+                            let summary = generator.summary();
+                            metrics.loaded(
+                                summary.model,
+                                load_started.elapsed(),
+                                summary.mapped_bytes,
+                                summary.tensor_count,
+                                summary.vocabulary_size,
+                            );
+                            Some((command.engine.clone(), generator))
+                        }
                         Err(error) => {
+                            metrics.request_failed(request_started.elapsed());
                             let _ = command.events.send(GenerationEvent::Finished(Err(error)));
                             None
                         }
                     };
                 }
                 if let Some((_, generator)) = &mut loaded {
+                    let mut prefill_started = None::<(Instant, u32)>;
                     let mut emit = |reasoning, content| {
                         let _ = command
                             .events
                             .send(GenerationEvent::Chunk { reasoning, content });
                     };
                     let mut progress = |used, limit, tokens_per_second| {
+                        if let Some(speed) = tokens_per_second {
+                            metrics.generation_progress(used, limit, speed);
+                        } else {
+                            let (started, initial) =
+                                prefill_started.get_or_insert_with(|| (Instant::now(), used));
+                            let elapsed = started.elapsed().as_secs_f32();
+                            let speed = if elapsed > 0.0 {
+                                used.saturating_sub(*initial) as f32 / elapsed
+                            } else {
+                                0.0
+                            };
+                            metrics.prefill_progress(used, limit, speed);
+                        }
                         let _ = command.events.send(GenerationEvent::Context {
                             used,
                             limit,
@@ -136,6 +185,18 @@ fn run(commands: Receiver<Command>) {
                             &mut progress,
                         ),
                     };
+                    match &result {
+                        Ok(output) => metrics.request_finished(
+                            source,
+                            request_started.elapsed(),
+                            output.prompt_tokens,
+                            output.cached_tokens,
+                            output.completion_tokens,
+                            output.previous_checkpoint_bytes,
+                            output.checkpoint_bytes,
+                        ),
+                        Err(_) => metrics.request_failed(request_started.elapsed()),
+                    }
                     let _ = command.events.send(GenerationEvent::Finished(result));
                     last_used = Instant::now();
                 }
@@ -143,6 +204,7 @@ fn run(commands: Receiver<Command>) {
             Err(mpsc::RecvTimeoutError::Timeout) => {
                 if loaded.is_some() && last_used.elapsed() >= idle_timeout {
                     loaded = None;
+                    metrics.unloaded();
                 }
             }
             Err(mpsc::RecvTimeoutError::Disconnected) => break,