engine/crates/runtime/src/tests/mod.rs

//! Scheduler tests on the CPU `FakeModel`: prefix cache, image blocks, chunked prefill, KV exhaustion and eviction, budgets, the deterministic
//! thought cap, concurrency, and the sliding-window KV pool (swa.rs).
use crate::fake::*;
use crate::*;
use dllm_core::policy::Policy;
use dllm_core::serving::Item;
use dllm_core::tokenizer::Tok;
use std::sync::Arc;

mod engine;
mod swa;

pub fn tok() -> Arc<Tok> {
    let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../tokenizer.json");
    Arc::new(Tok::from_file(p).expect("tokenizer.json of the release folder"))
}

pub fn engine(max_len: usize, pages: usize, max_rows: usize, tweak: impl FnOnce(&mut EngineConfig)) -> (Engine, Arc<FakeCtl>) {
    let m = FakeModel::full_kv(pages, max_rows);
    let ctl = m.ctl.clone();
    let mut cfg = EngineConfig { max_len, prep: Some(Arc::new(FakePrep)), ..Default::default() };
    tweak(&mut cfg);
    (Engine::start(cfg, Box::new(m), Some(Box::new(FakeVision { encodes: 0 })), tok()).unwrap(), ctl)
}

/// a question over `nopt` options: prompt = prefix + tail + "\n"; the fake LM reads the option letters at ids 150..
pub fn text_item(key: &str, prefix: &[u32], tail: &[u32], nopt: usize) -> Item {
    let mut ids = prefix.to_vec();
    ids.extend_from_slice(tail);
    ids.push(107);
    Item { key: key.into(), noul: false, keys: (0..nopt).map(|i| format!("k{i}")).collect(), ids, label_ids: (150..150 + nopt as u32).collect(), ..Default::default() }
}

pub fn toks(seed: u32, n: usize) -> Vec<u32> {
    (0..n as u32).map(|i| 2 + (seed.wrapping_mul(2654435761).wrapping_add(i.wrapping_mul(40503)) >> 7) % 90 + 200).collect()
}

pub fn think_pol() -> Policy {
    Policy { tau: 1.01, cap: 30, budget_ms: 0, ..Policy::default() }
}