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;
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time::{Duration, Instant};

#[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>,
}

pub(crate) enum CheckpointTarget {
    Local(PathBuf),
    Transient(PathBuf),
}

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 = match &checkpoint {
            CheckpointTarget::Local(_) => WorkSource::LocalChat,
            CheckpointTarget::Transient(_) => WorkSource::Http,
        };
        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 = 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);
                    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,
                            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) => 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();
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {
                if loaded.is_some() && last_used.elapsed() >= idle_timeout {
                    loaded = None;
                    metrics.unloaded();
                }
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
        }
    }
}