use std::collections::BTreeMap;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::{Duration, Instant};
use haem_frames::embedding::Distance;
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::{Value, json};
use crate::serve::embed::job::{Bounds, Embed};
use crate::serve::embed::question::Ask;
use crate::serve::embed::settings::Restart;
use crate::serve::embed::{Answer, Embedded, ModelId};
use crate::serve::embedding::Worker;
use crate::serve::protocol::{
Change, CollectionWrite, DocumentWrite, Metric, Mutation, Operation, QueryPage, QueryRequest,
Selection, StoreRead, VectorDeclaration, VectorSearchText,
};
use crate::serve::query::{Rows, Settings, probe};
use crate::serve::service::Service;
use crate::serve::storage::embedded::{Declaration, ModelName, VectorKind};
use crate::serve::storage::{Failure, Storage};
pub type TestResult = Result<(), Box<dyn std::error::Error>>;
const MANIFEST: &str = "0000000000000000000000000000000000000000000000000000000000005eed";
pub const TEXT: &str = "chest pain";
pub const QUERY: &str = "notes | similar(.body_vec; \"chest pain\"; cosine) | limit(10)";
pub struct WorkerFixture {
pub asked: AtomicUsize,
service: Mutex<Option<Weak<Service>>>,
}
impl WorkerFixture {
fn embedded(texts: &[String]) -> Result<Embedded, Failure> {
Ok(Embedded {
model: ModelId {
name: "stand-in".to_owned(),
revision: "1".to_owned(),
manifest_sha256: MANIFEST.to_owned(),
dimensions: 2,
distance: Distance::Cosine,
normalised: false,
},
answers: texts
.iter()
.map(|text| {
let vector = match text.as_str() {
"chest pain" | "same" => vec![1.0, 0.0],
"near" => vec![1.0, 1.0],
"side" => vec![0.0, 1.0],
"opposite" => vec![-1.0, 0.0],
_ => return Err(Failure::new("fixture_text", "unserved fixture text")),
};
Ok(Answer::Vector(vector))
})
.collect::<Result<_, Failure>>()?,
})
}
}
impl Embed for WorkerFixture {
fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
Self::embedded(texts)
}
}
impl Ask for WorkerFixture {
fn question(&self, text: &str) -> Result<Embedded, Failure> {
if text == "ready" {
return Self::embedded(&[]);
}
if text != TEXT {
return Err(Failure::new(
"fixture_question",
"the question text changed",
));
}
let service = self
.service
.lock()
.map_err(|_| Failure::new("fixture_lock", "the fixture lock was poisoned"))?
.as_ref()
.and_then(Weak::upgrade)
.ok_or_else(|| Failure::new("fixture_service", "the fixture has no service"))?;
service.reclaim_status()?;
self.asked.fetch_add(1, Ordering::SeqCst);
Self::embedded(&[text.to_owned()])
}
}
pub struct Fixture {
pub service: Arc<Service>,
pub worker: Arc<WorkerFixture>,
directory: tempfile::TempDir,
}
impl Fixture {
pub fn new(configured: bool) -> Result<Self, Box<dyn std::error::Error>> {
Self::with_scan(configured, 100)
}
pub fn with_scan(configured: bool, scan: u64) -> Result<Self, Box<dyn std::error::Error>> {
let directory = tempfile::tempdir()?;
let storage = storage(directory.path())?;
let settings = settings(directory.path(), scan)?;
let mut service = Service::with_query(
storage,
NonZeroU32::MIN.saturating_add(99),
Some(probe::prepare(
settings,
NonZeroU32::MIN.saturating_add((1 << 20) - 1),
)),
)?;
initialize(&service)?;
let worker = Arc::new(WorkerFixture {
asked: AtomicUsize::new(0),
service: Mutex::new(None),
});
if configured {
let bounds = Bounds {
page_records: NonZeroUsize::MIN.saturating_add(9),
batch: NonZeroUsize::MIN.saturating_add(9),
input_bytes: NonZeroUsize::MIN.saturating_add(4095),
call: Duration::from_secs(1),
};
let embedder: Arc<dyn Worker> = worker.clone();
service.start_embedding_steps(embedder, bounds, Restart::default())?;
while service.embedding_step_at(bounds, Instant::now())? {}
}
let service = Arc::new(service);
*worker.service.lock().map_err(|_| "fixture lock poisoned")? =
Some(Arc::downgrade(&service));
Ok(Self {
service,
worker,
directory,
})
}
pub fn query(&self, request: QueryRequest) -> Result<Value, Box<dyn std::error::Error>> {
Ok(self.service.dispatch(Operation::QueryRun(request))?)
}
pub fn search(&self) -> Result<Value, Box<dyn std::error::Error>> {
Ok(self
.service
.dispatch(Operation::VectorSearchText(VectorSearchText {
selection: selection(),
collection: "notes".to_owned(),
field: "body_vec".to_owned(),
text: TEXT.to_owned(),
k: NonZeroU32::MIN.saturating_add(9),
}))?)
}
pub fn no_query_worker(&self) -> TestResult {
std::fs::remove_file(self.directory.path().join("query-worker"))?;
Ok(())
}
}
fn mutate(service: &Service, change: Change) -> TestResult {
let head = service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(),
}))?;
let expected = serde_json::from_value(head["selection"]["version"].clone())?;
service.dispatch(Operation::Mutate(Mutation {
request_id: format!("search-{}", serde_json::to_string(&expected)?),
branch: "main".to_owned(),
expected,
change,
}))?;
Ok(())
}
pub fn selection() -> Selection {
Selection {
branch: "main".to_owned(),
version: None,
}
}
pub fn request(text: &str, rows: u64) -> Result<QueryRequest, Box<dyn std::error::Error>> {
Ok(QueryRequest {
selection: selection(),
query: text.to_owned(),
parameters: BTreeMap::new(),
page: QueryPage {
rows: NonZeroU64::new(rows).ok_or("zero rows")?,
bytes: NonZeroU64::MIN.saturating_add((1 << 20) - 1),
},
continuation: None,
rows: Rows::Values,
})
}
fn executable(directory: &Path) -> Result<PathBuf, Box<dyn std::error::Error>> {
let executable = directory.join("query-worker");
std::fs::write(
&executable,
r"#!/usr/bin/python3
import json, struct, sys
def receive():
header = sys.stdin.buffer.read(4)
if not header:
return None
return sys.stdin.buffer.read(struct.unpack('>I', header)[0])
def send(tag, value):
data = bytes([tag]) + json.dumps(value, separators=(',', ':')).encode()
sys.stdout.buffer.write(struct.pack('>I', len(data)) + data)
sys.stdout.buffer.flush()
start = json.loads(receive()[1:])
reading = 'alternate' if start.get('alternate') is not None else 'candidate'
program = start['alternate'] if reading == 'alternate' else start['filter']
if program != '.':
raise ValueError('unserved fixture filter: ' + program)
send(3, dict(protocol=start['protocol'], worker='@WORKER_VERSION@', core_limit=0,
descriptors_closed=0, environment=0, engine='fixture',
jaq_core='0', jaq_std='0', jaq_json='0', reading=reading))
while True:
frame = receive()
if frame is None:
break
send(4, json.loads(frame[1:])['value'])
send(5, dict(outputs=1))
"
.replace("@WORKER_VERSION@", env!("CARGO_PKG_VERSION")),
)?;
std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o700))?;
Ok(executable)
}
fn storage(directory: &Path) -> Result<Storage, Box<dyn std::error::Error>> {
Ok(Storage::create(
DatabaseConfig {
data_dir: directory.join("database"),
shard_count: 2,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
Duration::from_secs(1),
)?)
}
fn settings(directory: &Path, scan: u64) -> Result<Settings, Box<dyn std::error::Error>> {
Ok(Settings {
time_ms: NonZeroU64::MIN.saturating_add(59_999),
memory_bytes: NonZeroU64::MIN.saturating_add((1 << 26) - 1),
outputs_per_record: NonZeroU64::MIN.saturating_add(31),
scan_records: NonZeroU64::new(scan).ok_or("zero scan bound")?,
max_running: NonZeroUsize::MIN,
worker: executable(directory)?,
})
}
fn initialize(service: &Service) -> TestResult {
mutate(
service,
Change::Collection(CollectionWrite {
collection: "notes".to_owned(),
schema: None,
}),
)?;
for (id, body) in [
("c", "same"),
("a", "same"),
("b", "near"),
("d", "side"),
("e", "opposite"),
] {
mutate(
service,
Change::Put(DocumentWrite {
collection: "notes".to_owned(),
document_id: id.to_owned(),
document: serde_json::value::to_raw_value(&json!({"body": body}))?,
}),
)?;
}
mutate(
service,
Change::VectorDeclare(VectorDeclaration {
collection: "notes".to_owned(),
field: "body_vec".to_owned(),
declaration: Declaration {
kind: VectorKind::Dense,
source: "body".to_owned(),
model: ModelName {
name: "stand-in".to_owned(),
revision: "1".to_owned(),
manifest_sha256: MANIFEST.to_owned(),
},
dimensions: NonZeroUsize::MIN.saturating_add(1),
distance: Metric::Cosine,
},
}),
)?;
Ok(())
}