src/vector/remote.rs
//! **An embedder on another box** — a GPU worker lends its forward passes to
//! a server that has none.
//!
//! ```text
//! server (no GPU) worker (GPU)
//! embeddings feed: cut fn texts, content hash, `nornir vector serve-embed`
//! what has no vector yet ──── Embedder.Embed ────▶ ort CUDA forward passes
//! store vectors + links + ledger ◀── f32 rows ──── (no warehouse, no state)
//! ```
//!
//! Everything that decides WHAT is embedded stays where the room and the
//! checkout are: the cut by the driver's span, the content hash, the delta
//! against the link state, the ledger row. The worker only turns texts into
//! vectors, so a commit that changes one function sends one text, and the
//! vectors land in the server's own warehouse under the same keys a local
//! GPU would have written.
//!
//! **One model space.** The worker says its [`ModelProfile`] first; a server
//! refuses a worker whose profile id is not the one it stores under (another
//! model, other weights), and the worker refuses a request that names
//! another space. Two spaces never meet in one table.
//!
//! Selected by `$NORNIR_EMBED_REMOTE` (`host:port` or `http://host:port`) on
//! the server; `$NORNIR_EMBED_REMOTE_TOKEN` or `$NORNIR_EMBED_REMOTE_TOKEN_FILE`
//! is the worker's bearer (none = an open worker).
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use anyhow::{Context, Result, anyhow, bail, ensure};
use nornir_wire::pb;
use super::store::{Embedder, ModelProfile};
/// The worker's address on the server that borrows it.
pub const REMOTE_ENV: &str = "NORNIR_EMBED_REMOTE";
/// The worker's bearer, inline.
pub const TOKEN_ENV: &str = "NORNIR_EMBED_REMOTE_TOKEN";
/// The worker's bearer, in a file (first line).
pub const TOKEN_FILE_ENV: &str = "NORNIR_EMBED_REMOTE_TOKEN_FILE";
/// Texts per RPC: a slice of the feed (2048 functions, several pieces each)
/// travels as several calls, each well inside one message.
pub const TEXTS_PER_CALL: usize = 256;
/// Bytes of text per RPC (pieces are ≤ 8 kB, so a call is a few MB at most).
pub const BYTES_PER_CALL: usize = 4 << 20;
/// gRPC message cap both ends set: 256 texts × 8 kB in, 256 × 768 × 4 B out,
/// with room to spare.
pub const MESSAGE_LIMIT: usize = 64 << 20;
/// The worker this process is told to use, normalized to a URL, or `None`.
pub fn configured() -> Option<String> {
let raw = std::env::var(REMOTE_ENV).ok()?;
let raw = raw.trim();
if raw.is_empty() {
return None;
}
Some(if raw.starts_with("http://") || raw.starts_with("https://") { raw.to_string() } else { format!("http://{raw}") })
}
/// The bearer for the worker: `$NORNIR_EMBED_REMOTE_TOKEN`, else the first
/// line of `$NORNIR_EMBED_REMOTE_TOKEN_FILE`, else none (an open worker).
pub fn token() -> Result<String> {
if let Ok(t) = std::env::var(TOKEN_ENV) {
if !t.trim().is_empty() {
return Ok(t.trim().to_string());
}
}
match std::env::var(TOKEN_FILE_ENV) {
Ok(f) if !f.trim().is_empty() => {
let text = std::fs::read_to_string(f.trim()).with_context(|| format!("read the embed worker's token from {} (${TOKEN_FILE_ENV})", f.trim()))?;
Ok(text.lines().next().unwrap_or("").trim().to_string())
}
_ => Ok(String::new()),
}
}
/// The wire form of a profile.
pub fn profile_to_wire(p: &ModelProfile, model_id: &str, backend: &str, device: &str) -> pb::EmbedProfile {
pb::EmbedProfile {
model_name: p.model_name.clone(),
weights_sha: p.weights_sha.clone(),
tokenizer_sha: p.tokenizer_sha.clone(),
pooling: p.pooling.clone(),
normalize: p.normalize,
dim: p.dim as u32,
dtype: p.dtype.clone(),
profile_id: p.id(),
model_id: model_id.to_string(),
backend: backend.to_string(),
device: device.to_string(),
}
}
/// The profile a worker described, checked: its fields hash to the id it says.
pub fn profile_from_wire(w: &pb::EmbedProfile) -> Result<ModelProfile> {
let p = ModelProfile {
model_name: w.model_name.clone(),
weights_sha: w.weights_sha.clone(),
tokenizer_sha: w.tokenizer_sha.clone(),
pooling: w.pooling.clone(),
normalize: w.normalize,
dim: w.dim as usize,
dtype: w.dtype.clone(),
};
ensure!(p.id() == w.profile_id, "the embed worker's profile does not hash to the id it names ({} vs {})", p.id(), w.profile_id);
Ok(p)
}
/// `count × dim` little-endian f32 into rows.
pub fn rows_from_bytes(values: &[u8], count: usize, dim: usize) -> Result<Vec<Vec<f32>>> {
ensure!(values.len() == count * dim * 4, "embed reply carries {} bytes for {count} × {dim} f32", values.len());
Ok(values
.chunks_exact(dim * 4)
.map(|row| row.chunks_exact(4).map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])).collect())
.collect())
}
/// Rows into `count × dim` little-endian f32.
pub fn rows_to_bytes(rows: &[Vec<f32>]) -> Vec<u8> {
let mut out = Vec::with_capacity(rows.iter().map(|r| r.len() * 4).sum());
for r in rows {
for x in r {
out.extend_from_slice(&x.to_le_bytes());
}
}
out
}
/// Cut `texts` into calls of at most [`TEXTS_PER_CALL`] texts and
/// [`BYTES_PER_CALL`] bytes (a single bigger text is a call of its own).
pub fn calls(texts: &[&str]) -> Vec<std::ops::Range<usize>> {
let mut out = Vec::new();
let (mut start, mut bytes) = (0usize, 0usize);
for (i, t) in texts.iter().enumerate() {
if i > start && (i - start == TEXTS_PER_CALL || bytes + t.len() > BYTES_PER_CALL) {
out.push(start..i);
(start, bytes) = (i, 0);
}
bytes += t.len();
}
if start < texts.len() {
out.push(start..texts.len());
}
out
}
type Bearer = Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>;
/// **The client**: an [`Embedder`] whose forward passes run on the worker.
pub struct RemoteEmbedder {
endpoint: String,
bearer: Bearer,
profile: ModelProfile,
worker: pb::EmbedProfile,
rt: Rt,
channel: tonic::transport::Channel,
/// Forward-pass time the worker reported, summed (ms).
worker_ms: AtomicU64,
}
impl RemoteEmbedder {
/// Connect to the worker at `endpoint`, read its profile and refuse it
/// unless it is `expect` (the space this process stores under).
pub fn connect(endpoint: &str, token: &str, expect: &ModelProfile) -> Result<Self> {
let bearer = nornir_service::client::bearer(token).context("the embed worker's bearer")?;
let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().context("tokio runtime for the embed worker client")?;
let endpoint = endpoint.to_string();
let (channel, worker) = off_runtime(|| {
rt.block_on(async {
let channel = tonic::transport::Channel::from_shared(endpoint.clone())
.with_context(|| format!("embed worker url `{endpoint}`"))?
.connect_timeout(std::time::Duration::from_secs(10))
.timeout(std::time::Duration::from_secs(600))
.connect()
.await
.with_context(|| format!("connect to the embed worker at {endpoint}"))?;
let mut c = client(channel.clone(), bearer.clone());
let worker = c.profile(pb::Empty {}).await.map_err(|s| anyhow!("embed worker {endpoint}: Profile: {}", s.message()))?.into_inner();
Ok::<_, anyhow::Error>((channel, worker))
})
})?;
let profile = profile_from_wire(&worker)?;
ensure!(
profile.id() == expect.id(),
"the embed worker at {endpoint} embeds in another model space: {} ({}, weights {}) — this process stores under {} ({}, weights {}). \
Run the worker with the same model ($NORNIR_EMBED_MODEL) and weights",
&worker.profile_id[..worker.profile_id.len().min(12)],
worker.model_name,
&worker.weights_sha[..worker.weights_sha.len().min(12)],
&expect.id()[..12],
expect.model_name,
&expect.weights_sha[..expect.weights_sha.len().min(12)],
);
Ok(Self { endpoint, bearer, profile, worker, rt: Rt(Some(rt)), channel, worker_ms: AtomicU64::new(0) })
}
/// [`Self::connect`] to the worker `$NORNIR_EMBED_REMOTE` names, in the
/// space this build stores under.
#[cfg(any(feature = "embed-tract", feature = "embed-ort"))]
pub fn from_env() -> Result<Self> {
let endpoint = configured().ok_or_else(|| anyhow!("${REMOTE_ENV} is not set"))?;
Self::connect(&endpoint, &token()?, &super::embed_support::profile())
}
/// What the worker said it is.
pub fn worker(&self) -> &pb::EmbedProfile {
&self.worker
}
pub fn endpoint(&self) -> &str {
&self.endpoint
}
/// Forward-pass time the worker reported for this client so far.
pub fn worker_ms(&self) -> u64 {
self.worker_ms.load(Ordering::Relaxed)
}
/// One line naming the worker: where, on what.
pub fn label(&self) -> String {
format!("remote {} ({})", self.endpoint, self.worker.backend)
}
}
/// The client's own runtime, let go without blocking: the embedder may be
/// dropped on a thread inside another runtime (a server's), where a plain
/// drop of a runtime panics.
struct Rt(Option<tokio::runtime::Runtime>);
impl Rt {
fn get(&self) -> &tokio::runtime::Runtime {
self.0.as_ref().expect("the runtime lives as long as the client")
}
}
impl Drop for Rt {
fn drop(&mut self) {
if let Some(rt) = self.0.take() {
rt.shutdown_background();
}
}
}
/// A gRPC client with the bearer stamped and the message caps raised.
fn client(channel: tonic::transport::Channel, bearer: Bearer) -> pb::embedder_client::EmbedderClient<tonic::service::interceptor::InterceptedService<tonic::transport::Channel, impl tonic::service::Interceptor>> {
pb::embedder_client::EmbedderClient::with_interceptor(channel, move |mut req: tonic::Request<()>| {
if let Some(b) = &bearer {
req.metadata_mut().insert("authorization", b.clone());
}
Ok(req)
})
.max_decoding_message_size(MESSAGE_LIMIT)
.max_encoding_message_size(MESSAGE_LIMIT)
}
/// Run `f` where a `block_on` is allowed: on this thread, or — when this
/// thread is inside a tokio runtime — on a scoped thread of its own.
fn off_runtime<T: Send>(f: impl FnOnce() -> T + Send) -> T {
if tokio::runtime::Handle::try_current().is_err() {
return f();
}
std::thread::scope(|s| s.spawn(f).join().unwrap_or_else(|p| std::panic::resume_unwind(p)))
}
impl Embedder for RemoteEmbedder {
fn profile(&self) -> ModelProfile {
self.profile.clone()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
let dim = self.profile.dim;
let profile_id = self.worker.profile_id.clone();
let mut out: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
off_runtime(|| {
self.rt.get().block_on(async {
let mut c = client(self.channel.clone(), self.bearer.clone());
for range in calls(texts) {
let n = range.len();
let req = pb::EmbedRequest { texts: texts[range].iter().map(|t| t.to_string()).collect(), profile_id: profile_id.clone() };
let reply = c.embed(req).await.map_err(|s| anyhow!("embed worker {}: Embed of {n} text(s): {}", self.endpoint, s.message()))?.into_inner();
ensure!(reply.count as usize == n && reply.dim as usize == dim, "embed worker {} answered {} × {} for {n} × {dim}", self.endpoint, reply.count, reply.dim);
out.extend(rows_from_bytes(&reply.values, n, dim)?);
self.worker_ms.fetch_add(reply.embed_ms, Ordering::Relaxed);
}
Ok::<_, anyhow::Error>(())
})
})?;
Ok(out)
}
}
// ── the worker ──────────────────────────────────────────────────────────────
/// What a worker has served since it started.
#[derive(Default)]
struct Served {
calls: AtomicU64,
texts: AtomicU64,
bytes: AtomicU64,
embed_ms: AtomicU64,
}
/// **The worker**: one loaded embedder behind `Embedder.Embed`. Requests are
/// embedded one at a time (the device is one; the embedder batches inside a
/// request).
pub struct Worker {
embedder: Arc<dyn Embedder>,
profile: pb::EmbedProfile,
turn: Arc<tokio::sync::Mutex<()>>,
served: Arc<Served>,
/// The same model on the CPU, for the rare text the device cannot hold
/// by itself (attention is quadratic in tokens: an 8k-token text needs
/// more than 16 GB). Loaded on the first such text.
fallback: Arc<std::sync::OnceLock<std::result::Result<Arc<dyn Embedder>, String>>>,
load_fallback: Option<fn() -> Result<Box<dyn Embedder>>>,
}
impl Worker {
pub fn new(embedder: Arc<dyn Embedder>, model_id: &str, backend: &str, device: &str) -> Self {
let profile = profile_to_wire(&embedder.profile(), model_id, backend, device);
Self { embedder, profile, turn: Arc::new(tokio::sync::Mutex::new(())), served: Arc::new(Served::default()), fallback: Arc::default(), load_fallback: None }
}
/// A text the device cannot hold alone is embedded by what `load` gives
/// (the CPU backend of the same weights; the same model space — refused
/// otherwise).
pub fn with_fallback(mut self, load: fn() -> Result<Box<dyn Embedder>>) -> Self {
self.load_fallback = Some(load);
self
}
pub fn profile(&self) -> &pb::EmbedProfile {
&self.profile
}
}
#[tonic::async_trait]
impl pb::embedder_server::Embedder for Worker {
async fn profile(&self, _req: tonic::Request<pb::Empty>) -> Result<tonic::Response<pb::EmbedProfile>, tonic::Status> {
Ok(tonic::Response::new(self.profile.clone()))
}
async fn embed(&self, req: tonic::Request<pb::EmbedRequest>) -> Result<tonic::Response<pb::EmbedReply>, tonic::Status> {
let req = req.into_inner();
if !req.profile_id.is_empty() && req.profile_id != self.profile.profile_id {
return Err(tonic::Status::failed_precondition(format!(
"this worker embeds in model space {} ({}); the request names {}",
self.profile.profile_id, self.profile.model_name, req.profile_id
)));
}
let n = req.texts.len();
let bytes: usize = req.texts.iter().map(String::len).sum();
let _turn = self.turn.lock().await;
let embedder = self.embedder.clone();
let (fallback, load_fallback, profile_id) = (self.fallback.clone(), self.load_fallback, self.profile.profile_id.clone());
let t0 = std::time::Instant::now();
let rows = tokio::task::spawn_blocking(move || {
let texts: Vec<&str> = req.texts.iter().map(String::as_str).collect();
embedder.embed(&texts).or_else(|whole| {
// one text too big for the device next to the others (the
// batch pads to its longest): each alone; a text the device
// cannot hold even alone goes to the CPU fallback
eprintln!("nornir-embed-worker: a batch of {} failed ({}); embedding them one by one", texts.len(), first_line(&whole));
texts
.iter()
.map(|t| {
let one = std::slice::from_ref(t);
match embedder.embed(one) {
Ok(mut v) => Ok(v.remove(0)),
Err(e) => {
let Some(load) = load_fallback else { return Err(e) };
let cpu = fallback
.get_or_init(|| {
let f: Arc<dyn Embedder> = Arc::from(load().map_err(|e| format!("{e:#}"))?);
if f.profile().id() != profile_id {
return Err("the fallback embedder is in another model space".to_string());
}
Ok(f)
})
.as_ref()
.map_err(|why| anyhow!("{} — and no CPU fallback: {why}", first_line(&e)))?;
eprintln!("nornir-embed-worker: a {}-byte text does not fit the device alone ({}); embedded on the CPU", t.len(), first_line(&e));
Ok(cpu.embed(one)?.remove(0))
}
}
})
.collect::<Result<Vec<_>>>()
})
})
.await
.map_err(|e| tonic::Status::internal(format!("embed task: {e}")))?
.map_err(|e| tonic::Status::internal(format!("embed: {e:#}")))?;
let ms = t0.elapsed().as_millis() as u64;
let dim = self.profile.dim as usize;
if rows.len() != n || rows.iter().any(|r| r.len() != dim) {
return Err(tonic::Status::internal(format!("the embedder returned {} row(s) for {n} text(s)", rows.len())));
}
let s = &self.served;
let calls = s.calls.fetch_add(1, Ordering::Relaxed) + 1;
let texts = s.texts.fetch_add(n as u64, Ordering::Relaxed) + n as u64;
let all_bytes = s.bytes.fetch_add(bytes as u64, Ordering::Relaxed) + bytes as u64;
let all_ms = s.embed_ms.fetch_add(ms, Ordering::Relaxed) + ms;
eprintln!(
"nornir-embed-worker: {n} text(s) {} kB in {ms} ms ({:.0} text/s) · served {calls} call(s), {texts} text(s), {} MB, {:.1} s of forward passes",
bytes / 1024,
if ms > 0 { n as f64 * 1000.0 / ms as f64 } else { 0.0 },
all_bytes >> 20,
all_ms as f64 / 1000.0
);
Ok(tonic::Response::new(pb::EmbedReply { dim: dim as u32, count: n as u32, values: rows_to_bytes(&rows), embed_ms: ms }))
}
}
/// The first line of an error, for a log line.
fn first_line(e: &anyhow::Error) -> String {
let s = format!("{e:#}");
s.lines().next().unwrap_or("").chars().take(200).collect()
}
/// Serve `worker` on `listener` until the future is dropped; `token` empty =
/// open.
pub async fn serve(worker: Worker, listener: tokio::net::TcpListener, token: &str) -> Result<()> {
let svc = pb::embedder_server::EmbedderServer::new(worker).max_decoding_message_size(MESSAGE_LIMIT).max_encoding_message_size(MESSAGE_LIMIT);
let svc = tonic::service::interceptor::InterceptedService::new(svc, nornir_service::auth_interceptor(Arc::new(token.to_string())));
let incoming = tonic::transport::server::TcpIncoming::from(listener);
tonic::transport::Server::builder().add_service(svc).serve_with_incoming(incoming).await.context("embed worker: serve")
}
/// Load the embedder a worker serves: the GPU strictly (a worker that
/// silently runs on the CPU would hand a server hours it did not ask for),
/// or — `allow_cpu` — whatever the selector picks. Returns the embedder, its
/// backend label and the device.
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
pub fn load_worker_embedder(allow_cpu: bool) -> Result<(Arc<dyn Embedder>, String, String)> {
use super::select::{Backend, chosen_backend};
let chosen = chosen_backend();
if chosen == Backend::OrtCuda {
let e = super::embed_ort::OrtEmbedder::load_gpu_strict().context("load the model on the CUDA execution provider")?;
// (the CPU fallback for over-long texts: `cpu_fallback`)
ensure!(e.cuda_ready(), "the CUDA runtime set (cudart + cuDNN) was not found — the CUDA EP would run on the CPU");
return Ok((Arc::new(e), chosen.label().to_string(), gpu_name()));
}
if !allow_cpu {
bail!(
"no GPU backend here (the selector picked {}): a worker serves a GPU. `--allow-cpu` serves {} anyway",
chosen.label(),
chosen.label()
);
}
Ok((Arc::from(super::select::load_backend(chosen)?), chosen.label().to_string(), "cpu".to_string()))
}
/// The same weights on ort's CPU execution provider (same model space): what
/// a worker embeds a text with that its device cannot hold alone.
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
pub fn cpu_fallback() -> Result<Box<dyn Embedder>> {
Ok(Box::new(super::embed_ort::OrtEmbedder::load_cpu()?))
}
/// The NVIDIA device's name, as the driver reports it (best effort).
#[cfg(all(feature = "embed-tract", feature = "embed-ort"))]
fn gpu_name() -> String {
std::process::Command::new("nvidia-smi")
.args(["--query-gpu=name,memory.total", "--format=csv,noheader"])
.output()
.ok()
.filter(|o| o.status.success())
.map(|o| String::from_utf8_lossy(&o.stdout).lines().next().unwrap_or("").trim().to_string())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "NVIDIA GPU".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
/// Vector = `[bytes, lines, 1]` normalized: tells two texts apart and
/// is the same on both ends.
struct Shape(&'static str);
impl Embedder for Shape {
fn profile(&self) -> ModelProfile {
ModelProfile { model_name: self.0.into(), weights_sha: "w".into(), tokenizer_sha: "t".into(), pooling: "mean".into(), normalize: true, dim: 3, dtype: "f32".into() }
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
Ok(texts
.iter()
.map(|t| {
let v = [t.len() as f32, t.lines().count() as f32, 1.0];
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
v.iter().map(|x| x / n).collect()
})
.collect())
}
}
/// A worker on a loopback port with `token`; returns its URL. The
/// server runs on a runtime of its own thread for the test's life.
fn start(model: &'static str, token: &'static str) -> String {
let (tx, rx) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(2).enable_all().build().unwrap();
rt.block_on(async move {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
tx.send(listener.local_addr().unwrap()).unwrap();
serve(Worker::new(Arc::new(Shape(model)), "shape", "test", "none"), listener, token).await.unwrap();
});
});
format!("http://{}", rx.recv().unwrap())
}
/// The vectors that come back over the wire are the worker's own, in
/// request order — also when the texts travel as several calls.
#[test]
fn a_remote_embed_hands_back_the_workers_vectors_in_order() {
let url = start("shape", "s3cret-token-for-test");
let remote = RemoteEmbedder::connect(&url, "s3cret-token-for-test", &Shape("shape").profile()).unwrap();
let owned: Vec<String> = (0..TEXTS_PER_CALL * 2 + 7).map(|i| "x\n".repeat(i % 13 + 1)).collect();
let texts: Vec<&str> = owned.iter().map(String::as_str).collect();
assert_eq!(calls(&texts).len(), 3, "travels as three calls");
let got = remote.embed(&texts).unwrap();
assert_eq!(got, Shape("shape").embed(&texts).unwrap());
assert_eq!(remote.profile(), Shape("shape").profile());
assert_eq!(remote.worker().backend, "test");
}
/// A worker in another model space is refused at connect: its vectors
/// would land among ours under our profile.
#[test]
fn a_worker_in_another_model_space_is_refused() {
let url = start("other-model", "");
let err = RemoteEmbedder::connect(&url, "", &Shape("shape").profile()).err().expect("must refuse");
assert!(format!("{err:#}").contains("another model space"), "{err:#}");
}
/// Without the worker's bearer nothing is embedded.
#[test]
fn a_worker_with_a_token_refuses_a_caller_without_it() {
let url = start("shape", "s3cret-token-for-test");
let err = RemoteEmbedder::connect(&url, "wrong-token-xxxxxxxx", &Shape("shape").profile()).err().expect("must refuse");
assert!(format!("{err:#}").contains("bearer"), "{err:#}");
}
/// The client works from inside a tokio runtime too (the server's own
/// threads): the block_on moves to a thread of its own.
#[test]
fn a_remote_embed_works_from_inside_a_runtime() {
let url = start("shape", "");
let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(1).enable_all().build().unwrap();
let got = rt.block_on(async move {
let remote = RemoteEmbedder::connect(&url, "", &Shape("shape").profile()).unwrap();
remote.embed(&["fn a() {}"]).unwrap()
});
assert_eq!(got, Shape("shape").embed(&["fn a() {}"]).unwrap());
}
/// Fails any batch holding a text over 100 bytes (a device too small).
struct Small;
impl Embedder for Small {
fn profile(&self) -> ModelProfile {
Shape("shape").profile()
}
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
ensure!(texts.iter().all(|t| t.len() <= 100), "out of memory");
Shape("shape").embed(texts)
}
}
fn shape_box() -> Result<Box<dyn Embedder>> {
Ok(Box::new(Shape("shape")))
}
/// A text the device cannot hold fails neither the batch it came in nor
/// itself: the others go to the device one by one, it goes to the CPU —
/// the same vectors, in order.
#[test]
fn a_text_too_big_for_the_device_is_embedded_by_the_fallback() {
let (tx, rx) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_multi_thread().worker_threads(2).enable_all().build().unwrap();
rt.block_on(async move {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
tx.send(listener.local_addr().unwrap()).unwrap();
serve(Worker::new(Arc::new(Small), "shape", "test", "none").with_fallback(shape_box), listener, "").await.unwrap();
});
});
let url = format!("http://{}", rx.recv().unwrap());
let remote = RemoteEmbedder::connect(&url, "", &Shape("shape").profile()).unwrap();
let long = "z".repeat(500);
let texts = ["fn a() {}", long.as_str(), "fn b() {}"];
assert_eq!(remote.embed(&texts).unwrap(), Shape("shape").embed(&texts).unwrap());
}
#[test]
fn calls_respect_both_caps_and_lose_nothing() {
let big = "y".repeat(BYTES_PER_CALL / 2 + 1);
let texts = vec!["a", big.as_str(), big.as_str(), "b"];
let cut = calls(&texts);
assert_eq!(cut, vec![0..2, 2..4]);
assert!(calls(&[]).is_empty());
}
}