use crate::session::Ids;
use crate::{SkipReason, chunk, graph, image, keyscan};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, mpsc};
use std::time::{Duration, Instant};
pub const LANE_DEPTH: usize = 2;
const LANE_GONE: Duration = Duration::from_secs(2);
const ACTIVE_WINDOWS: usize = 3;
const WRITE_AHEAD: usize = 16;
const LOOKAHEAD_PER_THREAD: usize = 4;
const BATCH_SECONDS: f64 = 0.2;
#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
pub struct Stages {
pub wall_ms: u64,
pub walk_ms: u64,
pub read_ms: u64,
pub extract_ms: u64,
pub parse_ms: u64,
pub tokenize_ms: u64,
pub write_ms: u64,
pub embed_wait_ms: BTreeMap<String, u64>,
pub prepare_cpu_ms: BTreeMap<String, u64>,
pub write_parts_ms: BTreeMap<String, u64>,
}
impl Stages {
pub fn add(&mut self, other: &Stages) {
self.wall_ms += other.wall_ms;
self.walk_ms += other.walk_ms;
self.read_ms += other.read_ms;
self.extract_ms += other.extract_ms;
self.parse_ms += other.parse_ms;
self.tokenize_ms += other.tokenize_ms;
self.write_ms += other.write_ms;
for (lane, ms) in &other.embed_wait_ms {
*self.embed_wait_ms.entry(lane.clone()).or_default() += ms;
}
for (part, ms) in &other.prepare_cpu_ms {
*self.prepare_cpu_ms.entry(part.clone()).or_default() += ms;
}
for (part, ms) in &other.write_parts_ms {
*self.write_parts_ms.entry(part.clone()).or_default() += ms;
}
}
pub fn parts_ms(&self) -> u64 {
self.walk_ms
+ self.read_ms
+ self.extract_ms
+ self.parse_ms
+ self.tokenize_ms
+ self.write_ms
+ self.embed_wait_ms.values().sum::<u64>()
}
pub fn line(&self) -> String {
let waits: Vec<String> = self
.embed_wait_ms
.iter()
.map(|(lane, ms)| format!("embed wait ({lane}) {}", secs(*ms)))
.collect();
let mut parts = vec![
format!("walk {}", secs(self.walk_ms)),
format!("read+hash {}", secs(self.read_ms)),
format!("extract+scan {}", secs(self.extract_ms)),
format!("parse+chunk {}", secs(self.parse_ms)),
format!("tokenize {}", secs(self.tokenize_ms)),
];
parts.extend(waits);
let split: Vec<String> = self
.write_parts_ms
.iter()
.filter(|(_, ms)| **ms > 0)
.map(|(part, ms)| format!("{part} {}", secs(*ms)))
.collect();
if split.is_empty() {
parts.push(format!("write {}", secs(self.write_ms)));
} else {
parts.push(format!(
"write {} ({})",
secs(self.write_ms),
split.join(", ")
));
}
format!("stages over {}: {}", secs(self.wall_ms), parts.join(", "))
}
}
fn secs(ms: u64) -> String {
format!("{:.1} s", ms as f64 / 1000.0)
}
#[derive(Default)]
pub struct Clocks {
read: AtomicU64,
extract: AtomicU64,
parse: AtomicU64,
tokenize: AtomicU64,
}
impl Clocks {
fn tick(slot: &AtomicU64, since: Instant) -> Instant {
let now = Instant::now();
slot.fetch_add(
now.duration_since(since).as_micros() as u64,
Ordering::Relaxed,
);
now
}
fn read_all(&self) -> [u64; 4] {
[
self.read.load(Ordering::Relaxed),
self.extract.load(Ordering::Relaxed),
self.parse.load(Ordering::Relaxed),
self.tokenize.load(Ordering::Relaxed),
]
}
}
#[derive(Debug, Default)]
pub struct WriteParts {
micros: [std::cell::Cell<u64>; WRITE_PARTS.len()],
}
pub const WRITE_PARTS: [&str; 9] = [
"rows", "evict", "graph", "vectors", "commit", "save", "hashes", "cache", "images",
];
impl WriteParts {
pub fn add(&self, part: &str, since: Instant) {
if let Some(at) = WRITE_PARTS.iter().position(|p| *p == part) {
let cell = &self.micros[at];
cell.set(cell.get() + since.elapsed().as_micros() as u64);
}
}
pub fn take(&self) -> BTreeMap<String, u64> {
WRITE_PARTS
.iter()
.zip(&self.micros)
.map(|(part, cell)| (part.to_string(), cell.replace(0) / 1000))
.collect()
}
}
pub struct Charge<'a>(pub &'a WriteParts, pub &'static str, pub Instant);
impl Drop for Charge<'_> {
fn drop(&mut self) {
self.0.add(self.1, self.2);
}
}
pub struct WriterClock {
started: Instant,
pub walk_ms: u64,
prepare_wait: Duration,
embed_wait: Duration,
}
impl WriterClock {
pub fn start(walk_ms: u64) -> Self {
Self {
started: Instant::now(),
walk_ms,
prepare_wait: Duration::ZERO,
embed_wait: Duration::ZERO,
}
}
pub fn waited_on_prepare(&mut self, since: Instant) {
self.prepare_wait += since.elapsed();
}
pub fn waited_on_embed(&mut self, since: Instant) {
self.embed_wait += since.elapsed();
}
pub fn waited_on_walk(&mut self, since: Instant) {
self.walk_ms += since.elapsed().as_millis() as u64;
}
pub fn stages(&self, clocks: &Clocks, lanes: &BTreeMap<String, usize>) -> Stages {
let wall = self.started.elapsed().as_millis() as u64 + self.walk_ms;
let prepare = self.prepare_wait.as_millis() as u64;
let embed = self.embed_wait.as_millis() as u64;
let cpu = clocks.read_all();
let cpu_total: u64 = cpu.iter().sum();
let share = |part: u64| {
if cpu_total == 0 {
0
} else {
(prepare as u128 * part as u128 / cpu_total as u128) as u64
}
};
let mut stages = Stages {
wall_ms: wall,
walk_ms: self.walk_ms,
read_ms: share(cpu[0]),
extract_ms: share(cpu[1]),
parse_ms: share(cpu[2]),
tokenize_ms: share(cpu[3]),
..Stages::default()
};
if cpu_total == 0 {
stages.read_ms = prepare;
}
let chunks: usize = lanes.values().sum();
if embed > 0 {
if chunks == 0 {
stages.embed_wait_ms.insert("cpu".to_string(), embed);
} else {
let mut given = 0;
let last = lanes.len().saturating_sub(1);
for (i, (lane, n)) in lanes.iter().enumerate() {
let ms = if i == last {
embed - given
} else {
(embed as u128 * *n as u128 / chunks as u128) as u64
};
given += ms;
stages.embed_wait_ms.insert(lane.clone(), ms);
}
}
}
let accounted = stages.parts_ms();
stages.write_ms = wall.saturating_sub(accounted);
for (name, us) in ["read+hash", "extract+scan", "parse+chunk", "tokenize"]
.iter()
.zip(cpu)
{
stages.prepare_cpu_ms.insert(name.to_string(), us / 1000);
}
stages
}
}
#[derive(Debug, Clone)]
pub enum Piece {
Ids(Ids),
Text(String),
Cached {
vector: Vec<f32>,
variant: &'static str,
},
}
impl Piece {
pub fn len(&self) -> usize {
match self {
Piece::Ids(ids) => ids.len(),
Piece::Text(text) => text.len() / 4 + 1,
Piece::Cached { .. } => 0,
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
pub struct Context {
pub rechunk: bool,
pub rescan: bool,
pub regraph: bool,
pub allow_secrets: bool,
pub hashes: HashMap<String, String>,
pub uncounted: HashSet<String>,
pub over_cap: HashSet<String>,
pub prehashed: HashMap<PathBuf, (u64, i64, String)>,
pub each: Option<crate::ResolvedBoundary>,
pub tokenizer: Option<tokenizers::Tokenizer>,
pub clocks: Clocks,
pub cache: Option<crate::cache::Scope>,
pub cache_at: Option<crate::cache::Location>,
pub variants: Vec<&'static str>,
pub lookups: AtomicU64,
pub hits: AtomicU64,
pub journal: Option<crate::worker::Journal>,
}
pub enum Prepared {
Refused(crate::Refusal),
Unchanged {
size: u64,
mtime: i64,
read_mtime: Option<i64>,
file_bytes: u64,
},
Unusable {
why: SkipReason,
gone: bool,
file_bytes: u64,
},
Same {
len: u64,
mtime: Option<i64>,
read_mtime: Option<i64>,
file_bytes: u64,
rescan: Option<(String, Vec<keyscan::Match>)>,
regraph: Option<graph::Extraction>,
units: Option<i64>,
},
Image {
bytes: Vec<u8>,
hash: String,
file_bytes: u64,
},
Crashed { reason: String, file_bytes: u64 },
Failed {
error: anyhow::Error,
file_bytes: u64,
},
Skipped { why: SkipReason, file_bytes: u64 },
Text {
len: u64,
hash: String,
file_bytes: u64,
read_mtime: Option<i64>,
text: String,
units: Option<i64>,
found: Vec<keyscan::Match>,
ready: Option<anyhow::Result<Ready>>,
},
}
pub struct Ready {
pub extraction: Option<graph::Extraction>,
pub chunks: Vec<chunk::Chunk>,
pub pieces: Vec<Piece>,
pub hashes: Vec<Option<[u8; 32]>>,
}
fn mtime_of(meta: &std::fs::Metadata) -> i64 {
meta.modified()
.ok()
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map_or(0, |d| d.as_secs() as i64)
}
fn parse<T>(
ctx: &Context,
key: &str,
hash: &str,
call: impl FnOnce() -> Result<T, crate::worker::Stop>,
) -> Result<T, crate::worker::Stop> {
match &ctx.journal {
Some(journal) => journal.around(key, hash, call),
None => call(),
}
}
pub fn prepare(path: &Path, ctx: &Context, cache: Option<&crate::cache::Cache>) -> Prepared {
if let Some(boundary) = &ctx.each
&& let Some(refusal) = boundary.refuses(path, true)
{
return Prepared::Refused(refusal);
}
let _cap = crate::cpucap::slot();
let started = Instant::now();
let key = path.to_string_lossy();
let opened = crate::open_regular(path);
let measured = opened.as_ref().ok().map(|(_, m)| m.clone());
let file_bytes = crate::embeddable_bytes(measured.clone());
if !ctx.rechunk
&& !ctx.rescan
&& !ctx.regraph
&& let (Some((size, mtime, hash)), Some(meta)) = (ctx.prehashed.get(path), &measured)
&& meta.len() == *size
&& mtime_of(meta) == *mtime
&& ctx.hashes.get(key.as_ref()) == Some(hash)
&& !ctx.uncounted.contains(key.as_ref())
{
Clocks::tick(&ctx.clocks.read, started);
return Prepared::Unchanged {
size: *size,
mtime: *mtime,
read_mtime: crate::store::mtime_ns(meta),
file_bytes,
};
}
let cap = if ctx.over_cap.contains(key.as_ref()) {
u64::MAX - 1
} else {
chunk::MAX_FILE_BYTES
};
let unusable = match (&opened, &measured) {
(Err(e), _) => Some(SkipReason::Unreadable(e.to_string())),
(_, None) => Some(SkipReason::Unreadable(
"its metadata could not be read".to_string(),
)),
(_, Some(m)) if !m.is_file() => Some(SkipReason::NotRegular),
(_, Some(m)) if m.len() == 0 => Some(SkipReason::Empty),
(_, Some(m)) if m.len() > cap => Some(SkipReason::TooLarge),
_ => None,
};
if let Some(why) = unusable {
Clocks::tick(&ctx.clocks.read, started);
return Prepared::Unusable {
why,
gone: !path.exists(),
file_bytes,
};
}
let read = opened.and_then(|(file, _)| {
use std::io::Read;
let mut bytes = Vec::new();
file.take(cap + 1).read_to_end(&mut bytes).map(|_| bytes)
});
let bytes = match read {
Ok(bytes) if bytes.len() as u64 <= cap => bytes,
Ok(_) => {
Clocks::tick(&ctx.clocks.read, started);
return Prepared::Unusable {
why: SkipReason::TooLarge,
gone: false,
file_bytes,
};
}
Err(e) => {
Clocks::tick(&ctx.clocks.read, started);
return Prepared::Unusable {
why: SkipReason::Unreadable(e.to_string()),
gone: false,
file_bytes,
};
}
};
let hash = blake3::hash(&bytes).to_hex().to_string();
#[cfg(debug_assertions)]
if std::env::var("SEMLITH_FAULT_DELETE").is_ok_and(|name| {
path.file_name()
.is_some_and(|n| n.to_string_lossy() == name)
}) {
let _ = std::fs::remove_file(path);
}
let now = Clocks::tick(&ctx.clocks.read, started);
let len = bytes.len() as u64;
if let Some(reason) = ctx.journal.as_ref().and_then(|j| j.refused(&key, &hash)) {
return Prepared::Crashed {
reason: reason.to_string(),
file_bytes,
};
}
let same = !ctx.rechunk && ctx.hashes.get(key.as_ref()) == Some(&hash);
if same {
let is_image = image::is_image(path);
let counting = ctx.uncounted.contains(key.as_ref());
let scanning = ctx.rescan && !ctx.allow_secrets && !is_image;
let graphing = ctx.regraph && !is_image && graph::language_of(path).is_some();
let mut extracted = None;
let mut regraph = None;
if scanning || counting || graphing {
match parse(ctx, &key, &hash, || {
crate::worker::read(path, &bytes, graphing)
}) {
Err(crate::worker::Stop::Crashed(reason)) => {
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Crashed { reason, file_bytes };
}
Err(crate::worker::Stop::Unavailable(_)) => {}
Ok(read) => {
regraph = read.graph.and_then(|g| g.ok()).flatten();
extracted = read.text.ok().and_then(|t| t.ok());
}
}
}
let units = counting.then(|| extracted.as_ref().and_then(|(_, n)| *n).unwrap_or(0));
let rescan = extracted.filter(|_| scanning).map(|(text, _)| {
let found = keyscan::scan(&key, &text);
(text, found)
});
let now = Clocks::tick(&ctx.clocks.extract, now);
Clocks::tick(&ctx.clocks.parse, now);
return Prepared::Same {
len,
mtime: measured.as_ref().map(mtime_of),
read_mtime: measured.as_ref().and_then(crate::store::mtime_ns),
file_bytes,
rescan,
regraph,
units,
};
}
if image::is_image(path) {
if image::too_large(&bytes).is_none()
&& let Err(crate::worker::Stop::Crashed(reason)) =
parse(ctx, &key, &hash, || crate::worker::image(path, &bytes))
{
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Crashed { reason, file_bytes };
}
return Prepared::Image {
bytes,
hash,
file_bytes,
};
}
let read = match parse(ctx, &key, &hash, || crate::worker::read(path, &bytes, true)) {
Ok(read) => read,
Err(crate::worker::Stop::Crashed(reason)) => {
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Crashed { reason, file_bytes };
}
Err(crate::worker::Stop::Unavailable(reason)) => {
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Failed {
error: anyhow::anyhow!(reason),
file_bytes,
};
}
};
let (text, units) = match read.text {
Err(panic) => {
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Failed {
error: anyhow::anyhow!(panic),
file_bytes,
};
}
Ok(Err(why)) => {
Clocks::tick(&ctx.clocks.extract, now);
return Prepared::Skipped { why, file_bytes };
}
Ok(Ok(extracted)) => extracted,
};
let graph = read.graph;
drop(bytes);
let found = keyscan::scan(&key, &text);
let now = Clocks::tick(&ctx.clocks.extract, now);
let live = !ctx.allow_secrets && keyscan::refuses(&found);
let ready = (!live).then(|| take_ready(path, &text, graph, ctx, cache, now));
Prepared::Text {
len,
hash,
file_bytes,
read_mtime: measured.as_ref().and_then(crate::store::mtime_ns),
text,
units,
found,
ready,
}
}
pub fn take_ready(
path: &Path,
text: &str,
graph: Option<Result<Option<graph::Extraction>, String>>,
ctx: &Context,
cache: Option<&crate::cache::Cache>,
since: Instant,
) -> anyhow::Result<Ready> {
let graph = match graph {
Some(graph) => graph,
None => {
crate::worker::graph(path, text).map_err(|stop| anyhow::anyhow!("{}", stop.reason()))?
}
};
let extraction = graph.map_err(|e| anyhow::anyhow!(e))?;
let chunks = crate::contained(|| {
chunk::chunk_file(
path,
text,
extraction.as_ref().map_or(&[][..], |e| &e.symbols),
)
})?;
let now = Clocks::tick(&ctx.clocks.parse, since);
let mut pieces = Vec::with_capacity(chunks.len());
let mut hashes = Vec::with_capacity(chunks.len());
for c in &chunks {
let text = c.embedded();
let hash = ctx.cache.as_ref().map(|_| crate::cache::Scope::text(&text));
if let (Some(scope), Some(cache), Some(hash)) = (&ctx.cache, cache, &hash) {
ctx.lookups.fetch_add(1, Ordering::Relaxed);
let found = ctx.variants.iter().find_map(|variant| {
cache
.get(&scope.key(hash, variant))
.filter(|v| v.len() == crate::session::DIM)
.map(|vector| Piece::Cached { vector, variant })
});
if let Some(piece) = found {
ctx.hits.fetch_add(1, Ordering::Relaxed);
pieces.push(piece);
hashes.push(Some(*hash));
continue;
}
}
pieces.push(match &ctx.tokenizer {
Some(tokenizer) => Piece::Ids(crate::session::encode(tokenizer, &text)?),
None => Piece::Text(text),
});
hashes.push(hash);
}
Clocks::tick(&ctx.clocks.tokenize, now);
Ok(Ready {
extraction,
chunks,
pieces,
hashes,
})
}
pub struct Prefetch {
shared: Arc<Shared>,
}
struct Shared {
next: AtomicUsize,
list: Mutex<(Vec<PathBuf>, bool)>,
grown: Condvar,
done: Mutex<(BTreeMap<usize, Prepared>, usize)>,
ready: Condvar,
room: Condvar,
stop: AtomicBool,
lookahead: usize,
}
impl Prefetch {
pub fn start<'scope, 'env>(
scope: &'scope std::thread::Scope<'scope, 'env>,
paths: Vec<PathBuf>,
sealed: bool,
ctx: &'env Context,
threads: usize,
) -> Self {
let threads = threads.clamp(1, 32);
let shared = Arc::new(Shared {
next: AtomicUsize::new(0),
list: Mutex::new((paths, sealed)),
grown: Condvar::new(),
done: Mutex::new((BTreeMap::new(), 0)),
ready: Condvar::new(),
room: Condvar::new(),
stop: AtomicBool::new(false),
lookahead: threads * LOOKAHEAD_PER_THREAD,
});
for n in 0..threads {
let shared = Arc::clone(&shared);
let _ = std::thread::Builder::new()
.name(format!("semlith-prepare-{n}"))
.spawn_scoped(scope, move || {
let _class = crate::priority::indexing_thread();
let cache = ctx.cache_at.as_ref().and_then(crate::cache::Cache::open_in);
loop {
let i = shared.next.fetch_add(1, Ordering::Relaxed);
let path = {
let mut list = shared.list.lock().unwrap_or_else(|e| e.into_inner());
loop {
if shared.stop.load(Ordering::Relaxed) {
return;
}
if let Some(path) = list.0.get(i) {
break path.clone();
}
if list.1 {
return;
}
list = shared
.grown
.wait_timeout(list, Duration::from_millis(100))
.unwrap_or_else(|e| e.into_inner())
.0;
}
};
{
let mut done = shared.done.lock().unwrap_or_else(|e| e.into_inner());
while i >= done.1 + shared.lookahead
&& !shared.stop.load(Ordering::Relaxed)
{
done = shared.room.wait(done).unwrap_or_else(|e| e.into_inner());
}
}
if shared.stop.load(Ordering::Relaxed) {
return;
}
let prepared =
match crate::contained(|| prepare(&path, ctx, cache.as_ref())) {
Ok(prepared) => prepared,
Err(error) => Prepared::Failed {
error,
file_bytes: 0,
},
};
let mut done = shared.done.lock().unwrap_or_else(|e| e.into_inner());
done.0.insert(i, prepared);
shared.ready.notify_all();
}
});
}
Self { shared }
}
pub fn extend_and_seal(&self, more: Vec<PathBuf>) {
let mut list = self.shared.list.lock().unwrap_or_else(|e| e.into_inner());
list.0.extend(more);
list.1 = true;
self.shared.grown.notify_all();
}
pub fn len(&self) -> (usize, bool) {
let list = self.shared.list.lock().unwrap_or_else(|e| e.into_inner());
(list.0.len(), list.1)
}
pub fn path(&self, i: usize) -> Option<PathBuf> {
let list = self.shared.list.lock().unwrap_or_else(|e| e.into_inner());
list.0.get(i).cloned()
}
pub fn from(&self, i: usize) -> Vec<PathBuf> {
let list = self.shared.list.lock().unwrap_or_else(|e| e.into_inner());
list.0.get(i..).map(<[PathBuf]>::to_vec).unwrap_or_default()
}
pub fn take_ticking(&self, i: usize, tick: &mut dyn FnMut()) -> Prepared {
let mut last = Instant::now();
loop {
if let Some(prepared) = self.try_take(i, Duration::from_millis(250)) {
return prepared;
}
if last.elapsed() >= Duration::from_secs(2) {
tick();
last = Instant::now();
}
}
}
fn try_take(&self, i: usize, wait: Duration) -> Option<Prepared> {
let until = Instant::now() + wait;
let mut done = self.shared.done.lock().unwrap_or_else(|e| e.into_inner());
loop {
if let Some(prepared) = done.0.remove(&i) {
done.1 = done.1.max(i + 1);
self.shared.room.notify_all();
return Some(prepared);
}
if Arc::strong_count(&self.shared) == 1 {
return Some(Prepared::Failed {
error: anyhow::anyhow!("the prepare stage lost this file"),
file_bytes: 0,
});
}
let left = until.saturating_duration_since(Instant::now());
if left.is_zero() {
return None;
}
done = self
.shared
.ready
.wait_timeout(done, left.min(Duration::from_millis(50)))
.unwrap_or_else(|e| e.into_inner())
.0;
}
}
pub fn take(&self, i: usize) -> Prepared {
let mut done = self.shared.done.lock().unwrap_or_else(|e| e.into_inner());
loop {
if let Some(prepared) = done.0.remove(&i) {
done.1 = done.1.max(i + 1);
self.shared.room.notify_all();
return prepared;
}
if Arc::strong_count(&self.shared) == 1 {
return Prepared::Failed {
error: anyhow::anyhow!("the prepare stage lost this file"),
file_bytes: 0,
};
}
done = self
.shared
.ready
.wait_timeout(done, Duration::from_millis(50))
.unwrap_or_else(|e| e.into_inner())
.0;
}
}
}
impl Drop for Prefetch {
fn drop(&mut self) {
self.shared.stop.store(true, Ordering::Relaxed);
let _list = self.shared.list.lock().unwrap_or_else(|e| e.into_inner());
self.shared.grown.notify_all();
drop(_list);
let _guard = self.shared.done.lock().unwrap_or_else(|e| e.into_inner());
self.shared.room.notify_all();
self.shared.ready.notify_all();
}
}
pub struct Window {
pub ids: Vec<u64>,
pub pieces: Vec<Piece>,
pub hashes: Vec<Option<[u8; 32]>>,
}
pub struct Embedded {
pub ids: Vec<u64>,
pub vectors: Vec<Vec<f32>>,
pub rows: Vec<(&'static str, bool)>,
pub hashes: Vec<Option<[u8; 32]>>,
pub variants: Vec<(&'static str, usize)>,
pub lanes: Vec<(&'static str, usize)>,
}
pub enum CpuWork {
Ids {
main: crate::session::CpuSession,
alt: Option<crate::session::CpuSession>,
},
Text(Box<fastembed::TextEmbedding>),
}
type Vectors = Result<(Vec<Vec<f32>>, &'static str), String>;
type CpuAnswer = mpsc::Receiver<Vectors>;
enum Answer {
Cpu(CpuAnswer),
Lane(mpsc::Receiver<Result<Vec<Vec<f32>>, String>>, &'static str),
}
impl Answer {
fn wait(&self, timeout: Option<Duration>) -> Result<Vectors, bool> {
let lift = |e: mpsc::RecvTimeoutError| matches!(e, mpsc::RecvTimeoutError::Timeout);
let lift_try = |e: mpsc::TryRecvError| matches!(e, mpsc::TryRecvError::Empty);
match (self, timeout) {
(Answer::Cpu(rx), Some(t)) => rx.recv_timeout(t).map_err(lift),
(Answer::Cpu(rx), None) => rx.try_recv().map_err(lift_try),
(Answer::Lane(rx, variant), Some(t)) => rx
.recv_timeout(t)
.map(|r| r.map(|v| (v, *variant)))
.map_err(lift),
(Answer::Lane(rx, variant), None) => rx
.try_recv()
.map(|r| r.map(|v| (v, *variant)))
.map_err(lift_try),
}
}
}
struct CpuJob {
batch: Vec<Piece>,
reply: mpsc::Sender<Vectors>,
}
pub static CPU_THREADS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
fn cpu_lane(mut work: CpuWork, jobs: mpsc::Receiver<CpuJob>) -> CpuWork {
let mut batches = 0u64;
while let Ok(job) = jobs.recv() {
batches += 1;
crate::cpucap::pace();
let answer = match &mut work {
CpuWork::Ids { main, alt } => {
main.follow_threads();
if let Some(alt) = alt.as_mut() {
alt.follow_threads();
}
CPU_THREADS.store(main.threads(), Ordering::Relaxed);
let rows: Vec<&[u32]> = job
.batch
.iter()
.map(|piece| match piece {
Piece::Ids(ids) => ids.as_slice(),
_ => &[],
})
.collect();
let session = match alt {
Some(alt) if batches.is_multiple_of(2) => alt,
_ => main,
};
let variant = session.variant().name();
session
.embed(&rows)
.map(|v| (v, variant))
.map_err(|e| format!("{e:#}"))
}
CpuWork::Text(model) => {
let texts: Vec<&str> = job
.batch
.iter()
.map(|piece| match piece {
Piece::Text(text) => text.as_str(),
_ => "",
})
.collect();
model
.embed(&texts, Some(texts.len().max(1)))
.map(|mut out| {
for v in &mut out {
crate::normalize(v);
}
(out, crate::embed::Variant::Int8.name())
})
.map_err(|e| format!("{e}"))
}
};
crate::accel::count_cpu(job.batch.len());
let _ = job.reply.send(answer);
}
work
}
#[derive(Clone)]
enum Target {
Cpu,
Lane(Arc<crate::accel::Lane>),
}
impl Target {
fn id(&self) -> &'static str {
match self {
Target::Cpu => "cpu",
Target::Lane(lane) => lane.id,
}
}
}
enum Taken {
CpuOnly,
Retry,
Order,
}
pub fn sound(vector: &[f32]) -> bool {
vector.iter().all(|v| v.is_finite()) && vector.iter().any(|v| *v != 0.0)
}
struct Active {
ids: Vec<u64>,
pieces: Vec<Piece>,
order: Vec<usize>,
front: usize,
back: usize,
retry: VecDeque<usize>,
cpu_only: VecDeque<usize>,
vectors: Vec<Option<Vec<f32>>>,
rows: Vec<(&'static str, bool)>,
hashes: Vec<Option<[u8; 32]>>,
left: usize,
variants: BTreeMap<&'static str, usize>,
lanes: BTreeMap<&'static str, usize>,
}
impl Active {
fn new(mut window: Window) -> Self {
let n = window.pieces.len();
let mut vectors = vec![None; n];
let mut rows = vec![("", false); n];
let mut variants: BTreeMap<&'static str, usize> = BTreeMap::new();
let mut lanes: BTreeMap<&'static str, usize> = BTreeMap::new();
for (at, piece) in window.pieces.iter_mut().enumerate() {
if let Piece::Cached { vector, variant } = piece {
vectors[at] = Some(std::mem::take(vector));
rows[at] = (*variant, true);
*variants.entry(*variant).or_default() += 1;
*lanes.entry("cache").or_default() += 1;
}
}
let mut order: Vec<usize> = (0..n).filter(|i| vectors[*i].is_none()).collect();
if std::env::var(crate::UNSORTED_ENV).is_err() {
order.sort_by_key(|&i| window.pieces[i].len());
}
let left = order.len();
if window.hashes.len() != n {
window.hashes = vec![None; n];
}
Self {
ids: window.ids,
pieces: window.pieces,
back: order.len(),
order,
front: 0,
retry: VecDeque::new(),
cpu_only: VecDeque::new(),
vectors,
rows,
hashes: window.hashes,
left,
variants,
lanes,
}
}
fn untaken(&self, cpu: bool) -> bool {
(cpu && !self.cpu_only.is_empty()) || !self.retry.is_empty() || self.front < self.back
}
fn take(&mut self, from_back: bool, budget: usize, per_call: usize) -> Vec<usize> {
let mut out = Vec::new();
let mut longest = 0;
loop {
let next = if !from_back && let Some(&at) = self.cpu_only.front() {
Some((at, Taken::CpuOnly))
} else if let Some(&at) = self.retry.front() {
Some((at, Taken::Retry))
} else if self.front < self.back {
Some((
if from_back {
self.order[self.back - 1]
} else {
self.order[self.front]
},
Taken::Order,
))
} else {
None
};
let Some((at, taken)) = next else { break };
let len = self.pieces[at].len().max(1);
let cost = (out.len() + 1) * longest.max(len);
if !out.is_empty() && cost > budget && out.len() % per_call.max(1) == 0 {
break;
}
longest = longest.max(len);
out.push(at);
match taken {
Taken::CpuOnly => {
self.cpu_only.pop_front();
}
Taken::Retry => {
self.retry.pop_front();
}
Taken::Order if from_back => self.back -= 1,
Taken::Order => self.front += 1,
}
}
out
}
}
struct Pace {
tokens: u64,
busy: Duration,
}
fn budget(target: &Target, pace: Option<&Pace>) -> usize {
let (floor, start, ceiling) = match target {
Target::Cpu => (512, 2_048, 8_192),
Target::Lane(lane) => lane.token_budget(),
};
match pace {
Some(p) if p.busy.as_secs_f64() > 0.5 => {
let per_second = p.tokens as f64 / p.busy.as_secs_f64();
((per_second * BATCH_SECONDS) as usize).clamp(floor, ceiling)
}
_ => start,
}
}
pub struct Embedder {
windows: Option<mpsc::SyncSender<Window>>,
done: mpsc::Receiver<Result<Embedded, String>>,
pub abort: Arc<AtomicBool>,
pub batched: Arc<AtomicUsize>,
pub outstanding: usize,
}
impl Embedder {
pub fn start<'scope>(
scope: &'scope std::thread::Scope<'scope, '_>,
cpu: CpuWork,
lanes: bool,
cpu_back: mpsc::Sender<CpuWork>,
paused: Arc<AtomicBool>,
) -> Self {
let (windows_tx, windows_rx) = mpsc::sync_channel::<Window>(WRITE_AHEAD);
let (done_tx, done_rx) = mpsc::channel();
let (cpu_tx, cpu_rx) = mpsc::channel::<CpuJob>();
let abort = Arc::new(AtomicBool::new(false));
std::thread::Builder::new()
.name("semlith-cpu-lane".to_string())
.spawn_scoped(scope, move || {
let _class = crate::priority::indexing_thread();
let _ = cpu_back.send(cpu_lane(cpu, cpu_rx));
})
.expect("spawning the CPU lane");
let batched = Arc::new(AtomicUsize::new(0));
let (a, p, b) = (
Arc::clone(&abort),
Arc::clone(&paused),
Arc::clone(&batched),
);
std::thread::Builder::new()
.name("semlith-embed".to_string())
.spawn_scoped(scope, move || {
let _class = crate::priority::indexing_thread();
schedule(windows_rx, done_tx, cpu_tx, lanes, &a, &p, &b);
})
.expect("spawning the embed stage");
Self {
windows: Some(windows_tx),
done: done_rx,
abort,
batched,
outstanding: 0,
}
}
pub fn send(&mut self, window: Window, mut wait: impl FnMut() -> bool) -> bool {
let Some(tx) = &self.windows else {
return false;
};
let mut window = window;
loop {
match tx.try_send(window) {
Ok(()) => {
self.outstanding += 1;
return true;
}
Err(mpsc::TrySendError::Full(back)) => {
window = back;
if !wait() {
return false;
}
}
Err(mpsc::TrySendError::Disconnected(_)) => return false,
}
}
}
pub fn next(&mut self, timeout: Duration) -> Option<Result<Embedded, String>> {
match self.done.recv_timeout(timeout) {
Ok(done) => {
self.outstanding = self.outstanding.saturating_sub(1);
Some(done)
}
Err(mpsc::RecvTimeoutError::Timeout) => None,
Err(mpsc::RecvTimeoutError::Disconnected) => {
self.outstanding = 0;
Some(Err("the embed stage ended".to_string()))
}
}
}
pub fn close(&mut self) {
self.windows = None;
}
pub fn cancel(&mut self) {
self.abort.store(true, Ordering::Relaxed);
self.windows = None;
}
}
fn schedule(
windows: mpsc::Receiver<Window>,
done: mpsc::Sender<Result<Embedded, String>>,
cpu: mpsc::Sender<CpuJob>,
may_use_lanes: bool,
abort: &AtomicBool,
paused: &AtomicBool,
batched: &AtomicUsize,
) {
let mut active: VecDeque<Active> = VecDeque::new();
let mut closed = false;
let mut flying: Vec<(Target, usize, Vec<usize>, Answer, Instant, usize)> = Vec::new();
let mut base = 0usize;
let mut pace: HashMap<&'static str, Pace> = HashMap::new();
let mut last_back: HashMap<&'static str, Instant> = HashMap::new();
loop {
if abort.load(Ordering::Relaxed) {
return;
}
while !closed && active.len() < ACTIVE_WINDOWS {
let got = if active.is_empty() && flying.is_empty() {
windows.recv_timeout(Duration::from_millis(20))
} else {
windows.try_recv().map_err(|e| match e {
mpsc::TryRecvError::Empty => mpsc::RecvTimeoutError::Timeout,
mpsc::TryRecvError::Disconnected => mpsc::RecvTimeoutError::Disconnected,
})
};
match got {
Ok(window) => active.push_back(Active::new(window)),
Err(mpsc::RecvTimeoutError::Timeout) => break,
Err(mpsc::RecvTimeoutError::Disconnected) => closed = true,
}
}
if closed && active.is_empty() && flying.is_empty() {
return;
}
if !paused.load(Ordering::Relaxed) {
let (lanes, cpu_on) = if may_use_lanes {
crate::accel::for_run()
} else {
(Vec::new(), true)
};
let mut targets: Vec<Target> = Vec::new();
for lane in lanes {
if lane.ready() {
targets.push(Target::Lane(lane));
} else {
lane.wake();
}
}
let guard = active.iter().any(|w| !w.cpu_only.is_empty());
if cpu_on || guard {
targets.push(Target::Cpu);
}
if targets.is_empty() && flying.is_empty() {
std::thread::sleep(Duration::from_millis(50));
}
for target in &targets {
let from_back = !matches!(target, Target::Cpu);
let guard_only =
!from_back && !cpu_on && !targets.iter().all(|t| matches!(t, Target::Cpu));
while flying
.iter()
.filter(|(t, ..)| t.id() == target.id())
.count()
< LANE_DEPTH
{
let Some(slot) = active.iter().position(|w| {
if guard_only {
!w.cpu_only.is_empty()
} else {
w.untaken(!from_back)
}
}) else {
break;
};
let want = budget(target, pace.get(target.id()));
let window = &mut active[slot];
let group = if guard_only {
let n = window.cpu_only.len();
window.cpu_only.drain(..n).collect()
} else {
let per_call = match target {
Target::Lane(lane) => lane.rows_per_call(),
Target::Cpu => 1,
};
window.take(from_back, want, per_call)
};
if group.is_empty() {
break;
}
let tokens: usize = group
.iter()
.map(|&i| window.pieces[i].len())
.max()
.unwrap_or(0)
* group.len();
let batch: Vec<Piece> =
group.iter().map(|&i| window.pieces[i].clone()).collect();
let answer: Answer = match target {
Target::Cpu => {
let (reply, answer) = mpsc::channel();
if cpu.send(CpuJob { batch, reply }).is_err() {
let _ = done.send(Err("the CPU lane has gone".to_string()));
return;
}
Answer::Cpu(answer)
}
Target::Lane(lane) => {
let ids: Vec<Ids> = batch
.into_iter()
.map(|piece| match piece {
Piece::Ids(ids) => ids,
_ => Vec::new(),
})
.collect();
Answer::Lane(lane.submit(ids), lane.variant())
}
};
flying.push((
target.clone(),
base + slot,
group,
answer,
Instant::now(),
tokens,
));
}
}
}
let mut moved = false;
let mut still = Vec::with_capacity(flying.len());
let oldest = flying.first().map(|f| f.4);
for (target, seq, group, answer, sent, tokens) in flying.drain(..) {
let got =
answer.wait((Some(sent) == oldest && !moved).then_some(Duration::from_millis(5)));
let window = &mut active[seq - base];
match got {
Ok(Ok((vectors, variant))) if vectors.len() == group.len() => {
moved = true;
let entry = pace.entry(target.id()).or_insert(Pace {
tokens: 0,
busy: Duration::ZERO,
});
entry.tokens += tokens as u64;
let now = Instant::now();
let from = last_back
.get(target.id())
.map_or(sent, |&back| back.max(sent));
entry.busy += now.saturating_duration_since(from);
last_back.insert(target.id(), now);
let mut kept = 0;
for (at, vector) in group.iter().zip(vectors) {
if !sound(&vector) && !matches!(target, Target::Cpu) {
window.cpu_only.push_back(*at);
continue;
}
window.vectors[*at] = Some(vector);
window.rows[*at] = (variant, false);
kept += 1;
}
window.left -= kept;
batched.fetch_add(kept, Ordering::Relaxed);
*window.variants.entry(variant).or_default() += kept;
*window.lanes.entry(target.id()).or_default() += kept;
}
Err(true)
if matches!(&target, Target::Lane(lane) if !lane.ready())
&& sent.elapsed() > LANE_GONE =>
{
moved = true;
window.retry.extend(group);
}
Err(true) => still.push((target, seq, group, answer, sent, tokens)),
_ => {
moved = true;
window.retry.extend(group);
}
}
}
flying = still;
while active.front().is_some_and(|w| w.left == 0) {
let window = active.pop_front().expect("checked above");
base += 1;
let vectors: Vec<Vec<f32>> = window
.vectors
.into_iter()
.map(|v| v.unwrap_or_default())
.collect();
let finished = Embedded {
ids: window.ids,
vectors,
rows: window.rows,
hashes: window.hashes,
variants: window.variants.into_iter().collect(),
lanes: window.lanes.into_iter().collect(),
};
if done.send(Ok(finished)).is_err() {
return;
}
}
if paused.load(Ordering::Relaxed) && flying.is_empty() {
std::thread::sleep(Duration::from_millis(20));
}
}
}
pub fn prepare_threads() -> usize {
crate::embed::embed_threads().max(1)
}
#[cfg(test)]
mod tests {
use super::*;
fn window(lens: &[usize]) -> Active {
Active::new(Window {
ids: (0..lens.len() as u64).collect(),
pieces: lens.iter().map(|n| Piece::Ids(vec![1; *n])).collect(),
hashes: Vec::new(),
})
}
#[test]
fn a_cached_piece_is_done_before_the_window_starts() {
let mut w = Active::new(Window {
ids: vec![7, 8],
pieces: vec![
Piece::Cached {
vector: vec![1.0; 3],
variant: "fp16-ane",
},
Piece::Ids(vec![1; 5]),
],
hashes: vec![Some([1; 32]), Some([2; 32])],
});
assert_eq!(w.left, 1);
assert_eq!(w.take(true, 100, 1), vec![1]);
assert!(!w.untaken(true));
assert_eq!(w.rows[0], ("fp16-ane", true));
assert_eq!(w.lanes["cache"], 1);
}
#[test]
fn a_fixed_batch_lane_takes_whole_calls_of_rows() {
let mut w = window(&[100; 20]);
assert_eq!(w.take(true, 200, 8).len(), 8);
assert_eq!(w.take(true, 300, 8).len(), 8);
assert_eq!(w.take(true, 200, 8).len(), 4);
}
#[test]
fn the_cpu_takes_the_shortest_and_a_lane_the_longest_within_a_budget() {
let mut w = window(&[10, 300, 20, 200, 30]);
assert_eq!(w.take(false, 90, 1), vec![0, 2, 4]);
assert_eq!(w.take(true, 500, 1), vec![1]);
assert_eq!(w.take(true, 500, 1), vec![3]);
assert!(!w.untaken(false));
}
#[test]
fn a_batch_always_takes_one_even_over_budget_and_retries_come_first() {
let mut w = window(&[1000, 5]);
assert_eq!(w.take(true, 10, 1), vec![0]);
w.retry.push_back(0);
assert_eq!(w.take(false, 10, 1), vec![0]);
assert_eq!(w.take(false, 10, 1), vec![1]);
}
#[test]
fn a_zero_or_non_finite_vector_is_unsound() {
assert!(sound(&[0.1, 0.0]));
assert!(!sound(&[0.0, 0.0]));
assert!(!sound(&[f32::NAN, 1.0]));
assert!(!sound(&[f32::INFINITY, 1.0]));
}
#[test]
fn only_the_cpu_takes_what_the_guard_handed_back() {
let mut w = window(&[10, 20, 30]);
assert_eq!(w.take(true, 30, 1), vec![2]);
w.cpu_only.push_back(2);
assert!(!w.untaken(false) || w.front < w.back);
assert_eq!(w.take(true, 1000, 1), vec![1, 0]);
assert!(w.untaken(true));
assert_eq!(w.take(false, 1000, 1), vec![2]);
}
#[test]
fn the_stages_sum_to_the_wall_time() {
let clock = WriterClock {
started: Instant::now() - Duration::from_millis(1_000),
walk_ms: 100,
prepare_wait: Duration::from_millis(300),
embed_wait: Duration::from_millis(400),
};
let clocks = Clocks::default();
clocks.read.store(100_000, Ordering::Relaxed);
clocks.parse.store(300_000, Ordering::Relaxed);
let lanes: BTreeMap<String, usize> =
[("ane".to_string(), 300), ("cpu".to_string(), 100)].into();
let stages = clock.stages(&clocks, &lanes);
assert_eq!(stages.parts_ms(), stages.wall_ms);
assert_eq!(stages.read_ms, 75);
assert_eq!(stages.parse_ms, 225);
assert_eq!(stages.embed_wait_ms["ane"], 300);
assert_eq!(stages.embed_wait_ms["cpu"], 100);
}
}