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
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,6 @@
|
||||
pub mod agent;
|
||||
pub mod autotune;
|
||||
pub mod lance_backend;
|
||||
pub mod chunker;
|
||||
pub mod embedding_cache;
|
||||
pub mod harness;
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user