DS4Server/src/runtime.rs

use crate::engine::{ChatTurn, GenerationOutput, Generator};
use crate::metrics::{Metrics, WorkSource};
use crate::settings::{EngineSettings, TurnSettings};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time::{Duration, Instant};

const RUNTIME_PANIC_ERROR: &str = "The model runtime hit an internal error and was reset.";

#[derive(Clone)]
pub(crate) struct GenerationService {
    commands: Sender<Command>,
    metrics: Arc<Metrics>,
}

pub(crate) struct ActiveGeneration {
    pub(crate) events: Receiver<GenerationEvent>,
    pub(crate) cancel: Arc<AtomicBool>,
}

impl Drop for ActiveGeneration {
    fn drop(&mut self) {
        self.cancel.store(true, Ordering::Relaxed);
    }
}

pub(crate) enum CheckpointTarget {
    Local(PathBuf),
    Transient(PathBuf),
    /// Same transient KV handling as [`CheckpointTarget::Transient`], but asked
    /// for by the app itself (session titling) rather than by an HTTP client.
    OneShot(PathBuf),
}

impl CheckpointTarget {
    fn source(&self) -> WorkSource {
        match self {
            Self::Local(_) | Self::OneShot(_) => WorkSource::LocalChat,
            Self::Transient(_) => WorkSource::Http,
        }
    }
}

pub(crate) enum GenerationEvent {
    Loading,
    Chunk {
        reasoning: bool,
        content: String,
    },
    Context {
        used: u32,
        limit: u32,
        tokens_per_second: Option<f32>,
    },
    Finished(Result<GenerationOutput, String>),
}

struct Command {
    engine: EngineSettings,
    turn: TurnSettings,
    messages: Vec<ChatTurn>,
    checkpoint: CheckpointTarget,
    idle_timeout: Duration,
    cancel: Arc<AtomicBool>,
    events: Sender<GenerationEvent>,
}

impl GenerationService {
    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, worker_metrics))
            .map_err(|error| format!("Could not start the model runtime: {error}"))?;
        Ok(Self { commands, metrics })
    }

    pub(crate) fn generate(
        &self,
        engine: EngineSettings,
        turn: TurnSettings,
        messages: Vec<ChatTurn>,
        checkpoint: CheckpointTarget,
        idle_timeout: Duration,
    ) -> Result<ActiveGeneration, String> {
        let source = checkpoint.source();
        let cancel = Arc::new(AtomicBool::new(false));
        let (events, receiver) = mpsc::channel();
        self.metrics.request_queued(source);
        if self
            .commands
            .send(Command {
                engine,
                turn,
                messages,
                checkpoint,
                idle_timeout,
                cancel: Arc::clone(&cancel),
                events,
            })
            .is_err()
        {
            self.metrics.request_rejected();
            return Err("The model runtime stopped unexpectedly.".to_owned());
        }
        Ok(ActiveGeneration {
            events: receiver,
            cancel,
        })
    }
}

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 = command.checkpoint.source();
                let events = command.events.clone();
                metrics.request_started(source);
                if let Err(error) = catch_runtime_panic(|| {
                    run_command(
                        command,
                        &mut loaded,
                        &metrics,
                        request_started,
                        source,
                        &mut last_used,
                        &mut idle_timeout,
                    );
                }) {
                    metrics.request_failed(request_started.elapsed());
                    if loaded.take().is_some() {
                        metrics.unloaded();
                    }
                    let _ = events.send(GenerationEvent::Finished(Err(error)));
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {
                if loaded.is_some() && last_used.elapsed() >= idle_timeout {
                    loaded = None;
                    metrics.unloaded();
                }
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn run_command(
    command: Command,
    loaded: &mut Option<(EngineSettings, Generator)>,
    metrics: &Arc<Metrics>,
    request_started: Instant,
    source: WorkSource,
    last_used: &mut Instant,
    idle_timeout: &mut Duration,
) {
    *idle_timeout = command.idle_timeout;
    if command.cancel.load(Ordering::Relaxed) {
        metrics.request_failed(request_started.elapsed());
        let _ = command.events.send(GenerationEvent::Finished(
            Err("generation cancelled".into()),
        ));
        return;
    }
    if loaded
        .as_ref()
        .is_none_or(|(settings, _)| settings != &command.engine)
    {
        if loaded.take().is_some() {
            metrics.unloaded();
        }
        let _ = command.events.send(GenerationEvent::Loading);
        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)) = 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,
                tokens_per_second,
            });
        };
        let result = match command.checkpoint {
            CheckpointTarget::Local(path) => generator.generate(
                &path,
                &command.messages,
                &command.turn,
                &command.cancel,
                &mut emit,
                &mut progress,
            ),
            CheckpointTarget::Transient(directory) | CheckpointTarget::OneShot(directory) => {
                generator.generate_transient(
                    &directory,
                    &command.messages,
                    &command.turn,
                    &command.cancel,
                    &mut emit,
                    &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();
    }
}

fn catch_runtime_panic<T>(operation: impl FnOnce() -> T) -> Result<T, String> {
    // Generator owns GPU handles and is not unwind-safe; callers discard it on error.
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(operation))
        .map_err(|_| RUNTIME_PANIC_ERROR.to_owned())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dropping_active_generation_cancels_it() {
        let cancel = Arc::new(AtomicBool::new(false));
        let (_, events) = mpsc::channel();

        drop(ActiveGeneration {
            events,
            cancel: Arc::clone(&cancel),
        });

        assert!(cancel.load(Ordering::Relaxed));
    }

    #[test]
    fn runtime_panic_recovery_accepts_the_next_operation() {
        assert_eq!(
            catch_runtime_panic(|| panic!("injected runtime panic")),
            Err(RUNTIME_PANIC_ERROR.to_owned())
        );
        assert_eq!(catch_runtime_panic(|| 42), Ok(42));
    }
}