Phases 16.2 + L2 + 17 VRAM gate + MySQL + 18 Lance hybrid milestone

Five threads of work landing as one milestone — all individually
verified end-to-end against real data, full release build clean,
46 unit tests pass.

## Phase 16.2 / 16.5 — autotune agent + ingest triggers

`vectord::agent` is a long-running tokio task that watches the trial
journal and autonomously proposes + runs new HNSW configs. Distinct
from `autotune::run_autotune` (synchronous one-shot grid). Triggered
on POST /vectors/agent/enqueue/{idx} or by the periodic wake; ingest
paths now push DatasetAppended events when an index's source dataset
gets re-ingested. Rate-limited (max_trials_per_hour) and cooldown-
gated so it can't saturate Ollama under live load.

The proposer is ε-greedy around the current champion: with prob 0.25
sample random from full bounds, otherwise perturb champion ± small
delta on both axes. Dedup against history. Deterministic — RNG seeded
from history.len() so the same journal state proposes the same next
config (helps offline replay debugging).

`[agent]` config section in lakehouse.toml; opt-in via enabled=true.

## Federation Layer 2 — runtime bucket lifecycle + per-index scoping

`BucketRegistry.buckets` moved to `std::sync::RwLock<HashMap>` so
buckets can be added/removed after startup. POST /storage/buckets
provisions at runtime; DELETE /storage/buckets/{name} unregisters
(refuses primary/rescue with 403). Local-backend buckets get their
root directory auto-created.

`IndexMeta.bucket` (default "primary" via serde) records each index's
home bucket. `TrialJournal` and `PromotionRegistry` now hold
Arc<BucketRegistry> + IndexRegistry; they resolve target store per-
index via IndexMeta.bucket. PromotionRegistry::list_all scans every
bucket and dedups by index_name. Pre-federation indexes keep working
unchanged — they just default to primary.

`ModelProfile.bucket: Option<String>` declares per-profile artifact
home. POST /vectors/profile/{id}/activate auto-provisions the
profile's bucket under storage.profile_root if not yet registered.

EvalSets stay primary-only for now — noted gap, low-risk to extend
later with the same resolver pattern.

## Phase 17 — VRAM-aware two-profile gate

Sidecar gains POST /admin/unload (Ollama keep_alive=0 trick — forces
immediate VRAM release), POST /admin/preload (keep_alive=5m with
empty prompt, takes the slot warm), and GET /admin/vram (combines
nvidia-smi snapshot with Ollama /api/ps). Exposed via aibridge as
unload_model / preload_model / vram_snapshot.

`VectorState.active_profile` is the GPU-slot singleton —
Arc<RwLock<Option<ActiveProfileSlot>>>. activate_profile checks for
a previous profile with a different ollama_name and unloads it
before preloading the new one; same-model reactivations skip the
unload (Ollama no-ops). New routes: POST /vectors/profile/{id}/
deactivate (unload + clear slot), GET /vectors/profile/active.

Verified live: staffing-recruiter (qwen2.5) → docs-assistant
(mistral) swap freed qwen2.5 from VRAM and loaded mistral. nomic-
embed-text persists across swaps because both profiles use it —
free optimization that fell out of the design. Scoped search
correctly 403s cross-profile in both directions.

## MySQL streaming connector

`crates/ingestd/src/my_stream.rs` mirrors pg_stream.rs for MySQL.
Pure-rust `mysql_async` driver (default-features=false to avoid C
deps). Same OFFSET pagination, same Parquet-streaming write shape.
Type mapping per ADR-010: int/bigint → Int32/Int64, decimal/float
→ Float64, tinyint(1)/bool → Boolean, everything else → Utf8 with
fallback parsers for date/time/json/uuid via Display.

POST /ingest/mysql parallel to /ingest/db. Same PII auto-detection,
same lineage capture (source_system="mysql"), same agent-trigger
hook. `redact_dsn` generalized — was hardcoded to "postgresql://"
length, now works for any scheme://user:pass@host/path URL (latent
PII leak fix for MySQL DSNs).

Verified live against MariaDB on localhost: 10 rows × 9 columns of
test data round-tripped through datatypes int/varchar/decimal/
tinyint/datetime/text. PII detection auto-flagged name + email.
Aggregation queries through DataFusion match the source values
exactly.

## Phase 18 — Hybrid Parquet+HNSW ⊕ Lance backend (ADR-019)

`vectord-lance` is a new firewall crate. Lance pulls Arrow 57 and
DataFusion 52 — incompatible with the rest of the workspace's
Arrow 55 / DataFusion 47. The firewall isolates that dep tree:
public API uses only std types (Vec<f32>, Vec<String>, Hit, Row,
*Stats), so no Arrow types cross the crate boundary and nothing
propagates to vectord. The ADR-019 path that didn't ship until now.

`vectord::lance_backend::LanceRegistry` lazy-creates a
LanceVectorStore per index, resolving bucket → URI via the
conventional local-bucket layout. `IndexMeta.vector_backend` and
`ModelProfile.vector_backend` carry the choice (default Parquet so
existing indexes unchanged).

Six routes under /vectors/lance/*:
- migrate/{idx}: convert binary-blob Parquet → Lance FixedSizeList
- index/{idx}: build IVF_PQ
- search/{idx}: vector search (embed via sidecar)
- doc/{idx}/{doc_id}: random row fetch
- append/{idx}: native fragment append
- stats/{idx}: row count + index presence

Verified live on the real resumes_100k_v2 corpus (100K × 768d):
- Migrate: 0.57s
- Build IVF_PQ index: 16.2s (matches ADR-019 bench; 14× faster than
  HNSW's 230s for the same data)
- Search end-to-end (Ollama embed + Lance scan): 23-53ms
- Random doc_id fetch: 5-7ms (filter scan; faster than Parquet's
  ~35ms full-file scan, slower than the bench's 311us positional
  take — would close that gap with a scalar btree on doc_id)
- Append 100 rows: 3.3ms / +320KB on disk vs Parquet's required
  full ~330MB rewrite — the structural win
- Index survives append; both backends coexist cleanly

## Known follow-ups not in this milestone

- ModelProfile.vector_backend doesn't yet auto-route /vectors/profile/
  {id}/search to Lance; callers go through /vectors/lance/* directly
- Scalar btree on doc_id (closes the 5-7ms → ~300us gap)
- vectord-lance built default-features=false → no S3 yet
- IVF_PQ recall not measured (ADR-019 caveat) — needs a Lance-aware
  variant of the eval harness
- Watcher-path ingest doesn't push agent triggers (HTTP paths do)
- EvalSets still primary-only (federation gap)
- No PATCH endpoint to move an existing index between buckets
- The pre-existing storaged::append_log doctest fails to compile
  (malformed `{prefix}/` parses as code fence) — pre-existing bug,
  left for a focused fix

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
2026-04-16 20:24:46 -05:00
co-authored by Claude Opus 4.6
parent 4e1c400f5d
commit 0d037cfac1
28 changed files with 3240 additions and 111 deletions
+43
View File
@@ -134,4 +134,47 @@ impl AiClient {
}
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
}
/// Force Ollama to unload the named model from VRAM (keep_alive=0).
/// Used for predictable profile swaps — without this, Ollama holds a
/// model for its configured TTL (default 5min) and the previous
/// profile's model can linger in VRAM next to the new one.
pub async fn unload_model(&self, model: &str) -> Result<serde_json::Value, String> {
let resp = self.client
.post(format!("{}/admin/unload", self.base_url))
.json(&serde_json::json!({ "model": model }))
.send().await
.map_err(|e| format!("unload request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("unload error: {text}"));
}
resp.json().await.map_err(|e| format!("unload parse error: {e}"))
}
/// Ask Ollama to load the named model into VRAM proactively. Makes
/// the first real request after profile activation fast (no cold-load
/// latency).
pub async fn preload_model(&self, model: &str) -> Result<serde_json::Value, String> {
let resp = self.client
.post(format!("{}/admin/preload", self.base_url))
.json(&serde_json::json!({ "model": model }))
.send().await
.map_err(|e| format!("preload request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("preload error: {text}"));
}
resp.json().await.map_err(|e| format!("preload parse error: {e}"))
}
/// GPU + loaded-model snapshot from the sidecar. Combines nvidia-smi
/// output (if available) with Ollama's /api/ps.
pub async fn vram_snapshot(&self) -> Result<serde_json::Value, String> {
let resp = self.client
.get(format!("{}/admin/vram", self.base_url))
.send().await
.map_err(|e| format!("vram request failed: {e}"))?;
resp.json().await.map_err(|e| format!("vram parse error: {e}"))
}
}
+8
View File
@@ -16,9 +16,17 @@ pub fn router(client: AiClient) -> Router {
.route("/embed", post(embed))
.route("/generate", post(generate))
.route("/rerank", post(rerank))
.route("/vram", get(vram))
.with_state(client)
}
async fn vram(State(client): State<AiClient>) -> impl IntoResponse {
match client.vram_snapshot().await {
Ok(snap) => Ok(Json(snap)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
match client.health().await {
Ok(info) => Ok(Json(info)),
+10
View File
@@ -357,6 +357,14 @@ struct CreateProfileRequest {
embed_model: String,
#[serde(default)]
created_by: String,
/// Federation: optional per-profile bucket (`profile:{id}` by convention).
/// Omitting keeps artifacts in primary.
#[serde(default)]
bucket: Option<String>,
/// ADR-019 hybrid: which vector backend to route this profile's
/// indexes to. Defaults to Parquet+HNSW.
#[serde(default)]
vector_backend: shared::types::VectorBackend,
}
fn default_embed_model_req() -> String { "nomic-embed-text".to_string() }
@@ -374,6 +382,8 @@ async fn create_profile(
embed_model: req.embed_model,
created_at: chrono::Utc::now(),
created_by: req.created_by,
bucket: req.bucket,
vector_backend: req.vector_backend,
};
match registry.put_profile(profile).await {
Ok(p) => Ok((StatusCode::CREATED, Json(p))),
+53 -14
View File
@@ -75,6 +75,41 @@ async fn main() {
// AI sidecar client
let ai_client = aibridge::client::AiClient::new(&config.sidecar.url);
// Vector service components — built before the router because both the
// /vectors service AND ingestd need the agent handle to enqueue triggers.
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
let _ = index_reg.rebuild().await;
let hnsw = vectord::hnsw::HnswStore::new();
let emb_cache = vectord::embedding_cache::EmbeddingCache::new(store.clone());
// Phase B (federation layer 2): trial journals + promotion files now
// live in each index's recorded bucket (IndexMeta.bucket), not always
// primary. Journal / registry resolve per-call via the bucket registry.
let tj = vectord::trial::TrialJournal::new(bucket_registry.clone(), index_reg.clone());
let pr = vectord::promotion::PromotionRegistry::new(bucket_registry.clone(), index_reg.clone());
// Phase 16.2: spawn the autotune agent. When config.agent.enabled=false
// this returns a handle that drops triggers silently — no surprise load.
let agent_cfg = vectord::agent::AgentConfig {
enabled: config.agent.enabled,
cycle_interval_secs: config.agent.cycle_interval_secs,
cooldown_between_trials_secs: config.agent.cooldown_between_trials_secs,
min_recall: config.agent.min_recall,
max_trials_per_hour: config.agent.max_trials_per_hour,
};
let agent_handle = vectord::agent::spawn(
agent_cfg,
vectord::agent::AgentDeps {
store: store.clone(),
ai_client: ai_client.clone(),
catalog: registry.clone(),
index_registry: index_reg.clone(),
hnsw_store: hnsw.clone(),
embedding_cache: emb_cache.clone(),
trial_journal: tj.clone(),
promotion_registry: pr.clone(),
},
);
// HTTP router
let mut app = Router::new()
.route("/health", get(health))
@@ -87,21 +122,25 @@ async fn main() {
store: store.clone(),
registry: registry.clone(),
buckets: bucket_registry.clone(),
agent_handle: agent_handle.clone(),
index_registry: index_reg.clone(),
}))
.nest("/vectors", vectord::service::router({
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
let _ = index_reg.rebuild().await;
vectord::service::VectorState {
store: store.clone(),
ai_client: ai_client.clone(),
job_tracker: vectord::jobs::JobTracker::new(),
index_registry: index_reg,
hnsw_store: vectord::hnsw::HnswStore::new(),
embedding_cache: vectord::embedding_cache::EmbeddingCache::new(store.clone()),
trial_journal: vectord::trial::TrialJournal::new(store.clone()),
catalog: registry.clone(),
promotion_registry: vectord::promotion::PromotionRegistry::new(store.clone()),
}
.nest("/vectors", vectord::service::router(vectord::service::VectorState {
store: store.clone(),
ai_client: ai_client.clone(),
job_tracker: vectord::jobs::JobTracker::new(),
index_registry: index_reg.clone(),
hnsw_store: hnsw,
embedding_cache: emb_cache,
trial_journal: tj,
catalog: registry.clone(),
promotion_registry: pr,
agent_handle,
bucket_registry: bucket_registry.clone(),
active_profile: std::sync::Arc::new(tokio::sync::RwLock::new(None)),
lance: vectord::lance_backend::LanceRegistry::new(
bucket_registry.clone(), index_reg.clone(),
),
}))
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
.nest("/journal", journald::service::router(journal))
+2
View File
@@ -7,6 +7,7 @@ edition = "2024"
shared = { path = "../shared" }
storaged = { path = "../storaged" }
catalogd = { path = "../catalogd" }
vectord = { path = "../vectord" }
tokio = { workspace = true }
axum = { workspace = true, features = ["multipart"] }
lopdf = { workspace = true }
@@ -21,4 +22,5 @@ csv = { workspace = true }
chrono = { workspace = true }
object_store = { workspace = true }
tokio-postgres = { workspace = true }
mysql_async = { workspace = true }
uuid = { workspace = true }
+1
View File
@@ -1,4 +1,5 @@
pub mod db_ingest;
pub mod my_stream;
pub mod pg_stream;
pub mod detect;
pub mod csv_ingest;
+410
View File
@@ -0,0 +1,410 @@
/// Streaming MySQL ingest.
///
/// Mirrors `pg_stream` for MySQL sources. Same OFFSET-paginated strategy,
/// same Parquet-streaming shape. Uses `mysql_async` (pure-rust) so we
/// don't need a C client library at build time.
///
/// Type mapping follows ADR-010 (default to string on ambiguity):
/// - Booleans (TINYINT(1)) and integer types map to Arrow Int32/Int64
/// - Floating point and decimals → Float64
/// - Everything else (text, varchar, json, date, timestamp) → Utf8
///
/// What's deliberately not supported (yet):
/// - TLS connections — `minimal` feature is plain TCP only. Upgrade
/// when a tenant actually needs it.
/// - Keyset pagination — OFFSET scans are O(N²) at multi-million-row
/// scale; same caveat as `pg_stream`.
/// - BINARY/BLOB columns — currently rendered as base64 or empty string
/// via Display-fallback.
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use mysql_async::prelude::*;
use mysql_async::{Conn, Opts, Row, Value};
use parquet::arrow::ArrowWriter;
use std::sync::Arc;
/// Request shape for MySQL streaming ingest.
#[derive(Debug, Clone, serde::Deserialize)]
pub struct MyStreamRequest {
/// mysql://user:pass@host:port/db
pub dsn: String,
pub table: String,
#[serde(default)]
pub dataset_name: Option<String>,
/// Rows per fetch. Default 10_000.
#[serde(default)]
pub batch_size: Option<usize>,
/// Column to ORDER BY for stable pagination. If omitted, the first
/// column returned by the schema probe is used.
#[serde(default)]
pub order_by: Option<String>,
/// Hard cap on total rows (for sampling / previews).
#[serde(default)]
pub limit: Option<usize>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct MyStreamResult {
pub table: String,
pub rows: usize,
pub batches: usize,
pub columns: usize,
pub schema: Vec<String>,
pub parquet_bytes: u64,
pub duration_secs: f32,
}
/// Parsed DSN pieces. Kept local (rather than reusing pg's DbConfig) so
/// the MySQL connector doesn't depend on the PG path.
#[derive(Debug, Clone)]
pub struct MyConfig {
pub host: String,
pub port: u16,
pub user: String,
pub password: String,
pub database: String,
}
impl MyConfig {
/// Build a mysql_async `Opts` from the parsed config.
pub fn to_opts(&self) -> Opts {
let url = if self.password.is_empty() {
format!(
"mysql://{}@{}:{}/{}",
percent(&self.user), percent(&self.host), self.port, percent(&self.database),
)
} else {
format!(
"mysql://{}:{}@{}:{}/{}",
percent(&self.user), percent(&self.password),
percent(&self.host), self.port, percent(&self.database),
)
};
Opts::from_url(&url).expect("MyConfig.to_opts: rebuilt URL should parse")
}
}
/// Minimal URL-encoder for the few characters that commonly appear in
/// MySQL passwords. mysql_async's URL parser expects valid URL-encoded
/// components.
fn percent(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match c {
'@' | ':' | '/' | '?' | '#' | '[' | ']' | ' ' | '%' => {
out.push_str(&format!("%{:02X}", c as u32));
}
_ => out.push(c),
}
}
out
}
/// Parse a mysql:// DSN.
/// Supports: mysql://[user[:password]@]host[:port][/db]
pub fn parse_dsn(dsn: &str) -> Result<MyConfig, String> {
let rest = dsn
.strip_prefix("mysql://")
.ok_or_else(|| "DSN must start with mysql://".to_string())?;
let (auth_host, database) = match rest.split_once('/') {
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
None => (rest, String::new()),
};
let (userpass, hostport) = match auth_host.rsplit_once('@') {
Some((up, hp)) => (Some(up), hp),
None => (None, auth_host),
};
let (user, password) = match userpass {
Some(up) => match up.split_once(':') {
Some((u, p)) => (u.to_string(), p.to_string()),
None => (up.to_string(), String::new()),
},
None => ("root".to_string(), String::new()),
};
let (host, port) = match hostport.rsplit_once(':') {
Some((h, p)) => (
h.to_string(),
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
),
None => (hostport.to_string(), 3306),
};
if host.is_empty() {
return Err("DSN has no host".into());
}
if database.is_empty() {
return Err("DSN has no database path (mysql://... /db)".into());
}
Ok(MyConfig { host, port, user, password, database })
}
/// Stream a MySQL table as Parquet bytes.
pub async fn stream_table_to_parquet(
req: &MyStreamRequest,
) -> Result<(bytes::Bytes, MyStreamResult), String> {
let t0 = std::time::Instant::now();
let config = parse_dsn(&req.dsn)?;
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
let mut conn = Conn::new(config.to_opts()).await
.map_err(|e| format!("mysql connect: {e}"))?;
// Probe columns via information_schema — gives real type names that
// we can map to Arrow dtypes without needing to fetch a row first.
let columns = probe_columns(&mut conn, &config.database, &req.table).await?;
if columns.is_empty() {
return Err(format!("table '{}' not found or has no columns", req.table));
}
let arrow_fields: Vec<Field> = columns
.iter()
.map(|(name, ty)| Field::new(name, mysql_type_to_arrow(ty), true))
.collect();
let schema = Arc::new(Schema::new(arrow_fields));
let schema_report: Vec<String> = columns
.iter()
.map(|(n, t)| format!("{}:{}", n, t))
.collect();
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
.map_err(|e| format!("ArrowWriter init: {e}"))?;
let mut total_rows: usize = 0;
let mut batch_count: usize = 0;
let row_cap = req.limit.unwrap_or(usize::MAX);
loop {
let remaining = row_cap.saturating_sub(total_rows);
if remaining == 0 { break; }
let fetch = remaining.min(batch_size);
// Backticks are MySQL's identifier quote. Forbid backticks in
// table/column names to prevent injection — neither pg nor us
// should be accepting those anyway.
if req.table.contains('`') || order_col.contains('`') {
return Err("table or order_by column contains backticks — refused".into());
}
let sql = format!(
"SELECT * FROM `{}` ORDER BY `{}` LIMIT {} OFFSET {}",
req.table, order_col, fetch, total_rows,
);
let rows: Vec<Row> = conn.query(&sql).await
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
if rows.is_empty() { break; }
let n = rows.len();
let arrays: Vec<ArrayRef> = columns
.iter()
.enumerate()
.map(|(idx, (_, ty))| rows_to_column(&rows, idx, ty))
.collect::<Result<_, _>>()?;
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch: {e}"))?;
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
total_rows += n;
batch_count += 1;
tracing::info!(
"mysql stream '{}': fetched batch {} ({} rows, total {})",
req.table, batch_count, n, total_rows,
);
if n < fetch { break; }
}
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
conn.disconnect().await.ok();
let result = MyStreamResult {
table: req.table.clone(),
rows: total_rows,
batches: batch_count,
columns: columns.len(),
schema: schema_report,
parquet_bytes: buf.len() as u64,
duration_secs: t0.elapsed().as_secs_f32(),
};
Ok((bytes::Bytes::from(buf), result))
}
async fn probe_columns(
conn: &mut Conn,
schema: &str,
table: &str,
) -> Result<Vec<(String, String)>, String> {
let sql = format!(
"SELECT COLUMN_NAME, DATA_TYPE FROM information_schema.columns \
WHERE table_schema = '{}' AND table_name = '{}' \
ORDER BY ORDINAL_POSITION",
schema.replace('\'', "''"),
table.replace('\'', "''"),
);
let rows: Vec<(String, String)> = conn.query(&sql).await
.map_err(|e| format!("probe columns: {e}"))?;
Ok(rows)
}
/// MySQL data type string → Arrow DataType. Conservative: anything we
/// don't recognize becomes Utf8 (ADR-010).
fn mysql_type_to_arrow(ty: &str) -> DataType {
match ty.to_lowercase().as_str() {
"tinyint" | "smallint" | "mediumint" | "int" | "integer" => DataType::Int32,
"bigint" => DataType::Int64,
"float" | "double" | "decimal" | "numeric" | "real" => DataType::Float64,
"bit" | "bool" | "boolean" => DataType::Boolean,
_ => DataType::Utf8,
}
}
/// Convert a single column slice of MySQL rows into an Arrow array.
fn rows_to_column(
rows: &[Row],
idx: usize,
ty: &str,
) -> Result<ArrayRef, String> {
let arrow_ty = mysql_type_to_arrow(ty);
match arrow_ty {
DataType::Boolean => {
let v: Vec<Option<bool>> = rows.iter().map(|r| cell_as_bool(r, idx)).collect();
Ok(Arc::new(BooleanArray::from(v)))
}
DataType::Int32 => {
let v: Vec<Option<i32>> = rows.iter().map(|r| cell_as_i64(r, idx).map(|n| n as i32)).collect();
Ok(Arc::new(Int32Array::from(v)))
}
DataType::Int64 => {
let v: Vec<Option<i64>> = rows.iter().map(|r| cell_as_i64(r, idx)).collect();
Ok(Arc::new(Int64Array::from(v)))
}
DataType::Float64 => {
let v: Vec<Option<f64>> = rows.iter().map(|r| cell_as_f64(r, idx)).collect();
Ok(Arc::new(Float64Array::from(v)))
}
_ => {
let v: Vec<Option<String>> = rows.iter().map(|r| cell_as_string(r, idx)).collect();
Ok(Arc::new(StringArray::from(v)))
}
}
}
fn cell(r: &Row, idx: usize) -> &Value { r.as_ref(idx).unwrap_or(&Value::NULL) }
fn cell_as_bool(r: &Row, idx: usize) -> Option<bool> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n != 0),
Value::UInt(n) => Some(*n != 0),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| {
match s.to_ascii_lowercase().as_str() {
"true" | "1" | "y" | "yes" => Some(true),
"false" | "0" | "n" | "no" => Some(false),
_ => None,
}
}),
_ => None,
}
}
fn cell_as_i64(r: &Row, idx: usize) -> Option<i64> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n),
Value::UInt(n) => i64::try_from(*n).ok(),
Value::Float(f) => Some(*f as i64),
Value::Double(f) => Some(*f as i64),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
_ => None,
}
}
fn cell_as_f64(r: &Row, idx: usize) -> Option<f64> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n as f64),
Value::UInt(n) => Some(*n as f64),
Value::Float(f) => Some(*f as f64),
Value::Double(f) => Some(*f),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
_ => None,
}
}
fn cell_as_string(r: &Row, idx: usize) -> Option<String> {
match cell(r, idx) {
Value::NULL => None,
Value::Bytes(b) => Some(String::from_utf8_lossy(b).into_owned()),
Value::Int(n) => Some(n.to_string()),
Value::UInt(n) => Some(n.to_string()),
Value::Float(f) => Some(f.to_string()),
Value::Double(f) => Some(f.to_string()),
Value::Date(y, mo, d, h, mi, s, _us) => {
Some(format!("{:04}-{:02}-{:02} {:02}:{:02}:{:02}", y, mo, d, h, mi, s))
}
Value::Time(neg, days, h, mi, s, _us) => {
let sign = if *neg { "-" } else { "" };
Some(format!("{}{}d {:02}:{:02}:{:02}", sign, days, h, mi, s))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_dsn_full() {
let c = parse_dsn("mysql://daisy:secret@db.example.com:3307/my_db").unwrap();
assert_eq!(c.host, "db.example.com");
assert_eq!(c.port, 3307);
assert_eq!(c.user, "daisy");
assert_eq!(c.password, "secret");
assert_eq!(c.database, "my_db");
}
#[test]
fn parse_dsn_default_port() {
let c = parse_dsn("mysql://root@localhost/shop").unwrap();
assert_eq!(c.port, 3306);
assert_eq!(c.user, "root");
assert_eq!(c.password, "");
}
#[test]
fn parse_dsn_no_auth() {
let c = parse_dsn("mysql://127.0.0.1:3306/analytics").unwrap();
assert_eq!(c.user, "root");
assert_eq!(c.host, "127.0.0.1");
}
#[test]
fn parse_dsn_rejects_non_mysql() {
assert!(parse_dsn("postgresql://host/db").is_err());
}
#[test]
fn parse_dsn_requires_db() {
assert!(parse_dsn("mysql://localhost:3306").is_err());
}
#[test]
fn type_map_basics() {
assert_eq!(mysql_type_to_arrow("int"), DataType::Int32);
assert_eq!(mysql_type_to_arrow("BIGINT"), DataType::Int64);
assert_eq!(mysql_type_to_arrow("decimal"), DataType::Float64);
assert_eq!(mysql_type_to_arrow("varchar"), DataType::Utf8);
assert_eq!(mysql_type_to_arrow("json"), DataType::Utf8);
assert_eq!(mysql_type_to_arrow("bool"), DataType::Boolean);
}
}
+152 -13
View File
@@ -11,7 +11,7 @@ use serde::Deserialize;
use std::sync::Arc;
use catalogd::registry::Registry;
use crate::{db_ingest, pg_stream, pipeline};
use crate::{db_ingest, my_stream, pg_stream, pipeline};
use shared::arrow_helpers::record_batch_to_parquet;
use shared::types::{ObjectRef, SchemaFingerprint};
use storaged::ops;
@@ -23,6 +23,36 @@ pub struct IngestState {
pub registry: Registry,
/// Federation layer 2: lookup target bucket from request headers.
pub buckets: Arc<BucketRegistry>,
/// Phase 16.5: when ingest marks a dataset's embeddings stale, push
/// a `DatasetAppended` trigger so the autotune agent can schedule
/// re-trials. Holds the agent handle (no-op when agent disabled).
pub agent_handle: vectord::agent::AgentHandle,
/// Used to look up which HNSW indexes are attached to the
/// just-ingested dataset. Each matching index gets one trigger.
pub index_registry: vectord::index_registry::IndexRegistry,
}
/// Push `DatasetAppended` triggers for every HNSW index bound to this
/// dataset. Called after a successful ingest. Logs failures but never
/// fails the ingest — the agent is best-effort, not load-bearing.
async fn notify_agent_on_append(state: &IngestState, dataset_name: &str) {
let indexes = state.index_registry.list(Some(dataset_name), None).await;
if indexes.is_empty() {
return;
}
for meta in indexes {
let event = vectord::agent::TriggerEvent::dataset_appended(
meta.index_name.clone(), dataset_name,
);
if let Err(e) = state.agent_handle.enqueue(event).await {
tracing::warn!("agent enqueue failed for '{}': {}", meta.index_name, e);
} else {
tracing::info!(
"agent: enqueued DatasetAppended({}) for index '{}'",
dataset_name, meta.index_name,
);
}
}
}
/// Resolve the target bucket from `X-Lakehouse-Bucket` header.
@@ -55,6 +85,7 @@ pub fn router(state: IngestState) -> Router {
.route("/postgres/tables", post(list_pg_tables))
.route("/postgres/import", post(import_pg_table))
.route("/db", post(ingest_db_stream))
.route("/mysql", post(ingest_mysql_stream))
.with_state(state)
}
@@ -94,6 +125,9 @@ async fn ingest_file(
match pipeline::ingest_file_to_bucket(&filename, &content, dataset_name, &bucket, &store, &state.registry).await {
Ok(result) => {
if !result.deduplicated {
notify_agent_on_append(&state, &result.dataset_name).await;
}
if result.deduplicated {
Ok((StatusCode::OK, Json(result)))
} else {
@@ -309,6 +343,9 @@ async fn ingest_db_stream(
// index. No-op for newly-created datasets.
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
// Phase 16.5: notify autotune agent. No-op if no indexes attached.
notify_agent_on_append(&state, &dataset_name).await;
Ok((StatusCode::CREATED, Json(serde_json::json!({
"dataset_name": dataset_name,
"table": stream_result.table,
@@ -322,17 +359,119 @@ async fn ingest_db_stream(
}))))
}
/// Redact the password in a postgresql:// DSN for logging / lineage.
fn redact_dsn(dsn: &str) -> String {
// Simple approach: find @...: and replace password between : and @.
if let Some(at_idx) = dsn.rfind('@') {
let prefix = &dsn[..at_idx];
if let Some(colon_idx) = prefix.rfind(':') {
// But only if the colon is after the scheme colons (postgresql://)
if colon_idx > "postgresql://".len() {
return format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..]);
}
}
/// Streaming MySQL ingest. Mirrors `ingest_db_stream` — same pagination +
/// Parquet-footer schema recovery + PII/lineage metadata. Different driver,
/// otherwise identical.
async fn ingest_mysql_stream(
State(state): State<IngestState>,
headers: HeaderMap,
Json(req): Json<my_stream::MyStreamRequest>,
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
tracing::info!(
"mysql stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
req.table, req.dataset_name, bucket, req.batch_size,
);
let (parquet, stream_result) = my_stream::stream_table_to_parquet(&req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
if stream_result.rows == 0 {
return Ok((StatusCode::OK, Json(serde_json::json!({
"table": req.table,
"rows": 0,
"message": "table is empty",
}))));
}
dsn.to_string()
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
let storage_key = format!("datasets/{}.parquet", dataset_name);
let size_bytes = parquet.len() as u64;
ops::put(&store, &storage_key, parquet)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
let now = chrono::Utc::now();
state.registry.register(
dataset_name.clone(),
SchemaFingerprint(schema_fp.0),
vec![ObjectRef {
bucket: bucket.clone(),
key: storage_key.clone(),
size_bytes,
created_at: now,
}],
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
shared::types::ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
}
}).collect();
let lineage = shared::types::Lineage {
source_system: "mysql".to_string(),
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
ingest_job: format!("mysql-stream-{}", now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
sensitivity,
columns: Some(columns),
lineage: Some(lineage),
row_count: Some(stream_result.rows as u64),
..Default::default()
}).await;
// Phase C: mark stale + Phase 16.5: notify agent.
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
notify_agent_on_append(&state, &dataset_name).await;
Ok((StatusCode::CREATED, Json(serde_json::json!({
"dataset_name": dataset_name,
"table": stream_result.table,
"rows": stream_result.rows,
"batches": stream_result.batches,
"columns": stream_result.columns,
"schema": stream_result.schema,
"storage_key": storage_key,
"size_bytes": size_bytes,
"duration_secs": stream_result.duration_secs,
}))))
}
/// Redact the password in any `scheme://user:pass@host/...` DSN. Works
/// for postgres, mysql, and any other scheme using the same shape.
fn redact_dsn(dsn: &str) -> String {
let scheme_end = match dsn.find("://") {
Some(i) => i + 3,
None => return dsn.to_string(),
};
let at_idx = match dsn.rfind('@') {
Some(i) if i > scheme_end => i,
_ => return dsn.to_string(),
};
// Find the password separator (first `:` after scheme, before `@`).
let userpass = &dsn[scheme_end..at_idx];
let colon_offset = match userpass.find(':') {
Some(i) => i,
None => return dsn.to_string(), // no password in DSN
};
let colon_idx = scheme_end + colon_offset;
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
}
+39
View File
@@ -16,8 +16,46 @@ pub struct Config {
pub auth: AuthConfig,
#[serde(default)]
pub observability: ObservabilityConfig,
#[serde(default)]
pub agent: AgentSettings,
}
/// Phase 16.2 — background autotune agent settings.
///
/// Duplicated from `vectord::agent::AgentConfig` because `shared` can't
/// depend on `vectord` (vectord already depends on shared). The gateway
/// copies these into the vectord config at startup.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentSettings {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_cycle_interval_secs")]
pub cycle_interval_secs: u64,
#[serde(default = "default_cooldown_secs")]
pub cooldown_between_trials_secs: u64,
#[serde(default = "default_min_recall")]
pub min_recall: f32,
#[serde(default = "default_max_trials_per_hour")]
pub max_trials_per_hour: u32,
}
impl Default for AgentSettings {
fn default() -> Self {
Self {
enabled: false,
cycle_interval_secs: default_cycle_interval_secs(),
cooldown_between_trials_secs: default_cooldown_secs(),
min_recall: default_min_recall(),
max_trials_per_hour: default_max_trials_per_hour(),
}
}
}
fn default_cycle_interval_secs() -> u64 { 60 }
fn default_cooldown_secs() -> u64 { 30 }
fn default_min_recall() -> f32 { 0.9 }
fn default_max_trials_per_hour() -> u32 { 30 }
#[derive(Debug, Clone, Deserialize)]
pub struct GatewayConfig {
#[serde(default = "default_host")]
@@ -159,6 +197,7 @@ impl Default for Config {
ai: AiConfig::default(),
auth: AuthConfig::default(),
observability: ObservabilityConfig::default(),
agent: AgentSettings::default(),
}
}
}
+45
View File
@@ -255,10 +255,55 @@ pub struct ModelProfile {
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
/// Federation layer 2: which bucket this profile's artifacts
/// (trial journals, promotions for its indexes) live in. If set and
/// not yet registered, `POST /vectors/profile/{id}/activate` auto-
/// provisions a local bucket under `storage.profile_root` and
/// registers it before warming indexes. `None` = share the primary
/// bucket (pre-federation behavior).
#[serde(default)]
pub bucket: Option<String>,
/// ADR-019 hybrid: which storage backend serves this profile's
/// indexes. Default Parquet so pre-existing profiles unchanged.
/// New profiles can opt into Lance for append-heavy / random-access
/// / 5M+ vector workloads.
#[serde(default)]
pub vector_backend: VectorBackend,
}
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
/// Per-profile / per-index storage backend choice.
///
/// `Parquet` (default) — vectors as binary-blob in Parquet + in-RAM HNSW
/// for ANN. Best when:
/// - Index ≤ ~5M vectors (fits in RAM)
/// - Search latency dominates (HNSW p50 < 1ms at 100K)
/// - Workload is mostly read, not append
///
/// `Lance` — vectors as native FixedSizeList in a Lance dataset, IVF_PQ
/// on-disk index. Best when:
/// - Index will exceed RAM ceiling (5M+ vectors per index)
/// - Heavy random-row access for RAG (Lance: 311us; Parquet: ~35ms full
/// scan per fetch)
/// - Append-heavy ingest (Lance: 0.08s for 10K rows; Parquet: full
/// ~330MB rewrite per change)
/// - Need to rebuild index frequently (Lance IVF_PQ: 16s; HNSW: 230s
/// at the same scale)
///
/// Defaults to Parquet for backward compat with pre-ADR-019 indexes.
/// See `docs/ADR-019-vector-storage.md` for the full scorecard.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum VectorBackend {
Parquet,
Lance,
}
impl Default for VectorBackend {
fn default() -> Self { VectorBackend::Parquet }
}
/// Soft-delete marker (Phase E).
///
/// Tombstones live beside the dataset in `_catalog/tombstones/{dataset}/`
+57 -4
View File
@@ -13,10 +13,11 @@ use axum::{
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, put},
routing::{delete, get, post, put},
};
use chrono::{DateTime, Utc};
use serde::Deserialize;
use shared::config::BucketConfig;
use std::sync::Arc;
use crate::error_journal::BucketErrorEvent;
@@ -24,16 +25,68 @@ use crate::registry::{BucketInfo, BucketRegistry};
pub fn router(registry: Arc<BucketRegistry>) -> Router {
Router::new()
.route("/buckets", get(list_buckets))
.route("/buckets", get(list_buckets).post(add_bucket))
.route("/buckets/{name}", delete(remove_bucket))
.route("/errors", get(list_errors))
.route("/errors/flush", axum::routing::post(flush_errors))
.route("/errors/compact", axum::routing::post(compact_errors))
.route("/errors/flush", post(flush_errors))
.route("/errors/compact", post(compact_errors))
.route("/bucket-health", get(get_health))
.route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
.route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
.with_state(registry)
}
/// Provision + register a bucket at runtime. Body is a `BucketConfig`.
///
/// Federation layer 2: profile buckets (`profile:alice`) and tenant
/// buckets can be added after startup without a service restart.
async fn add_bucket(
State(reg): State<Arc<BucketRegistry>>,
Json(bc): Json<BucketConfig>,
) -> impl IntoResponse {
// Safety net: local backends must land somewhere under the configured
// profile_root (for profile:*) or be otherwise explicitly rooted.
// Refuse empty / missing root to avoid "bucket created at ./" surprises.
if bc.backend == "local" {
match bc.root.as_deref() {
Some(r) if !r.trim().is_empty() => {}
_ => return Err((StatusCode::BAD_REQUEST,
"local bucket requires a non-empty 'root' path".to_string())),
}
}
match reg.add_bucket(bc).await {
Ok(info) => Ok((StatusCode::CREATED, Json(info))),
Err(e) => {
let code = if e.contains("already registered") {
StatusCode::CONFLICT
} else {
StatusCode::BAD_REQUEST
};
Err((code, e))
}
}
}
/// Unregister a bucket. Refused for primary / rescue / unknown names.
async fn remove_bucket(
State(reg): State<Arc<BucketRegistry>>,
Path(name): Path<String>,
) -> impl IntoResponse {
match reg.remove_bucket(&name) {
Ok(()) => Ok((StatusCode::OK, format!("unregistered: {name}"))),
Err(e) => {
let code = if e.contains("cannot remove") {
StatusCode::FORBIDDEN
} else if e.contains("not registered") {
StatusCode::NOT_FOUND
} else {
StatusCode::BAD_REQUEST
};
Err((code, e))
}
}
}
async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
Json(reg.list().await)
}
+111 -23
View File
@@ -15,7 +15,7 @@ use serde::Serialize;
use shared::config::{BucketConfig, StorageConfig};
use shared::secrets::{BucketCredentials, SecretsProvider};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::{Arc, RwLock};
use crate::error_journal::{BucketErrorEvent, ErrorJournal};
@@ -55,10 +55,16 @@ pub enum BucketRole {
}
pub struct BucketRegistry {
buckets: HashMap<String, Arc<BucketEntry>>,
/// RwLock because federation layer 2 adds runtime bucket lifecycle
/// (`POST /storage/buckets` / `DELETE /storage/buckets/{name}`).
/// Almost all accesses are reads — writes happen at provision /
/// deactivate time only.
buckets: RwLock<HashMap<String, Arc<BucketEntry>>>,
default: String,
rescue: Option<String>,
profile_root: String,
/// Held so runtime `add_bucket` can resolve secret_ref handles.
secrets: Arc<dyn SecretsProvider>,
journal: ErrorJournal,
}
@@ -113,10 +119,11 @@ impl BucketRegistry {
let _ = journal.load_recent().await;
Ok(Self {
buckets,
buckets: RwLock::new(buckets),
default: "primary".to_string(),
rescue: cfg.rescue_bucket.clone(),
profile_root: cfg.profile_root.clone(),
secrets,
journal,
})
}
@@ -125,12 +132,19 @@ impl BucketRegistry {
pub fn rescue_name(&self) -> Option<&str> { self.rescue.as_deref() }
pub fn profile_root(&self) -> &str { &self.profile_root }
pub fn journal(&self) -> &ErrorJournal { &self.journal }
/// Resolve a bucket name to its object store. Existing call sites use
/// this as a drop-in replacement for the old single-store pattern.
///
/// Uses `std::sync::RwLock` — every caller holds the guard for
/// microseconds (clones an Arc and drops), so there's no async-
/// blocking concern. Never hold across an await.
pub fn get(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>, String> {
self.buckets
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard
.get(bucket)
.map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))
@@ -139,22 +153,34 @@ impl BucketRegistry {
/// The default bucket's store — use for code paths that don't yet know
/// about buckets.
pub fn default_store(&self) -> Arc<dyn ObjectStore> {
self.buckets.get(&self.default).unwrap().store.clone()
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(&self.default).unwrap().store.clone()
}
/// True if this bucket name is registered.
pub fn contains(&self, bucket: &str) -> bool {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.contains_key(bucket)
}
/// List all registered buckets. Checks reachability by doing a trivial
/// `list` with limit 1 on each.
pub async fn list(&self) -> Vec<BucketInfo> {
let mut out = Vec::with_capacity(self.buckets.len());
for (name, entry) in &self.buckets {
let reachable = probe(&entry.store).await;
let role = self.classify(name);
out.push(BucketInfo {
name: name.clone(),
backend: entry.backend.clone(),
reachable,
role,
});
// Snapshot (name, backend, store_clone) under the lock, then probe
// outside — probing can be slow (network), don't hold the lock.
let snapshot: Vec<(String, String, Arc<dyn ObjectStore>)> = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard
.iter()
.map(|(n, e)| (n.clone(), e.backend.clone(), e.store.clone()))
.collect()
};
let mut out = Vec::with_capacity(snapshot.len());
for (name, backend, store) in snapshot {
let reachable = probe(&store).await;
let role = self.classify(&name);
out.push(BucketInfo { name, backend, reachable, role });
}
out.sort_by(|a, b| a.name.cmp(&b.name));
out
@@ -167,14 +193,69 @@ impl BucketRegistry {
else { BucketRole::Tenant }
}
/// Provision + register a bucket at runtime (federation layer 2).
/// Returns Err if the name is already registered. Local backends get
/// their root directory created automatically.
pub async fn add_bucket(&self, bc: BucketConfig) -> Result<BucketInfo, String> {
if self.contains(&bc.name) {
return Err(format!("bucket '{}' already registered", bc.name));
}
let store = build_store(&bc, self.secrets.as_ref()).await?;
let reachable = probe(&store).await;
let entry = Arc::new(BucketEntry {
name: bc.name.clone(),
backend: bc.backend.clone(),
store,
config: bc.clone(),
});
let role = self.classify(&bc.name);
let info = BucketInfo {
name: bc.name.clone(),
backend: bc.backend.clone(),
reachable,
role,
};
{
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
guard.insert(bc.name.clone(), entry);
}
tracing::info!("registered bucket '{}' backend={} reachable={}",
bc.name, bc.backend, reachable);
Ok(info)
}
/// Unregister a bucket. Refuses to remove `primary` or the configured
/// rescue bucket — those are load-bearing. Caller is responsible for
/// any final flush of bucket-local state before calling this.
pub fn remove_bucket(&self, name: &str) -> Result<(), String> {
if name == self.default {
return Err("cannot remove primary bucket".into());
}
if Some(name) == self.rescue.as_deref() {
return Err(format!("cannot remove rescue bucket '{name}'"));
}
let removed = {
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
guard.remove(name)
};
if removed.is_none() {
return Err(format!("bucket '{name}' not registered"));
}
tracing::info!("unregistered bucket '{}'", name);
Ok(())
}
/// Read with rescue-bucket fallback. If the target bucket fails and a
/// rescue is configured, retries against rescue. Records every failure
/// in the error journal.
pub async fn read_smart(&self, bucket: &str, key: &str) -> Result<ReadOutcome, String> {
let target = self.buckets.get(bucket)
.ok_or_else(|| format!("unknown bucket: {bucket}"))?;
let target_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(bucket).map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
};
match crate::ops::get(&target.store, key).await {
match crate::ops::get(&target_store, key).await {
Ok(data) => Ok(ReadOutcome {
data, rescued: false,
original_bucket: bucket.to_string(),
@@ -187,8 +268,12 @@ impl BucketRegistry {
// Try rescue, if any.
if let Some(rescue_name) = &self.rescue {
if rescue_name != bucket {
if let Some(rescue) = self.buckets.get(rescue_name) {
match crate::ops::get(&rescue.store, key).await {
let rescue_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(rescue_name).map(|e| e.store.clone())
};
if let Some(rescue_store) = rescue_store {
match crate::ops::get(&rescue_store, key).await {
Ok(data) => {
self.journal.mark_rescued_last(bucket, key).await;
return Ok(ReadOutcome {
@@ -219,9 +304,12 @@ impl BucketRegistry {
key: &str,
data: bytes::Bytes,
) -> Result<(), String> {
let target = self.buckets.get(bucket)
.ok_or_else(|| format!("unknown bucket: {bucket}"))?;
match crate::ops::put(&target.store, key, data).await {
let target_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(bucket).map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
};
match crate::ops::put(&target_store, key, data).await {
Ok(()) => Ok(()),
Err(err) => {
self.journal.append(BucketErrorEvent::new_write(bucket, key, &err)).await;
+31
View File
@@ -0,0 +1,31 @@
[package]
name = "vectord-lance"
version = "0.1.0"
edition = "2024"
# Firewall crate: the Lance stack (Arrow 57, DataFusion 52) is isolated
# from the rest of vectord (Arrow 55, DataFusion 47). Public API uses
# only std types — no Arrow types cross the crate boundary — so no
# version conflict propagates outward.
#
# See docs/ADR-019-vector-storage.md for the rationale: "productionizing
# will need either workspace-wide upgrade or a firewall via a dedicated
# vectord-lance crate." This is that firewall.
[dependencies]
lance = { version = "4.0", default-features = false }
lance-index = { version = "4.0", default-features = false }
lance-linalg = { version = "4.0", default-features = false }
# These Arrow/Parquet versions MUST match Lance's expectations — Lance
# re-exports their types through its API so any mismatch is a hard
# compile error. Keep in sync with lance-bench.
arrow = "57"
arrow-array = "57"
arrow-schema = "57"
parquet = "57"
tokio = { version = "1", features = ["full"] }
futures = "0.3"
serde = { version = "1", features = ["derive"] }
bytes = "1"
+528
View File
@@ -0,0 +1,528 @@
//! Production Lance vector backend (ADR-019 — hybrid architecture).
//!
//! This is the firewall crate. It owns its own Arrow 57 / DataFusion 52
//! / Lance 4 dependency tree. The public API uses only std types
//! (`Vec<f32>`, `Vec<String>`, `String`, `bool`) so nothing Arrow-shaped
//! crosses the crate boundary. That keeps vectord (Arrow 55) from
//! picking up an incompatible dep.
//!
//! Responsibilities:
//! - Migrate an existing binary-blob Parquet vector file into a Lance
//! dataset with `FixedSizeList<Float32, dims>`. One-time cost.
//! - Append new rows natively (no full-file rewrite — Lance's structural win).
//! - Build an IVF_PQ ANN index on the vector column.
//! - Vector search (`search`) using the IVF_PQ index when present,
//! falling back to full scan otherwise.
//! - Random-access row fetch by `doc_id` (`get_by_doc_id`) — the O(1)
//! lookup that Parquet-on-object-store can't cheaply do.
//! - Cheap count + stats introspection.
use arrow_array::{
Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, Int32Array,
RecordBatch, RecordBatchIterator, StringArray,
};
use arrow_schema::{DataType, Field, Schema};
use futures::StreamExt;
use serde::Serialize;
use std::sync::Arc;
use std::time::Instant;
// ================= Public types =================
/// One search result. Mirrors vectord's existing `SearchResult` shape
/// structurally but carries simpler types so this crate stays firewalled.
#[derive(Debug, Clone, Serialize)]
pub struct Hit {
pub doc_id: String,
pub chunk_text: String,
pub score: f32,
/// Optional — set by search_with_vector, not by index-only search.
pub distance: Option<f32>,
}
/// A fully-hydrated row fetched by `get_by_doc_id` — includes the vector
/// so callers can do downstream work (rerank, cite, etc.) without a
/// second round trip.
#[derive(Debug, Clone, Serialize)]
pub struct Row {
pub doc_id: String,
pub chunk_text: String,
pub vector: Vec<f32>,
pub source: Option<String>,
pub chunk_idx: Option<i32>,
}
#[derive(Debug, Clone, Serialize)]
pub struct MigrationStats {
pub rows_written: usize,
pub dimensions: usize,
pub disk_bytes: u64,
pub duration_secs: f32,
}
#[derive(Debug, Clone, Serialize)]
pub struct AppendStats {
pub rows_appended: usize,
pub disk_bytes_added: u64,
pub duration_secs: f32,
}
#[derive(Debug, Clone, Serialize)]
pub struct IndexStats {
pub name: String,
pub num_partitions: u32,
pub num_bits: u32,
pub num_sub_vectors: u32,
pub build_time_secs: f32,
pub disk_bytes_added: u64,
}
#[derive(Debug, Clone, Serialize)]
pub struct DatasetStats {
pub path: String,
pub rows: usize,
pub disk_bytes: u64,
pub has_vector_index: bool,
}
// ================= The backend =================
/// Thin wrapper around a Lance dataset path. Lance handles the heavy
/// lifting — we just expose a narrow API.
#[derive(Clone)]
pub struct LanceVectorStore {
/// Local filesystem path or object-store URI (file:///..., s3://...).
/// Lance's internal URI parsing handles both.
path: String,
}
impl LanceVectorStore {
pub fn new(path: impl Into<String>) -> Self {
Self { path: path.into() }
}
pub fn path(&self) -> &str { &self.path }
/// Row count via Lance's fast metadata path (no scan).
pub async fn count(&self) -> Result<usize, String> {
let dataset = open_or_err(&self.path).await?;
let n = dataset.count_rows(None).await.map_err(e)?;
Ok(n)
}
/// True if the on-disk dataset exists AND has at least one `vector`
/// column index attached.
pub async fn has_vector_index(&self) -> Result<bool, String> {
use lance_index::DatasetIndexExt;
let dataset = match lance::dataset::Dataset::open(&self.path).await {
Ok(d) => d,
Err(_) => return Ok(false),
};
let indexes = dataset.load_indices().await.map_err(e)?;
Ok(indexes.iter().any(|ix| {
ix.fields.iter().any(|fid| {
dataset.schema().field_by_id(*fid)
.map(|f| f.name == "vector")
.unwrap_or(false)
})
}))
}
pub async fn stats(&self) -> Result<DatasetStats, String> {
let rows = self.count().await.unwrap_or(0);
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
let has_vector_index = self.has_vector_index().await.unwrap_or(false);
Ok(DatasetStats {
path: self.path.clone(),
rows,
disk_bytes,
has_vector_index,
})
}
/// Migrate a vectord-format Parquet file into a Lance dataset.
///
/// Input schema (vectord's binary-blob format):
/// - source : Utf8
/// - doc_id : Utf8
/// - chunk_idx : Int32
/// - chunk_text : Utf8
/// - vector : Binary (raw f32 little-endian bytes)
///
/// Output schema (Lance-friendly):
/// - source : Utf8
/// - doc_id : Utf8
/// - chunk_idx : Int32
/// - chunk_text : Utf8
/// - vector : FixedSizeList<Float32, dims>
///
/// Idempotent at the file level — if the target exists, it's
/// overwritten. Caller must manage destination paths.
pub async fn migrate_from_parquet_bytes(
&self,
parquet_bytes: &[u8],
) -> Result<MigrationStats, String> {
let t0 = Instant::now();
let (schema, batches, rows) = read_parquet(parquet_bytes)?;
let dims = detect_vector_dims(&batches)?;
let (new_schema, new_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
// Overwrite any prior dataset at this path.
let _ = std::fs::remove_dir_all(&strip_file_uri(&self.path));
write_dataset(&self.path, new_schema, new_batches).await?;
let disk_bytes = dir_size_bytes(&strip_file_uri(&self.path));
Ok(MigrationStats {
rows_written: rows,
dimensions: dims,
disk_bytes,
duration_secs: t0.elapsed().as_secs_f32(),
})
}
/// Native Lance append — does NOT rewrite existing files. New rows
/// land as a separate fragment; readers union across fragments at
/// query time. Contrast: our Parquet path requires rewriting the
/// entire vector file to add rows.
pub async fn append(
&self,
source: Option<String>,
doc_ids: Vec<String>,
chunk_idxs: Vec<i32>,
chunk_texts: Vec<String>,
vectors: Vec<Vec<f32>>,
) -> Result<AppendStats, String> {
let n = doc_ids.len();
if n == 0 {
return Ok(AppendStats { rows_appended: 0, disk_bytes_added: 0, duration_secs: 0.0 });
}
if chunk_idxs.len() != n || chunk_texts.len() != n || vectors.len() != n {
return Err(format!(
"append: length mismatch (doc_ids={n}, chunk_idxs={}, chunk_texts={}, vectors={})",
chunk_idxs.len(), chunk_texts.len(), vectors.len(),
));
}
let dims = vectors[0].len();
for (i, v) in vectors.iter().enumerate() {
if v.len() != dims {
return Err(format!("append: row {i} has {} dims, expected {}", v.len(), dims));
}
}
let t0 = Instant::now();
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
let src_arr = StringArray::from(
(0..n).map(|_| source.clone()).collect::<Vec<_>>()
);
let doc_id_arr = StringArray::from(doc_ids);
let chunk_idx_arr = Int32Array::from(chunk_idxs);
let chunk_text_arr = StringArray::from(chunk_texts);
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
for v in vectors { flat.extend(v); }
let values = Float32Array::from(flat);
let item_field = Arc::new(Field::new("item", DataType::Float32, true));
let vec_arr = FixedSizeListArray::try_new(
item_field.clone(), dims as i32, Arc::new(values), None,
).map_err(e)?;
let schema = Arc::new(Schema::new(vec![
Field::new("source", DataType::Utf8, true),
Field::new("doc_id", DataType::Utf8, false),
Field::new("chunk_idx", DataType::Int32, true),
Field::new("chunk_text", DataType::Utf8, true),
Field::new("vector", DataType::FixedSizeList(item_field, dims as i32), false),
]));
let arrays: Vec<ArrayRef> = vec![
Arc::new(src_arr), Arc::new(doc_id_arr), Arc::new(chunk_idx_arr),
Arc::new(chunk_text_arr), Arc::new(vec_arr),
];
let batch = RecordBatch::try_new(schema.clone(), arrays).map_err(e)?;
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
use lance::dataset::{Dataset, WriteMode, WriteParams};
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
Dataset::write(reader, &self.path, Some(params)).await.map_err(e)?;
Ok(AppendStats {
rows_appended: n,
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
duration_secs: t0.elapsed().as_secs_f32(),
})
}
/// Build an IVF_PQ vector index. Replaces any prior index with the
/// same name. Callers pass explicit params — sensible defaults for
/// ~100K × 768d: num_partitions=316 (≈√N), num_bits=8, num_sub_vectors=48.
pub async fn build_index(
&self,
num_partitions: u32,
num_bits: u32,
num_sub_vectors: u32,
) -> Result<IndexStats, String> {
use lance::dataset::Dataset;
use lance::index::vector::VectorIndexParams;
use lance_index::{DatasetIndexExt, IndexType};
use lance_linalg::distance::MetricType;
let pre_bytes = dir_size_bytes(&strip_file_uri(&self.path));
let t0 = Instant::now();
let mut dataset = Dataset::open(&self.path).await.map_err(e)?;
let params = VectorIndexParams::ivf_pq(
num_partitions as usize,
num_bits as u8,
num_sub_vectors as usize,
MetricType::Cosine,
50, // max_iterations — same as bench
);
dataset.create_index(
&["vector"],
IndexType::Vector,
Some("vec_idx".into()),
&params,
true, // replace
).await.map_err(e)?;
Ok(IndexStats {
name: "vec_idx".into(),
num_partitions,
num_bits,
num_sub_vectors,
build_time_secs: t0.elapsed().as_secs_f32(),
disk_bytes_added: dir_size_bytes(&strip_file_uri(&self.path)).saturating_sub(pre_bytes),
})
}
/// Search for top_k nearest neighbors of `query`. Uses the IVF_PQ
/// index if one exists; otherwise does a full scan (slow but
/// correct — useful during development before index build).
pub async fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<Hit>, String> {
use lance::dataset::Dataset;
let dataset = Dataset::open(&self.path).await.map_err(e)?;
let qarr = Float32Array::from(query.to_vec());
let mut scanner = dataset.scan();
scanner.nearest("vector", &qarr, top_k as usize).map_err(e)?;
scanner.project(&["doc_id", "chunk_text"]).map_err(e)?;
let mut stream = scanner.try_into_stream().await.map_err(e)?;
let mut hits: Vec<Hit> = Vec::with_capacity(top_k);
while let Some(batch) = stream.next().await {
let batch = batch.map_err(e)?;
let doc_ids = batch.column_by_name("doc_id")
.ok_or_else(|| "no doc_id column in search result".to_string())?
.as_any().downcast_ref::<StringArray>()
.ok_or_else(|| "doc_id is not StringArray".to_string())?;
let chunk_texts = batch.column_by_name("chunk_text")
.and_then(|c| c.as_any().downcast_ref::<StringArray>());
// Lance tacks on a `_distance` column for nearest() queries.
let distances = batch.column_by_name("_distance")
.and_then(|c| c.as_any().downcast_ref::<Float32Array>());
for row in 0..batch.num_rows() {
let d = distances.map(|a| a.value(row));
hits.push(Hit {
doc_id: doc_ids.value(row).to_string(),
chunk_text: chunk_texts.map(|a| a.value(row).to_string()).unwrap_or_default(),
score: d.map(|d| 1.0 - d).unwrap_or(0.0), // cosine distance → similarity
distance: d,
});
}
}
Ok(hits)
}
/// Fetch one row by doc_id. Implementation: Lance filter-pushdown
/// scan — O(1) with partition pruning on a proper btree index,
/// O(N) on vector-only datasets (still far faster than reading
/// the whole Parquet file). We don't build a scalar index on
/// doc_id yet; that's a future optimization.
pub async fn get_by_doc_id(&self, doc_id: &str) -> Result<Option<Row>, String> {
use lance::dataset::Dataset;
let dataset = Dataset::open(&self.path).await.map_err(e)?;
let filter = format!("doc_id = '{}'", doc_id.replace('\'', "''"));
let mut scanner = dataset.scan();
scanner.filter(&filter).map_err(e)?;
scanner.limit(Some(1), None).map_err(e)?;
let mut stream = scanner.try_into_stream().await.map_err(e)?;
while let Some(batch) = stream.next().await {
let batch = batch.map_err(e)?;
if batch.num_rows() == 0 { continue; }
return Ok(Some(row_from_batch(&batch, 0)?));
}
Ok(None)
}
}
// ================= Internal helpers =================
fn e<T: std::fmt::Display>(err: T) -> String { err.to_string() }
/// `file:///abs/path` → `/abs/path`. Leave other URI schemes as-is for
/// helpers that only work on local paths (dir_size_bytes, remove_dir_all).
fn strip_file_uri(uri: &str) -> String {
uri.strip_prefix("file://").unwrap_or(uri).to_string()
}
async fn open_or_err(path: &str) -> Result<lance::dataset::Dataset, String> {
lance::dataset::Dataset::open(path).await.map_err(e)
}
fn dir_size_bytes(path: &str) -> u64 {
fn recurse(p: &std::path::Path) -> u64 {
let Ok(meta) = std::fs::metadata(p) else { return 0; };
if meta.is_file() { return meta.len(); }
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
entries.filter_map(|e| e.ok()).map(|e| recurse(&e.path())).sum()
}
recurse(std::path::Path::new(path))
}
fn read_parquet(bytes: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize), String> {
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::copy_from_slice(bytes))
.map_err(e)?;
let schema = builder.schema().clone();
let reader = builder.build().map_err(e)?;
let batches: Vec<RecordBatch> = reader.collect::<Result<_, _>>().map_err(e)?;
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
Ok((schema, batches, rows))
}
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize, String> {
for batch in batches {
let idx = batch.schema().index_of("vector")
.map_err(|_| "no 'vector' column".to_string())?;
let col = batch.column(idx);
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
for i in 0..binary.len() {
if !binary.is_null(i) {
return Ok(binary.value(i).len() / 4);
}
}
} else if let Some(fsl) = col.as_any().downcast_ref::<FixedSizeListArray>() {
return Ok(fsl.value_length() as usize);
}
}
Err("could not determine vector dimensions".into())
}
fn convert_to_fixed_size_list(
schema: &Arc<Schema>,
batches: Vec<RecordBatch>,
dims: usize,
) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
let new_fields: Vec<Arc<Field>> = schema
.fields()
.iter()
.map(|f| {
if f.name() == "vector" {
Arc::new(Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
dims as i32,
),
false,
))
} else {
f.clone()
}
})
.collect();
let new_schema = Arc::new(Schema::new(new_fields));
let mut out = Vec::with_capacity(batches.len());
for batch in batches {
let vec_idx = batch.schema().index_of("vector").map_err(e)?;
let mut new_cols: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
for (i, col) in batch.columns().iter().enumerate() {
if i == vec_idx {
if let Some(bin) = col.as_any().downcast_ref::<BinaryArray>() {
new_cols.push(Arc::new(binary_to_fsl(bin, dims)?));
} else if col.as_any().is::<FixedSizeListArray>() {
// Already in the right shape — just clone.
new_cols.push(col.clone());
} else {
return Err("vector column is neither Binary nor FixedSizeList".into());
}
} else {
new_cols.push(col.clone());
}
}
out.push(RecordBatch::try_new(new_schema.clone(), new_cols).map_err(e)?);
}
Ok((new_schema, out))
}
fn binary_to_fsl(bin: &BinaryArray, dims: usize) -> Result<FixedSizeListArray, String> {
let n = bin.len();
let mut flat: Vec<f32> = Vec::with_capacity(n * dims);
for i in 0..n {
if bin.is_null(i) {
flat.extend(std::iter::repeat(0.0).take(dims));
continue;
}
let b = bin.value(i);
if b.len() != dims * 4 {
return Err(format!(
"row {i}: {} bytes vs expected {} ({} × f32)",
b.len(), dims * 4, dims,
));
}
for c in b.chunks_exact(4) {
flat.push(f32::from_le_bytes([c[0], c[1], c[2], c[3]]));
}
}
let values = Float32Array::from(flat);
let field = Arc::new(Field::new("item", DataType::Float32, true));
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None).map_err(e)
}
async fn write_dataset(
path: &str,
schema: Arc<Schema>,
batches: Vec<RecordBatch>,
) -> Result<(), String> {
use lance::dataset::{Dataset, WriteParams};
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
Dataset::write(reader, path, Some(WriteParams::default()))
.await.map_err(e)?;
Ok(())
}
fn row_from_batch(batch: &RecordBatch, row: usize) -> Result<Row, String> {
let doc_id = batch.column_by_name("doc_id")
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.map(|a| a.value(row).to_string())
.ok_or_else(|| "missing doc_id".to_string())?;
let chunk_text = batch.column_by_name("chunk_text")
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.map(|a| a.value(row).to_string())
.unwrap_or_default();
let source = batch.column_by_name("source")
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row).to_string()) });
let chunk_idx = batch.column_by_name("chunk_idx")
.and_then(|c| c.as_any().downcast_ref::<Int32Array>())
.and_then(|a| if a.is_null(row) { None } else { Some(a.value(row)) });
let vec_col = batch.column_by_name("vector")
.ok_or_else(|| "no vector column in row fetch".to_string())?;
let fsl = vec_col.as_any().downcast_ref::<FixedSizeListArray>()
.ok_or_else(|| "vector column is not FixedSizeList".to_string())?;
let inner = fsl.value(row);
let floats = inner.as_any().downcast_ref::<Float32Array>()
.ok_or_else(|| "vector inner not Float32".to_string())?;
let mut v = Vec::with_capacity(floats.len());
for i in 0..floats.len() { v.push(floats.value(i)); }
Ok(Row { doc_id, chunk_text, vector: v, source, chunk_idx })
}
+3
View File
@@ -8,6 +8,9 @@ shared = { path = "../shared" }
storaged = { path = "../storaged" }
aibridge = { path = "../aibridge" }
catalogd = { path = "../catalogd" }
# ADR-019 firewall — vectord-lance owns its own Arrow 57 / Lance 4 deps.
# Public API uses only std types so no version conflict propagates here.
vectord-lance = { path = "../vectord-lance" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
+711
View File
@@ -0,0 +1,711 @@
//! Phase 16.2 + 16.5 — The autotune agent.
//!
//! A long-running tokio task that watches the trial journal and
//! autonomously proposes + runs new HNSW configs. Distinct from
//! `autotune::run_autotune` which is synchronous (one HTTP call, grid
//! of trials, done). The agent is the continuous version: it sleeps,
//! wakes on triggers, proposes configs based on prior trial history,
//! runs them one at a time, and auto-promotes when it finds an
//! improvement.
//!
//! Design invariants:
//! - Trials are data (ADR-018). Every proposal reads the journal; every
//! attempt appends to it. The journal is the agent's memory.
//! - One trial at a time. Bounded Ollama load — the agent never fires
//! multiple parallel embeddings and respects `cooldown_between_trials_secs`.
//! - Rate-limited. `max_trials_per_hour` is a hard ceiling so a
//! misbehaving proposal function can't saturate the system.
//! - Never promotes below `min_recall`. Same safety gate as
//! `run_autotune` — we will not make the index worse.
//! - Triggered OR periodic. Ingest enqueues a `DatasetAppended` event
//! when a new batch lands; the agent also wakes periodically to keep
//! exploring even when nothing changed externally.
//! - Graceful shutdown via the `stop_tx` signal — the handle's Drop
//! doesn't force-kill, but `stop()` requests a clean exit after the
//! current trial.
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
use aibridge::client::AiClient;
use catalogd::registry::Registry as CatalogRegistry;
use crate::embedding_cache::EmbeddingCache;
use crate::harness;
use crate::hnsw::HnswStore;
use crate::index_registry::IndexRegistry;
use crate::promotion::{PromotionEntry, PromotionRegistry};
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
// -------- Public-facing types --------
/// Runtime configuration for the agent. Mirrored in shared::config under
/// `[agent]`. Defaults are conservative — designed to tune slowly in the
/// background without fighting real workloads for GPU time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConfig {
/// Master switch. When false, `spawn` returns a handle but the loop
/// doesn't run.
pub enabled: bool,
/// Periodic wake-up — even if the trigger queue is empty, every N
/// seconds the agent picks an index with trials and proposes one
/// more config. Keeps exploration alive on idle indexes.
pub cycle_interval_secs: u64,
/// Minimum gap between two trials on the SAME index. Prevents the
/// agent from hammering Ollama when a hot index has many pending
/// triggers in a row.
pub cooldown_between_trials_secs: u64,
/// Below this recall, a proposal is never promoted — even if it
/// beats the champion on latency.
pub min_recall: f32,
/// Budget cap: hard ceiling on trials per hour across all indexes.
/// When hit, the agent idles until the hour window rolls.
pub max_trials_per_hour: u32,
}
impl Default for AgentConfig {
fn default() -> Self {
Self {
enabled: false, // opt-in — don't auto-start until J turns it on
cycle_interval_secs: 60,
cooldown_between_trials_secs: 30,
min_recall: 0.9,
max_trials_per_hour: 30,
}
}
}
/// What caused the agent to look at a particular index. Recorded in the
/// trial's note field so we can tell "new data arrived" trials from
/// "periodic exploration" trials in the journal.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TriggerReason {
/// Ingest just appended to a dataset that has attached HNSW indexes.
DatasetAppended { dataset: String },
/// A human or another agent hit `/agent/enqueue/{index}`.
Manual,
/// Periodic wake — no external event, just keep exploring.
Periodic,
}
/// One unit of work for the agent.
#[derive(Debug, Clone)]
pub struct TriggerEvent {
pub index_name: String,
pub reason: TriggerReason,
pub enqueued_at: DateTime<Utc>,
}
impl TriggerEvent {
pub fn manual(index_name: impl Into<String>) -> Self {
Self { index_name: index_name.into(), reason: TriggerReason::Manual, enqueued_at: Utc::now() }
}
pub fn dataset_appended(index_name: impl Into<String>, dataset: impl Into<String>) -> Self {
Self {
index_name: index_name.into(),
reason: TriggerReason::DatasetAppended { dataset: dataset.into() },
enqueued_at: Utc::now(),
}
}
pub fn periodic(index_name: impl Into<String>) -> Self {
Self { index_name: index_name.into(), reason: TriggerReason::Periodic, enqueued_at: Utc::now() }
}
}
/// Observable snapshot of the agent's state — what `/agent/status` returns.
#[derive(Debug, Clone, Serialize)]
pub struct AgentStatus {
pub running: bool,
pub config: AgentConfig,
pub queue_depth: usize,
pub trials_run: u64,
pub promotions: u64,
pub trials_in_last_hour: u32,
pub last_event: Option<AgentEvent>,
pub started_at: Option<DateTime<Utc>>,
}
/// Last thing that happened — useful for "why didn't it do anything?" debugging.
#[derive(Debug, Clone, Serialize)]
pub struct AgentEvent {
pub at: DateTime<Utc>,
pub kind: String, // "trial_completed" | "promoted" | "skipped_rate_limit" | etc
pub index_name: Option<String>,
pub detail: String,
}
/// Handle returned by `spawn`. Holds the trigger sender + shared status +
/// stop signal.
#[derive(Clone)]
pub struct AgentHandle {
trigger_tx: mpsc::Sender<TriggerEvent>,
inner: Arc<AgentInner>,
}
struct AgentInner {
status: RwLock<AgentStatus>,
stop_tx: Mutex<Option<oneshot::Sender<()>>>,
queue_len: Mutex<usize>, // mirror of the channel capacity — for status reporting
recent_trials: Mutex<VecDeque<DateTime<Utc>>>, // ring of recent trial timestamps for rate limit
}
impl AgentHandle {
/// Enqueue a trigger. Returns Err if the agent isn't running or the
/// queue is full (backpressure — dropping events is correct here
/// since periodic exploration will pick up the slack).
pub async fn enqueue(&self, event: TriggerEvent) -> Result<(), String> {
self.trigger_tx.try_send(event).map_err(|e| format!("enqueue: {e}"))?;
let mut guard = self.inner.queue_len.lock().await;
*guard = guard.saturating_add(1);
// Update queue_depth in status for observability.
let mut s = self.inner.status.write().await;
s.queue_depth = *guard;
Ok(())
}
pub async fn status(&self) -> AgentStatus {
let mut s = self.inner.status.read().await.clone();
// Refresh rate-limit window from ring buffer.
let cutoff = Utc::now() - chrono::Duration::hours(1);
let ring = self.inner.recent_trials.lock().await;
s.trials_in_last_hour = ring.iter().filter(|t| **t >= cutoff).count() as u32;
s
}
/// Request a graceful stop. Returns immediately — the loop exits
/// after its current trial completes.
pub async fn stop(&self) -> bool {
let mut guard = self.inner.stop_tx.lock().await;
if let Some(tx) = guard.take() {
let _ = tx.send(());
true
} else {
false
}
}
}
// -------- Agent state holder --------
/// Everything the agent needs to run a trial. Mirrors the fields of
/// `VectorState` the agent actually uses. Kept separate so the service
/// layer builds it explicitly — no clone of unneeded state.
#[derive(Clone)]
pub struct AgentDeps {
pub store: Arc<dyn ObjectStore>,
pub ai_client: AiClient,
pub catalog: CatalogRegistry,
pub index_registry: IndexRegistry,
pub hnsw_store: HnswStore,
pub embedding_cache: EmbeddingCache,
pub trial_journal: TrialJournal,
pub promotion_registry: PromotionRegistry,
}
// -------- Spawn --------
/// Start the agent loop in a background tokio task. Returns a handle
/// the caller uses to enqueue events and read status.
pub fn spawn(config: AgentConfig, deps: AgentDeps) -> AgentHandle {
let (trigger_tx, trigger_rx) = mpsc::channel::<TriggerEvent>(256);
let (stop_tx, stop_rx) = oneshot::channel::<()>();
let status = AgentStatus {
running: config.enabled,
config: config.clone(),
queue_depth: 0,
trials_run: 0,
promotions: 0,
trials_in_last_hour: 0,
last_event: None,
started_at: if config.enabled { Some(Utc::now()) } else { None },
};
let inner = Arc::new(AgentInner {
status: RwLock::new(status),
stop_tx: Mutex::new(Some(stop_tx)),
queue_len: Mutex::new(0),
recent_trials: Mutex::new(VecDeque::with_capacity(64)),
});
if config.enabled {
tracing::info!(
"autotune agent started (cycle={}s, cooldown={}s, cap={}/hr, min_recall={})",
config.cycle_interval_secs, config.cooldown_between_trials_secs,
config.max_trials_per_hour, config.min_recall,
);
let loop_inner = inner.clone();
let loop_deps = deps.clone();
let loop_config = config.clone();
tokio::spawn(async move {
run_loop(loop_config, loop_deps, trigger_rx, stop_rx, loop_inner).await;
});
} else {
// Agent disabled — still drain the channel so sends don't back up.
tokio::spawn(async move {
let mut rx = trigger_rx;
while rx.recv().await.is_some() {}
});
tracing::info!("autotune agent configured but disabled (set [agent].enabled=true)");
}
AgentHandle { trigger_tx, inner }
}
// -------- Main loop --------
async fn run_loop(
config: AgentConfig,
deps: AgentDeps,
mut trigger_rx: mpsc::Receiver<TriggerEvent>,
mut stop_rx: oneshot::Receiver<()>,
inner: Arc<AgentInner>,
) {
let mut periodic = tokio::time::interval(std::time::Duration::from_secs(config.cycle_interval_secs));
// First tick fires immediately — skip it so we don't double-fire on startup.
periodic.tick().await;
loop {
let event = tokio::select! {
_ = &mut stop_rx => {
tracing::info!("autotune agent: stop signal received");
let mut s = inner.status.write().await;
s.running = false;
return;
}
maybe = trigger_rx.recv() => match maybe {
Some(ev) => {
let mut guard = inner.queue_len.lock().await;
*guard = guard.saturating_sub(1);
let mut s = inner.status.write().await;
s.queue_depth = *guard;
ev
}
None => {
tracing::info!("autotune agent: trigger channel closed");
return;
}
},
_ = periodic.tick() => {
// Periodic wake — pick an index with existing trials.
// If nothing's been tuned yet there's nothing to propose.
match pick_periodic_target(&deps).await {
Some(idx) => TriggerEvent::periodic(idx),
None => continue,
}
}
};
// Rate limit: count recent trials, skip if over budget.
if over_rate_limit(&inner, config.max_trials_per_hour).await {
record_event(&inner, "skipped_rate_limit", Some(&event.index_name),
format!("hit cap of {}/hour", config.max_trials_per_hour)).await;
continue;
}
// Per-index cooldown.
if cooling_down(&inner, &event.index_name, config.cooldown_between_trials_secs).await {
record_event(&inner, "skipped_cooldown", Some(&event.index_name),
format!("last trial too recent (<{}s)", config.cooldown_between_trials_secs)).await;
continue;
}
// Run one trial.
match run_one_cycle(&event, &deps, config.min_recall).await {
Ok(outcome) => {
mark_recent_trial(&inner).await;
{
let mut s = inner.status.write().await;
s.trials_run += 1;
if outcome.promoted { s.promotions += 1; }
}
record_event(&inner, if outcome.promoted { "promoted" } else { "trial_completed" },
Some(&event.index_name),
format!("config=ec{}/es{} recall={:.3} p50={:.0}us {}",
outcome.trial.config.ef_construction,
outcome.trial.config.ef_search,
outcome.trial.metrics.recall_at_k,
outcome.trial.metrics.search_latency_p50_us,
if outcome.promoted { "★ PROMOTED" } else { "" })).await;
}
Err(e) => {
record_event(&inner, "trial_error", Some(&event.index_name), e).await;
}
}
}
}
/// Result of one cycle — ran a trial, maybe promoted it.
struct CycleOutcome {
trial: Trial,
promoted: bool,
}
/// Core cycle: propose → build → bench → record → maybe promote.
async fn run_one_cycle(
event: &TriggerEvent,
deps: &AgentDeps,
min_recall: f32,
) -> Result<CycleOutcome, String> {
// Read history.
let history = deps.trial_journal.list(&event.index_name).await
.map_err(|e| format!("read journal: {e}"))?;
if history.is_empty() {
return Err(format!(
"no trials yet for '{}' — seed with at least one POST /hnsw/trial first",
event.index_name,
));
}
// Current champion (if any) is the promoted config.
let champion = deps.promotion_registry.get_current(&event.index_name).await;
let champion_trial = champion.as_ref().and_then(|p| {
history.iter().find(|t| t.id == p.trial_id).cloned()
});
// Propose the next config.
let Some(next_config) = propose_next_config(&history, champion_trial.as_ref()) else {
return Err("proposer returned None — search space exhausted".into());
};
// Validate bounds defensively.
if !(10..=400).contains(&next_config.ef_construction) {
return Err(format!("proposed ef_construction={} out of bounds", next_config.ef_construction));
}
if !(10..=200).contains(&next_config.ef_search) {
return Err(format!("proposed ef_search={} out of bounds", next_config.ef_search));
}
// Need a harness to measure. Use the most recent one in history.
// (A future refinement: remember per-index "canonical harness" on
// the index metadata. For now: latest wins.)
let harness_name = history.last().unwrap().eval_set.clone();
let mut harness_set = harness::EvalSet::load(&deps.store, &harness_name).await
.map_err(|e| format!("load harness '{harness_name}': {e}"))?;
let embeddings = deps.embedding_cache.get_or_load(&event.index_name).await
.map_err(|e| format!("embeddings: {e}"))?;
if !harness_set.ground_truth_built {
harness::compute_ground_truth(&mut harness_set, &embeddings, &deps.ai_client).await
.map_err(|e| format!("ground truth: {e}"))?;
harness_set.save(&deps.store).await.ok();
}
// Build + bench.
let trial_id = Trial::new_id();
let slot = format!("{}__{}", event.index_name, trial_id);
let build = deps.hnsw_store
.build_index_with_config(&slot, (*embeddings).clone(), &next_config)
.await?;
let query_vectors: Vec<Vec<f32>> = harness_set.queries
.iter().filter_map(|q| q.query_embedding.clone()).collect();
let bench = deps.hnsw_store.bench_search(&slot, &query_vectors, harness_set.k).await?;
let mut recalls = Vec::with_capacity(harness_set.queries.len());
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
if let Some(gt) = &q.ground_truth {
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
}
}
let mean_recall = if recalls.is_empty() { 0.0 } else {
recalls.iter().sum::<f32>() / recalls.len() as f32
};
let mut lats = bench.latencies_us.clone();
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let p = |pct: f32| -> f32 {
if lats.is_empty() { return 0.0; }
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
lats[idx.min(lats.len() - 1)]
};
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
let memory_bytes = (embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
let note = match &event.reason {
TriggerReason::DatasetAppended { dataset } => format!("agent: dataset_appended({dataset})"),
TriggerReason::Manual => "agent: manual".to_string(),
TriggerReason::Periodic => "agent: periodic".to_string(),
};
let trial = Trial {
id: trial_id,
index_name: event.index_name.clone(),
eval_set: harness_set.name.clone(),
config: next_config.clone(),
metrics: TrialMetrics {
build_time_secs: build.build_time_secs,
search_latency_p50_us: p(0.50),
search_latency_p95_us: p(0.95),
search_latency_p99_us: p(0.99),
recall_at_k: mean_recall,
memory_bytes,
vectors: build.vectors,
eval_queries: harness_set.queries.len(),
brute_force_latency_us: 0.0,
},
created_at: Utc::now(),
note: Some(note),
};
deps.trial_journal.append(&trial).await.ok();
deps.hnsw_store.drop(&slot).await;
// Promotion decision: the new trial must meet recall gate AND beat
// the current champion (higher recall OR same recall + lower p50).
let promoted = if trial.metrics.recall_at_k < min_recall {
false
} else {
let beats = match &champion_trial {
None => true, // no champion yet — anything passing the gate wins
Some(c) => beats_champion(&trial, c),
};
if beats {
let entry = PromotionEntry {
config: trial.config.clone(),
trial_id: trial.id.clone(),
promoted_at: Utc::now(),
promoted_by: "agent".to_string(),
note: Some(format!(
"auto-promote: recall={:.3} p50={:.0}us (was {:.3}/{:.0}us)",
trial.metrics.recall_at_k, trial.metrics.search_latency_p50_us,
champion_trial.as_ref().map(|t| t.metrics.recall_at_k).unwrap_or(0.0),
champion_trial.as_ref().map(|t| t.metrics.search_latency_p50_us).unwrap_or(0.0),
)),
};
deps.promotion_registry.promote(&event.index_name, entry).await.is_ok()
} else {
false
}
};
Ok(CycleOutcome { trial, promoted })
}
/// Champion-beat test: strictly higher recall, OR equal recall with
/// lower p50. Same rule as autotune::pick_winner — kept consistent so the
/// agent and the synchronous autotune agree on what "better" means.
fn beats_champion(candidate: &Trial, champion: &Trial) -> bool {
if candidate.metrics.recall_at_k > champion.metrics.recall_at_k {
return true;
}
if (candidate.metrics.recall_at_k - champion.metrics.recall_at_k).abs() < 1e-4
&& candidate.metrics.search_latency_p50_us < champion.metrics.search_latency_p50_us {
return true;
}
false
}
/// Propose the next HnswConfig given trial history and the current
/// champion.
///
/// ============================================================
/// J: THIS IS YOURS TO IMPLEMENT
/// ============================================================
///
/// Inputs:
/// - `history`: every trial ever run on this index, oldest first
/// - `champion`: the currently-promoted trial, if any
///
/// Output:
/// - `Some(HnswConfig)` with the config to try next
/// - `None` if you decide there's nothing worth trying (agent logs
/// "search space exhausted" and moves on)
///
/// Hard bounds the caller enforces: ef_construction ∈ [10, 400],
/// ef_search ∈ [10, 200]. Stay inside those — configs outside get
/// rejected and count as a wasted cycle.
///
/// Design options (pick one, or mix):
///
/// 1. LOCAL REFINEMENT (exploit-heavy):
/// Sample near champion ± small delta. Converges fast, risks local
/// minima. Good for "we know roughly where the sweet spot is."
///
/// 2. ε-GREEDY (mixed):
/// With prob ε, random sample from full bounds (explore). Otherwise
/// refinement around champion (exploit). ε=0.2 is a reasonable start.
/// Good for long-running tune with no prior knowledge.
///
/// 3. COARSE→FINE (annealed):
/// First N trials: wide random. Then shrink the neighborhood around
/// champion as more trials accumulate. Mimics simulated annealing.
///
/// 4. DEDUP-AWARE:
/// Whatever strategy, skip configs already in history. Prevents the
/// agent from re-running the same (ec, es) pair twice.
///
/// A simple starter implementation is provided below (local refinement
/// + dedup). Replace with your preferred strategy.
pub fn propose_next_config(history: &[Trial], champion: Option<&Trial>) -> Option<HnswConfig> {
// ε-greedy around the champion, dedup-aware.
//
// - With probability ε (≈0.25), sample a random config from the full
// bounds. Keeps exploration alive so we don't get stuck hill-climbing
// one axis.
// - Otherwise: perturb the champion symmetrically on BOTH axes (not
// just +20 / +40 like the starter did). Prefers small moves first
// so recall stays near the current level.
// - Always skip configs already in history — no point re-running.
// - Deterministic per-history: RNG is seeded from history length so
// the same journal state always proposes the same next config.
// Makes tests + offline replay reproducible.
let base = champion
.map(|t| t.config.clone())
.or_else(|| history.last().map(|t| t.config.clone()))
.unwrap_or_default();
let tried = |ec: usize, es: usize| -> bool {
history.iter().any(|t|
t.config.ef_construction == ec && t.config.ef_search == es
)
};
let clamp = |ec: i32, es: i32| -> (usize, usize) {
(ec.clamp(10, 400) as usize, es.clamp(10, 200) as usize)
};
// Tiny xorshift — no rand crate dep. Seeded from history length so the
// proposer is deterministic for a given journal state.
let mut rng = (history.len() as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
let mut next_u = || {
rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
rng
};
for _attempt in 0..32 {
let explore = (next_u() % 100) < 25; // ε = 0.25
let (ec, es) = if explore {
clamp(10 + (next_u() % 391) as i32, 10 + (next_u() % 191) as i32)
} else {
// Symmetric perturbation — signed steps on both axes.
let dec = [-40, -20, -10, 10, 20, 40][(next_u() % 6) as usize];
let des = [-20, -10, -5, 5, 10, 20][(next_u() % 6) as usize];
clamp(base.ef_construction as i32 + dec, base.ef_search as i32 + des)
};
if !tried(ec, es) {
return Some(HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) });
}
}
None // 32 attempts all landed on duplicates — likely saturated
}
// -------- Helpers --------
/// Find an index to poke on a periodic wake. Strategy: the one with the
/// most recent trial. If nothing's been trialed yet, return None.
async fn pick_periodic_target(deps: &AgentDeps) -> Option<String> {
// `/agent` runs against any index that has a trial journal. We don't
// have a "list all journals" helper, so we derive candidates from the
// promotion registry (indexes with a human ever promoting are live).
let promos = deps.promotion_registry.list_all().await.ok()?;
// Prefer the one most recently promoted — it's the one a human cares
// about right now.
promos.into_iter()
.filter_map(|f| f.current.map(|c| (f.index_name, c.promoted_at)))
.max_by_key(|(_, at)| *at)
.map(|(name, _)| name)
}
async fn over_rate_limit(inner: &Arc<AgentInner>, cap: u32) -> bool {
let cutoff = Utc::now() - chrono::Duration::hours(1);
let ring = inner.recent_trials.lock().await;
ring.iter().filter(|t| **t >= cutoff).count() as u32 >= cap
}
async fn cooling_down(inner: &Arc<AgentInner>, _index: &str, cooldown_secs: u64) -> bool {
// Minimal impl: gate on global most-recent trial rather than per-index.
// Per-index cooldown would be easy to add — keep a HashMap<String, DateTime>
// — but for Phase 16.2 MVP, global is fine. Ollama is the shared resource.
let ring = inner.recent_trials.lock().await;
if let Some(last) = ring.back() {
let since = Utc::now().signed_duration_since(*last);
return since < chrono::Duration::seconds(cooldown_secs as i64);
}
false
}
async fn mark_recent_trial(inner: &Arc<AgentInner>) {
let mut ring = inner.recent_trials.lock().await;
ring.push_back(Utc::now());
// Keep bounded.
while ring.len() > 256 {
ring.pop_front();
}
}
async fn record_event(
inner: &Arc<AgentInner>,
kind: &str,
index: Option<&str>,
detail: String,
) {
tracing::info!("agent: {} {}{}", kind,
index.map(|i| format!("[{i}] ")).unwrap_or_default(), detail);
let mut s = inner.status.write().await;
s.last_event = Some(AgentEvent {
at: Utc::now(),
kind: kind.to_string(),
index_name: index.map(String::from),
detail,
});
}
#[cfg(test)]
mod tests {
use super::*;
fn mk_trial(ec: usize, es: usize, recall: f32, p50: f32) -> Trial {
Trial {
id: format!("t-{ec}-{es}"),
index_name: "test".into(),
eval_set: "eval".into(),
config: HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) },
metrics: TrialMetrics {
build_time_secs: 1.0,
search_latency_p50_us: p50,
search_latency_p95_us: p50 * 1.5,
search_latency_p99_us: p50 * 2.0,
recall_at_k: recall,
memory_bytes: 0, vectors: 1000, eval_queries: 10,
brute_force_latency_us: 100.0,
},
created_at: Utc::now(),
note: None,
}
}
#[test]
fn propose_skips_duplicates() {
let hist = vec![
mk_trial(80, 30, 1.0, 500.0),
mk_trial(100, 30, 1.0, 520.0), // ec+20
];
let next = propose_next_config(&hist, Some(&hist[0])).unwrap();
// ec+20 is taken, so the proposer should skip it.
assert!(next.ef_construction != 100 || next.ef_search != 30);
}
#[test]
fn beats_champion_strict_recall() {
let champ = mk_trial(80, 30, 0.95, 500.0);
let better_recall = mk_trial(80, 30, 0.99, 600.0);
let worse_recall = mk_trial(80, 30, 0.90, 100.0);
assert!(beats_champion(&better_recall, &champ));
assert!(!beats_champion(&worse_recall, &champ));
}
#[test]
fn beats_champion_same_recall_lower_latency() {
let champ = mk_trial(80, 30, 1.0, 500.0);
let faster = mk_trial(60, 30, 1.0, 400.0);
let slower = mk_trial(60, 30, 1.0, 600.0);
assert!(beats_champion(&faster, &champ));
assert!(!beats_champion(&slower, &champ));
}
}
+13
View File
@@ -29,8 +29,21 @@ pub struct IndexMeta {
pub created_at: DateTime<Utc>,
pub build_time_secs: f32,
pub chunks_per_sec: f32,
/// Federation layer 2: which bucket holds this index's artifacts
/// (trial journal + promotion file). Defaults to "primary" for
/// pre-federation indexes — the serde default keeps old metadata
/// files readable without migration.
#[serde(default = "default_bucket")]
pub bucket: String,
/// ADR-019: which physical backend stores this index. `Parquet`
/// means storage_key points at our binary-blob Parquet file;
/// `Lance` means it points at a Lance dataset directory.
#[serde(default)]
pub vector_backend: shared::types::VectorBackend,
}
fn default_bucket() -> String { "primary".to_string() }
/// Registry of all vector indexes.
#[derive(Clone)]
pub struct IndexRegistry {
+112
View File
@@ -0,0 +1,112 @@
//! ADR-019 hybrid: routing layer for the Lance vector backend.
//!
//! Holds a `LanceRegistry` which maps an index name to its
//! `LanceVectorStore` instance, lazy-creating on first touch. Path
//! resolution: the index's bucket gives us a bucket root; we append
//! `lance/{index_name}` and use that as the dataset URI. For local
//! buckets that means a directory under the bucket's root.
//!
//! S3 buckets: in principle Lance accepts s3://... URIs, but the
//! firewall crate (`vectord-lance`) doesn't pull the S3 feature in by
//! default. When we promote profile buckets onto S3 in a future phase,
//! enable that feature and update `lance_uri_for` accordingly.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::registry::BucketRegistry;
use vectord_lance::LanceVectorStore;
use crate::index_registry::IndexRegistry;
/// Convert a bucket+index pair into the URI Lance should use as the
/// dataset path. Local-only for MVP; S3 when we wire that backend.
///
/// Path resolution mirrors lakehouse.toml's convention for local
/// buckets: ./data for primary, ./data/_rescue for rescue, ./data/_testing
/// for testing, ./data/_profiles/{sanitized} for profile:* buckets, and
/// ./data/_buckets/{sanitized} for everything else. Sanitization replaces
/// `:` with `_` so paths are filesystem-safe.
///
/// Refuses unknown buckets so a typo doesn't silently land Lance data
/// in a directory the rest of the system can't see.
pub fn lance_uri_for(
buckets: &BucketRegistry,
bucket: &str,
index_name: &str,
) -> Result<String, String> {
if !buckets.contains(bucket) {
return Err(format!("bucket '{bucket}' not registered"));
}
let root: PathBuf = match bucket {
"primary" => PathBuf::from("./data"),
"rescue" => PathBuf::from("./data/_rescue"),
"testing" => PathBuf::from("./data/_testing"),
b if b.starts_with("profile:") => {
let safe = b.replace(':', "_");
PathBuf::from(format!("./data/_profiles/{safe}"))
}
b => PathBuf::from(format!("./data/_buckets/{}", b.replace(':', "_"))),
};
let dataset_dir = root.join("lance").join(index_name);
// Pre-create the parent so Lance's first write doesn't trip on a
// missing ancestor. Lance handles the dataset directory itself.
let _ = std::fs::create_dir_all(root.join("lance"));
// Canonicalize after the parent is guaranteed to exist; if the
// dataset dir hasn't been created yet, canonicalize the parent and
// tack on the leaf name.
let abs = match std::fs::canonicalize(&root) {
Ok(p) => p.join("lance").join(index_name),
Err(_) => dataset_dir.clone(),
};
Ok(abs.to_string_lossy().to_string())
}
/// Lookup-by-index registry of `LanceVectorStore` handles. One handle
/// per index, lazy-created on first call. Cheap to keep alive — it's
/// just a path string + handle to Lance's metadata cache.
#[derive(Clone)]
pub struct LanceRegistry {
buckets: Arc<BucketRegistry>,
indexes: IndexRegistry,
stores: Arc<RwLock<HashMap<String, LanceVectorStore>>>,
}
impl LanceRegistry {
pub fn new(buckets: Arc<BucketRegistry>, indexes: IndexRegistry) -> Self {
Self {
buckets,
indexes,
stores: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Get (or lazy-create) the LanceVectorStore for an index.
/// Resolves bucket → URI via the registry; backend stays whatever
/// IndexMeta says (no enforcement here — caller decides whether to
/// route to Lance based on profile/index backend choice).
pub async fn store_for(&self, index_name: &str) -> Result<LanceVectorStore, String> {
if let Some(s) = self.stores.read().await.get(index_name) {
return Ok(s.clone());
}
let bucket = self.indexes
.get(index_name).await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string());
let uri = lance_uri_for(&self.buckets, &bucket, index_name)?;
let store = LanceVectorStore::new(uri);
self.stores.write().await.insert(index_name.to_string(), store.clone());
Ok(store)
}
/// For freshly-created indexes that don't have IndexMeta yet — caller
/// supplies the bucket explicitly.
pub async fn store_for_new(&self, index_name: &str, bucket: &str) -> Result<LanceVectorStore, String> {
let uri = lance_uri_for(&self.buckets, bucket, index_name)?;
let store = LanceVectorStore::new(uri);
self.stores.write().await.insert(index_name.to_string(), store.clone());
Ok(store)
}
}
+2
View File
@@ -1,4 +1,6 @@
pub mod agent;
pub mod autotune;
pub mod lance_backend;
pub mod chunker;
pub mod embedding_cache;
pub mod harness;
+49 -16
View File
@@ -19,13 +19,14 @@
//! - Agent loop — lives in `vectord::autotune`.
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use storaged::ops;
use storaged::registry::BucketRegistry;
use tokio::sync::RwLock;
use crate::index_registry::IndexRegistry;
use crate::trial::HnswConfig;
const PROMOTION_PREFIX: &str = "_hnsw_promotions";
@@ -55,14 +56,16 @@ pub struct PromotionFile {
#[derive(Clone)]
pub struct PromotionRegistry {
store: Arc<dyn ObjectStore>,
buckets: Arc<BucketRegistry>,
index_registry: IndexRegistry,
cache: Arc<RwLock<HashMap<String, PromotionFile>>>,
}
impl PromotionRegistry {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
Self {
store,
buckets,
index_registry,
cache: Arc::new(RwLock::new(HashMap::new())),
}
}
@@ -76,13 +79,26 @@ impl PromotionRegistry {
format!("{PROMOTION_PREFIX}/{safe}.json")
}
/// Resolve which bucket's store holds this index's promotion file.
/// Same rules as TrialJournal::bucket_for — follows IndexMeta.bucket,
/// defaults to primary when metadata is missing.
async fn store_for(&self, index_name: &str) -> Result<Arc<dyn object_store::ObjectStore>, String> {
let bucket = self.index_registry
.get(index_name)
.await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string());
self.buckets.get(&bucket)
}
/// Load (and cache) the promotion file for an index.
pub async fn load(&self, index_name: &str) -> Result<PromotionFile, String> {
if let Some(cached) = self.cache.read().await.get(index_name) {
return Ok(cached.clone());
}
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let file = match ops::get(&self.store, &key).await {
let file = match ops::get(&store, &key).await {
Ok(bytes) => serde_json::from_slice::<PromotionFile>(&bytes)
.map_err(|e| format!("parse promotion file: {e}"))?,
Err(_) => PromotionFile {
@@ -118,9 +134,10 @@ impl PromotionRegistry {
file.current = Some(entry);
file.index_name = index_name.to_string();
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
ops::put(&self.store, &key, json.into()).await?;
ops::put(&store, &key, json.into()).await?;
self.cache.write().await.insert(index_name.to_string(), file.clone());
tracing::info!(
@@ -146,9 +163,10 @@ impl PromotionRegistry {
file.current = None;
}
}
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
ops::put(&self.store, &key, json.into()).await?;
ops::put(&store, &key, json.into()).await?;
self.cache.write().await.insert(index_name.to_string(), file.clone());
tracing::info!("rolled back promotion for '{}'", index_name);
Ok(file)
@@ -169,19 +187,34 @@ impl PromotionRegistry {
}
/// List every index that has a promotion recorded (for operator UI).
///
/// Federation: scans EVERY registered bucket for promotion files.
/// Per-profile buckets each have their own `_hnsw_promotions/` so we
/// aggregate across them. Dedups by index_name — if the same index
/// somehow has promotion files in multiple buckets, the one from the
/// bucket recorded in IndexMeta wins.
pub async fn list_all(&self) -> Result<Vec<PromotionFile>, String> {
let keys = ops::list(&self.store, Some(&format!("{PROMOTION_PREFIX}/"))).await.unwrap_or_default();
let mut out = Vec::new();
for key in keys {
if !key.ends_with(".json") { continue; }
let bytes = match ops::get(&self.store, &key).await {
Ok(b) => b,
let bucket_infos = self.buckets.list().await;
let mut by_name: HashMap<String, PromotionFile> = HashMap::new();
for b in &bucket_infos {
let store = match self.buckets.get(&b.name) {
Ok(s) => s,
Err(_) => continue,
};
if let Ok(f) = serde_json::from_slice::<PromotionFile>(&bytes) {
out.push(f);
let keys = ops::list(&store, Some(&format!("{PROMOTION_PREFIX}/")))
.await.unwrap_or_default();
for key in keys {
if !key.ends_with(".json") { continue; }
let bytes = match ops::get(&store, &key).await {
Ok(b) => b,
Err(_) => continue,
};
if let Ok(f) = serde_json::from_slice::<PromotionFile>(&bytes) {
by_name.insert(f.index_name.clone(), f);
}
}
}
Ok(out)
Ok(by_name.into_values().collect())
}
}
+2
View File
@@ -305,6 +305,8 @@ async fn try_update_index_meta(
created_at: chrono::Utc::now(),
build_time_secs: 0.0,
chunks_per_sec: 0.0,
bucket: "primary".to_string(),
vector_backend: shared::types::VectorBackend::Parquet,
};
index_registry.register(meta).await
}
+410 -1
View File
@@ -11,7 +11,8 @@ use std::sync::Arc;
use aibridge::client::{AiClient, EmbedRequest};
use catalogd::registry::Registry as CatalogRegistry;
use crate::{autotune, chunker, embedding_cache, harness, hnsw, index_registry, jobs, promotion, rag, refresh, search, store, supervisor, trial};
use storaged::registry::BucketRegistry;
use crate::{agent, autotune, chunker, embedding_cache, harness, hnsw, index_registry, jobs, lance_backend, promotion, rag, refresh, search, store, supervisor, trial};
#[derive(Clone)]
pub struct VectorState {
@@ -27,6 +28,33 @@ pub struct VectorState {
pub catalog: CatalogRegistry,
/// Phase 16: promoted HNSW configs. Activation + autotune read/write here.
pub promotion_registry: promotion::PromotionRegistry,
/// Phase 16.2: handle to the background autotune agent. Always
/// present — if the agent is disabled in config, the handle drops
/// incoming triggers silently.
pub agent_handle: agent::AgentHandle,
/// Phase B (federation layer 2): bucket registry for per-profile
/// bucket auto-provisioning on activation.
pub bucket_registry: Arc<BucketRegistry>,
/// Phase C (two-profile VRAM gate): tracks which profile is currently
/// "active" on the GPU. Singleton — one profile at a time holds its
/// model in VRAM. Swapping profiles with different ollama_name unloads
/// the previous one (keep_alive=0) before preloading the new one.
///
/// `None` = no profile has been activated this session; any first
/// activation just preloads and takes the slot.
pub active_profile: Arc<tokio::sync::RwLock<Option<ActiveProfileSlot>>>,
/// ADR-019 hybrid: handles to Lance datasets keyed by index name.
/// Lazy-created on first /vectors/lance/* call.
pub lance: lance_backend::LanceRegistry,
}
/// What the active-profile singleton records. Narrow — we don't need the
/// full ModelProfile here, just enough to know what to unload on swap.
#[derive(Debug, Clone, Serialize)]
pub struct ActiveProfileSlot {
pub profile_id: String,
pub ollama_name: String,
pub activated_at: chrono::DateTime<chrono::Utc>,
}
pub fn router(state: VectorState) -> Router {
@@ -59,12 +87,26 @@ pub fn router(state: VectorState) -> Router {
// Phase 17: profile activation — pre-load caches + HNSW for this
// model's bound data. First search after activate is warm.
.route("/profile/{id}/activate", post(activate_profile))
.route("/profile/{id}/deactivate", post(deactivate_profile))
.route("/profile/{id}/search", post(profile_scoped_search))
// Phase 17 VRAM gate: which profile currently owns the GPU?
.route("/profile/active", get(get_active_profile))
// Phase 16: promotion + autotune
.route("/hnsw/promote/{index}/{trial_id}", post(promote_trial))
.route("/hnsw/rollback/{index}", post(rollback_promotion))
.route("/hnsw/promoted/{index}", get(get_promoted))
.route("/hnsw/autotune", post(run_autotune_endpoint))
// Phase 16.2: background autotune agent
.route("/agent/status", get(agent_status))
.route("/agent/stop", post(agent_stop))
.route("/agent/enqueue/{index_name}", post(agent_enqueue))
// ADR-019: Lance hybrid backend
.route("/lance/migrate/{index_name}", post(lance_migrate))
.route("/lance/index/{index_name}", post(lance_build_index))
.route("/lance/search/{index_name}", post(lance_search))
.route("/lance/doc/{index_name}/{doc_id}", get(lance_get_doc))
.route("/lance/append/{index_name}", post(lance_append))
.route("/lance/stats/{index_name}", get(lance_stats))
.with_state(state)
}
@@ -81,6 +123,11 @@ struct CreateIndexRequest {
documents: Vec<DocInput>,
chunk_size: Option<usize>,
overlap: Option<usize>,
/// Federation layer 2: optional bucket to hold this index's trial
/// journal + promotion file. Defaults to "primary" — pre-existing
/// clients that don't know about federation keep working unchanged.
#[serde(default)]
bucket: Option<String>,
}
#[derive(Deserialize)]
@@ -117,6 +164,7 @@ async fn create_index(
let n_docs = req.documents.len();
let n_chunks = chunks.len();
let index_name = req.index_name.clone();
let bucket = req.bucket.clone().unwrap_or_else(|| "primary".to_string());
// Create job and return immediately
let job_id = state.job_tracker.create(&index_name, n_chunks).await;
@@ -157,6 +205,8 @@ async fn create_index(
created_at: chrono::Utc::now(),
build_time_secs: elapsed,
chunks_per_sec: rate,
bucket: bucket.clone(),
vector_backend: shared::types::VectorBackend::Parquet,
};
let _ = registry.register(meta).await;
@@ -753,6 +803,11 @@ struct ActivateReport {
failures: Vec<String>,
total_vectors: usize,
duration_secs: f32,
/// Phase C: did we successfully preload the Ollama model?
model_preloaded: bool,
/// Phase C: which profile previously held the GPU slot, if any.
/// Useful for observability of the swap.
previous_profile: Option<String>,
}
#[derive(Serialize)]
@@ -788,6 +843,64 @@ async fn activate_profile(
let mut failures = Vec::new();
let mut total_vectors = 0usize;
// Phase 17 / C: VRAM-aware swap. If another profile currently holds
// the GPU and uses a DIFFERENT Ollama model than the one being
// activated, unload it first (keep_alive=0). Same-model activations
// skip the unload — no point churning a model that's already loaded.
let previous_slot = {
let guard = state.active_profile.read().await;
guard.clone()
};
if let Some(prev) = &previous_slot {
if prev.ollama_name != profile.ollama_name {
match state.ai_client.unload_model(&prev.ollama_name).await {
Ok(_) => tracing::info!(
"profile swap: unloaded '{}' ({} -> {})",
prev.ollama_name, prev.profile_id, profile.id,
),
Err(e) => failures.push(format!(
"unload previous model '{}': {e}", prev.ollama_name,
)),
}
}
}
// Federation layer 2: if this profile declares its own bucket and
// that bucket isn't registered yet, auto-provision it under the
// configured profile_root. This is the moment a "dormant" profile
// becomes live — its bucket exists and is readable/writable.
if let Some(bucket_name) = profile.bucket.clone() {
if !state.bucket_registry.contains(&bucket_name) {
let root = format!(
"{}/{}",
state.bucket_registry.profile_root().trim_end_matches('/'),
bucket_name.replace(':', "_"),
);
let bc = shared::config::BucketConfig {
name: bucket_name.clone(),
backend: "local".to_string(),
root: Some(root.clone()),
bucket: None,
region: None,
endpoint: None,
secret_ref: None,
};
match state.bucket_registry.add_bucket(bc).await {
Ok(info) => {
tracing::info!(
"profile '{}' activated bucket '{}' (root={}, reachable={})",
profile.id, bucket_name, root, info.reachable,
);
}
Err(e) => {
failures.push(format!(
"auto-provision bucket '{}': {}", bucket_name, e,
));
}
}
}
}
let all_indexes = state.index_registry.list(None, None).await;
for binding in &profile.bound_datasets {
@@ -847,6 +960,30 @@ async fn activate_profile(
}
}
// Preload the new profile's Ollama model proactively. Same-model
// re-activations are cheap (Ollama no-ops if already loaded).
let mut model_preloaded = false;
match state.ai_client.preload_model(&profile.ollama_name).await {
Ok(_) => {
model_preloaded = true;
tracing::info!("profile '{}' preloaded ollama model '{}'",
profile.id, profile.ollama_name);
}
Err(e) => failures.push(format!(
"preload ollama model '{}': {e}", profile.ollama_name,
)),
}
// Take the GPU slot.
{
let mut guard = state.active_profile.write().await;
*guard = Some(ActiveProfileSlot {
profile_id: profile.id.clone(),
ollama_name: profile.ollama_name.clone(),
activated_at: chrono::Utc::now(),
});
}
Ok(Json(ActivateReport {
profile_id: profile.id,
ollama_name: profile.ollama_name,
@@ -854,9 +991,52 @@ async fn activate_profile(
failures,
total_vectors,
duration_secs: t0.elapsed().as_secs_f32(),
model_preloaded,
previous_profile: previous_slot.map(|s| s.profile_id),
}))
}
/// Unload this profile's model and clear the active slot. No-op if the
/// caller isn't the currently-active profile.
async fn deactivate_profile(
State(state): State<VectorState>,
Path(profile_id): Path<String>,
) -> impl IntoResponse {
let profile = match state.catalog.get_profile(&profile_id).await {
Some(p) => p,
None => return Err((StatusCode::NOT_FOUND, format!("profile not found: {profile_id}"))),
};
let was_active = {
let mut guard = state.active_profile.write().await;
match guard.as_ref() {
Some(s) if s.profile_id == profile_id => {
let prev = s.clone();
*guard = None;
Some(prev)
}
_ => None,
}
};
// Regardless of whether it held the slot, we can still try to unload —
// the operator's intent is "get this model out of VRAM."
let unload_result = state.ai_client.unload_model(&profile.ollama_name).await;
Ok(Json(serde_json::json!({
"profile_id": profile.id,
"ollama_name": profile.ollama_name,
"was_active": was_active.is_some(),
"unloaded": unload_result.is_ok(),
"unload_error": unload_result.err(),
})))
}
async fn get_active_profile(State(state): State<VectorState>) -> impl IntoResponse {
let slot = state.active_profile.read().await.clone();
Json(slot)
}
#[derive(Deserialize)]
struct ProfileSearchRequest {
index_name: String,
@@ -1013,3 +1193,232 @@ async fn run_autotune_endpoint(
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- Phase 16.2: autotune agent endpoints ---
async fn agent_status(State(state): State<VectorState>) -> impl IntoResponse {
Json(state.agent_handle.status().await)
}
async fn agent_stop(State(state): State<VectorState>) -> impl IntoResponse {
let stopped = state.agent_handle.stop().await;
Json(serde_json::json!({ "stopped": stopped }))
}
async fn agent_enqueue(
State(state): State<VectorState>,
Path(index_name): Path<String>,
) -> impl IntoResponse {
let event = agent::TriggerEvent::manual(index_name);
match state.agent_handle.enqueue(event).await {
Ok(()) => Ok(Json(serde_json::json!({ "enqueued": true }))),
Err(e) => Err((StatusCode::SERVICE_UNAVAILABLE, e)),
}
}
// --- ADR-019: Lance hybrid backend HTTP surface ---
//
// Lance routes operate on the same `index_name` as the Parquet/HNSW path,
// but materialize the data as a Lance dataset on disk under
// `{bucket_root}/lance/{index_name}/`. The two backends are independent:
// you can have an index in both formats simultaneously. `IndexMeta.vector_backend`
// records which one is canonical for that index.
#[derive(Deserialize)]
struct LanceMigrateRequest {
/// Optional bucket override. Defaults to whatever the existing
/// IndexMeta says, or "primary" for indexes that don't exist yet.
#[serde(default)]
bucket: Option<String>,
}
/// Read the existing Parquet vector file for `index_name` from object
/// storage, hand the bytes to vectord-lance, return migration stats.
/// The original Parquet file is left intact — both backends coexist
/// after migration.
async fn lance_migrate(
State(state): State<VectorState>,
Path(index_name): Path<String>,
Json(req): Json<LanceMigrateRequest>,
) -> impl IntoResponse {
let meta = state.index_registry.get(&index_name).await
.ok_or((StatusCode::NOT_FOUND, format!("index not found: {index_name}")))?;
let bucket = req.bucket.unwrap_or(meta.bucket.clone());
// Pull the Parquet bytes via storaged::ops — same path as the
// existing embedding loader uses.
let store = state.bucket_registry.get(&bucket)
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let bytes = storaged::ops::get(&store, &meta.storage_key).await
.map_err(|e| (StatusCode::NOT_FOUND, format!("read parquet: {e}")))?;
let lance_store = state.lance.store_for_new(&index_name, &bucket).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let stats = lance_store.migrate_from_parquet_bytes(&bytes).await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
tracing::info!(
"lance migrate '{}': {} rows, {}d, {} bytes on disk, {:.2}s",
index_name, stats.rows_written, stats.dimensions,
stats.disk_bytes, stats.duration_secs,
);
Ok::<_, (StatusCode, String)>(Json(serde_json::json!({
"index_name": index_name,
"bucket": bucket,
"lance_path": lance_store.path(),
"stats": stats,
})))
}
#[derive(Deserialize)]
struct LanceIndexRequest {
#[serde(default = "default_partitions")]
num_partitions: u32,
#[serde(default = "default_bits")]
num_bits: u32,
#[serde(default = "default_subvectors")]
num_sub_vectors: u32,
}
fn default_partitions() -> u32 { 316 } // ≈√100K — sane for the reference dataset
fn default_bits() -> u32 { 8 }
fn default_subvectors() -> u32 { 48 } // 768/48 = 16 dims per subvector
/// Build the IVF_PQ index on the Lance dataset.
async fn lance_build_index(
State(state): State<VectorState>,
Path(index_name): Path<String>,
Json(req): Json<LanceIndexRequest>,
) -> impl IntoResponse {
let lance_store = state.lance.store_for(&index_name).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
match lance_store.build_index(req.num_partitions, req.num_bits, req.num_sub_vectors).await {
Ok(stats) => Ok(Json(stats)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct LanceSearchRequest {
/// Plain text query — embedded server-side for symmetry with the
/// existing /vectors/search path.
query: String,
#[serde(default = "default_top_k")]
top_k: usize,
}
fn default_top_k() -> usize { 5 }
/// Vector search against a Lance dataset. Embeds the query text via the
/// sidecar then calls Lance's nearest-neighbor scanner.
async fn lance_search(
State(state): State<VectorState>,
Path(index_name): Path<String>,
Json(req): Json<LanceSearchRequest>,
) -> impl IntoResponse {
let embed_resp = state.ai_client
.embed(EmbedRequest { texts: vec![req.query.clone()], model: None })
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("embed: {e}")))?;
if embed_resp.embeddings.is_empty() {
return Err((StatusCode::BAD_GATEWAY, "no embedding returned".into()));
}
let qv: Vec<f32> = embed_resp.embeddings[0].iter().map(|&x| x as f32).collect();
let lance_store = state.lance.store_for(&index_name).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let t0 = std::time::Instant::now();
let hits = lance_store.search(&qv, req.top_k).await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(serde_json::json!({
"index_name": index_name,
"query": req.query,
"method": "lance_ivf_pq",
"latency_us": t0.elapsed().as_micros() as u64,
"results": hits,
})))
}
/// Random-access fetch by doc_id — the O(1) lookup that's basically
/// impossible in our Parquet path without scanning the whole file.
async fn lance_get_doc(
State(state): State<VectorState>,
Path((index_name, doc_id)): Path<(String, String)>,
) -> impl IntoResponse {
let lance_store = state.lance.store_for(&index_name).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let t0 = std::time::Instant::now();
match lance_store.get_by_doc_id(&doc_id).await {
Ok(Some(row)) => Ok(Json(serde_json::json!({
"index_name": index_name,
"doc_id": doc_id,
"latency_us": t0.elapsed().as_micros() as u64,
"row": row,
}))),
Ok(None) => Err((StatusCode::NOT_FOUND, format!("doc_id not found: {doc_id}"))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct LanceAppendRequest {
/// Optional source tag — set on every appended row.
#[serde(default)]
source: Option<String>,
rows: Vec<LanceAppendRow>,
}
#[derive(Deserialize)]
struct LanceAppendRow {
doc_id: String,
#[serde(default)]
chunk_idx: Option<i32>,
chunk_text: String,
/// Pre-computed embedding. Caller is responsible for ensuring it
/// matches the dataset's dimensions and embedding model.
vector: Vec<f32>,
}
async fn lance_append(
State(state): State<VectorState>,
Path(index_name): Path<String>,
Json(req): Json<LanceAppendRequest>,
) -> impl IntoResponse {
if req.rows.is_empty() {
return Err((StatusCode::BAD_REQUEST, "rows array is empty".into()));
}
let lance_store = state.lance.store_for(&index_name).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let mut doc_ids = Vec::with_capacity(req.rows.len());
let mut chunk_idxs = Vec::with_capacity(req.rows.len());
let mut chunk_texts = Vec::with_capacity(req.rows.len());
let mut vectors = Vec::with_capacity(req.rows.len());
for r in req.rows {
doc_ids.push(r.doc_id);
chunk_idxs.push(r.chunk_idx.unwrap_or(0));
chunk_texts.push(r.chunk_text);
vectors.push(r.vector);
}
match lance_store.append(req.source, doc_ids, chunk_idxs, chunk_texts, vectors).await {
Ok(stats) => Ok(Json(stats)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn lance_stats(
State(state): State<VectorState>,
Path(index_name): Path<String>,
) -> impl IntoResponse {
let lance_store = state.lance.store_for(&index_name).await
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
match lance_store.stats().await {
Ok(s) => Ok(Json(s)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
+45 -18
View File
@@ -17,6 +17,9 @@ use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::append_log::AppendLog;
use storaged::registry::BucketRegistry;
use crate::index_registry::IndexRegistry;
/// HNSW build/search parameters the agent can tune.
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -86,18 +89,28 @@ impl Trial {
}
/// Per-index append log, lazy-created on first write.
///
/// Federation layer 2: the journal resolves each index's bucket from the
/// index registry and writes its JSONL batches to THAT bucket, not
/// primary. Back-compat is preserved by `IndexMeta::bucket` defaulting
/// to "primary" for pre-federation indexes. Indexes the registry has
/// never heard of (edge case — trials run before first register) fall
/// through to primary as well.
#[derive(Clone)]
pub struct TrialJournal {
store: Arc<dyn ObjectStore>,
/// Cache per-index AppendLog instances so the in-memory buffer persists
/// across calls.
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
buckets: Arc<BucketRegistry>,
index_registry: IndexRegistry,
/// Cache per (bucket, index) AppendLog so the in-memory buffer persists
/// across calls. Keyed by `(bucket, index_name)` so moving an index
/// between buckets is clean — the old journal stays intact.
logs: Arc<RwLock<HashMap<(String, String), Arc<AppendLog>>>>,
}
impl TrialJournal {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
Self {
store,
buckets,
index_registry,
logs: Arc::new(RwLock::new(HashMap::new())),
}
}
@@ -106,35 +119,49 @@ impl TrialJournal {
format!("_hnsw_trials/{}", index_name)
}
async fn log_for(&self, index_name: &str) -> Arc<AppendLog> {
if let Some(log) = self.logs.read().await.get(index_name) {
return log.clone();
/// Resolve which bucket holds this index's trial artifacts.
/// Falls back to primary for indexes without recorded metadata.
async fn bucket_for(&self, index_name: &str) -> String {
self.index_registry
.get(index_name)
.await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string())
}
async fn log_for(&self, index_name: &str) -> Result<Arc<AppendLog>, String> {
let bucket = self.bucket_for(index_name).await;
let key = (bucket.clone(), index_name.to_string());
if let Some(log) = self.logs.read().await.get(&key) {
return Ok(log.clone());
}
let mut guard = self.logs.write().await;
if let Some(log) = guard.get(index_name) {
return log.clone();
if let Some(log) = guard.get(&key) {
return Ok(log.clone());
}
let store = self.buckets.get(&bucket)?;
// Trials arrive one at a time during human/agent iteration — a low
// threshold gives "hit /trials and see my latest attempt" immediacy
// without creating one file per event.
let log = Arc::new(
AppendLog::new(self.store.clone(), Self::prefix(index_name))
AppendLog::new(store, Self::prefix(index_name))
.with_flush_threshold(4),
);
guard.insert(index_name.to_string(), log.clone());
log
guard.insert(key, log.clone());
Ok(log)
}
/// Append a trial record. In-memory buffered; persisted in batches.
pub async fn append(&self, trial: &Trial) -> Result<(), String> {
let line = serde_json::to_vec(trial).map_err(|e| e.to_string())?;
let log = self.log_for(&trial.index_name).await;
let log = self.log_for(&trial.index_name).await?;
log.append(line).await
}
/// Read all trials for an index (flushed batches + unflushed buffer).
pub async fn list(&self, index_name: &str) -> Result<Vec<Trial>, String> {
let log = self.log_for(index_name).await;
let log = self.log_for(index_name).await?;
let lines = log.read_all().await?;
let mut trials = Vec::with_capacity(lines.len());
for line in lines {
@@ -149,13 +176,13 @@ impl TrialJournal {
/// Explicit flush for callers that want write-through semantics
/// (e.g. an agent that wants to commit a trial before querying stats).
pub async fn flush(&self, index_name: &str) -> Result<(), String> {
let log = self.log_for(index_name).await;
let log = self.log_for(index_name).await?;
log.flush().await
}
/// Compact all batch files for an index into one.
pub async fn compact(&self, index_name: &str) -> Result<storaged::append_log::CompactStats, String> {
let log = self.log_for(index_name).await;
let log = self.log_for(index_name).await?;
log.compact().await
}