Phase 37: Hot-swap async + Phase 38: Universal API skeleton
- JobTracker extended with JobType::ProfileActivation + Embed - activate_profile returns job_id immediately, work spawns in background - /v1/chat, /v1/usage, /v1/sessions endpoints (OpenAI-compatible) - Langfuse trace integration (Phase 40 early deliverable) - 12 gateway unit tests green, curl gates pass
This commit is contained in:
root
+185
-205
@@ -7,6 +7,7 @@ use axum::{
|
||||
};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest, GenerateRequest};
|
||||
@@ -196,7 +197,7 @@ async fn create_index(
|
||||
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;
|
||||
let job_id = state.job_tracker.create_embed(&index_name, n_chunks).await;
|
||||
tracing::info!("job {job_id}: indexing '{}' — {} docs → {} chunks (background)", index_name, n_docs, n_chunks);
|
||||
|
||||
// Spawn supervised dual-pipeline embedding
|
||||
@@ -240,7 +241,7 @@ async fn create_index(
|
||||
};
|
||||
let _ = registry.register(meta).await;
|
||||
|
||||
tracker.complete(&jid, key).await;
|
||||
tracker.complete(&jid, Some(json!({ "storage_key": key }))).await;
|
||||
tracing::info!("job {jid}: completed — {n_chunks} chunks in {elapsed:.0}s ({rate:.0}/sec)");
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -482,7 +483,7 @@ async fn _run_embedding_job_legacy(
|
||||
// Update progress
|
||||
let elapsed = start.elapsed().as_secs_f32();
|
||||
let rate = if elapsed > 0.0 { all_vectors.len() as f32 / elapsed } else { 0.0 };
|
||||
tracker.update_progress(job_id, all_vectors.len(), rate).await;
|
||||
tracker.update_embed_progress(job_id, all_vectors.len(), rate).await;
|
||||
|
||||
// Log every 100 batches
|
||||
if (i + 1) % 100 == 0 {
|
||||
@@ -1371,239 +1372,218 @@ async fn activate_profile(
|
||||
State(state): State<VectorState>,
|
||||
Path(profile_id): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
let t0 = std::time::Instant::now();
|
||||
|
||||
tracing::info!("[activate_profile] START profile_id={}", profile_id);
|
||||
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 mut warmed = Vec::new();
|
||||
let mut failures = Vec::new();
|
||||
let mut total_vectors = 0usize;
|
||||
let job_id = state.job_tracker.create_profile_activation(&profile_id).await;
|
||||
let job_id_for_response = job_id.clone();
|
||||
let tracker = state.job_tracker.clone();
|
||||
let catalog = state.catalog.clone();
|
||||
let index_registry = state.index_registry.clone();
|
||||
let bucket_registry = state.bucket_registry.clone();
|
||||
let lance = state.lance.clone();
|
||||
let embedding_cache = state.embedding_cache.clone();
|
||||
let hnsw_store = state.hnsw_store.clone();
|
||||
let promotion_registry = state.promotion_registry.clone();
|
||||
let ai_client = state.ai_client.clone();
|
||||
let active_profile = state.active_profile.clone();
|
||||
let profile_name = profile.ollama_name.clone();
|
||||
let profile_id_clone = profile.id.clone();
|
||||
let profile_bucket = profile.bucket.clone();
|
||||
let profile_bound = profile.bound_datasets.clone();
|
||||
let profile_hnsw = profile.hnsw_config.clone();
|
||||
let profile_backend = profile.vector_backend.clone();
|
||||
let profile_full = profile.clone();
|
||||
|
||||
// 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,
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
tokio::spawn(async move {
|
||||
let t0 = std::time::Instant::now();
|
||||
let mut warmed = Vec::new();
|
||||
let mut failures = Vec::new();
|
||||
let mut total_vectors = 0usize;
|
||||
let job_id = job_id;
|
||||
|
||||
// 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 previous_slot = {
|
||||
let guard = active_profile.read().await;
|
||||
guard.clone()
|
||||
};
|
||||
if let Some(prev) = &previous_slot {
|
||||
if prev.ollama_name != profile_name {
|
||||
match ai_client.unload_model(&prev.ollama_name).await {
|
||||
Ok(_) => tracing::info!(
|
||||
"profile swap: unloaded '{}' ({} -> {})",
|
||||
prev.ollama_name, prev.profile_id, profile_id_clone,
|
||||
),
|
||||
Err(e) => failures.push(format!("unload previous model '{}': {e}", prev.ollama_name)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let all_indexes = state.index_registry.list(None, None).await;
|
||||
let use_lance = profile.vector_backend == shared::types::VectorBackend::Lance;
|
||||
|
||||
for binding in &profile.bound_datasets {
|
||||
let matched: Vec<_> = all_indexes
|
||||
.iter()
|
||||
.filter(|m| &m.source == binding)
|
||||
.collect();
|
||||
if matched.is_empty() {
|
||||
failures.push(format!(
|
||||
"no vector index found for binding '{}'", binding,
|
||||
));
|
||||
continue;
|
||||
}
|
||||
for meta in matched {
|
||||
if use_lance {
|
||||
// --- Lance activation path ---
|
||||
// Ensure a Lance dataset exists for this index. If it
|
||||
// doesn't, auto-migrate from the Parquet blob. Then
|
||||
// ensure an IVF_PQ index is built.
|
||||
let bucket = meta.bucket.clone();
|
||||
let lance_store = match state.lance.store_for_new(&meta.index_name, &bucket).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
failures.push(format!("{}: lance store init: {e}", meta.index_name));
|
||||
continue;
|
||||
}
|
||||
if let Some(bucket_name) = profile_bucket.clone() {
|
||||
if !bucket_registry.contains(&bucket_name) {
|
||||
let root = format!(
|
||||
"{}/{}",
|
||||
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,
|
||||
};
|
||||
let count = lance_store.count().await.unwrap_or(0);
|
||||
if count == 0 {
|
||||
// Auto-migrate from existing Parquet.
|
||||
let pq_store = match state.bucket_registry.get(&bucket) {
|
||||
match bucket_registry.add_bucket(bc).await {
|
||||
Ok(info) => {
|
||||
tracing::info!(
|
||||
"profile '{}' activated bucket '{}' (root={}, reachable={})",
|
||||
profile_id_clone, bucket_name, root, info.reachable,
|
||||
);
|
||||
}
|
||||
Err(e) => failures.push(format!("auto-provision bucket '{}': {}", bucket_name, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let all_indexes = index_registry.list(None, None).await;
|
||||
let use_lance = profile_backend == shared::types::VectorBackend::Lance;
|
||||
|
||||
for binding in &profile_bound {
|
||||
let matched: Vec<_> = all_indexes
|
||||
.iter()
|
||||
.filter(|m| &m.source == binding)
|
||||
.collect();
|
||||
if matched.is_empty() {
|
||||
failures.push(format!("no vector index found for binding '{}'", binding));
|
||||
continue;
|
||||
}
|
||||
for meta in matched {
|
||||
if use_lance {
|
||||
let bucket = meta.bucket.clone();
|
||||
let lance_store = match lance.store_for_new(&meta.index_name, &bucket).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => { failures.push(format!("{}: bucket: {e}", meta.index_name)); continue; }
|
||||
Err(e) => { failures.push(format!("{}: lance store init: {e}", meta.index_name)); continue; }
|
||||
};
|
||||
match storaged::ops::get(&pq_store, &meta.storage_key).await {
|
||||
Ok(bytes) => {
|
||||
let build_t = std::time::Instant::now();
|
||||
match lance_store.migrate_from_parquet_bytes(&bytes).await {
|
||||
Ok(ms) => {
|
||||
total_vectors += ms.rows_written;
|
||||
tracing::info!(
|
||||
"lance auto-migrate '{}': {} rows in {:.2}s",
|
||||
meta.index_name, ms.rows_written, ms.duration_secs,
|
||||
);
|
||||
warmed.push(WarmedIndex {
|
||||
index_name: meta.index_name.clone(),
|
||||
source: meta.source.clone(),
|
||||
vectors: ms.rows_written,
|
||||
hnsw_build_secs: build_t.elapsed().as_secs_f32(),
|
||||
});
|
||||
let count = lance_store.count().await.unwrap_or(0);
|
||||
if count == 0 {
|
||||
let pq_store = match bucket_registry.get(&bucket) {
|
||||
Ok(s) => s,
|
||||
Err(e) => { failures.push(format!("{}: bucket: {e}", meta.index_name)); continue; }
|
||||
};
|
||||
match storaged::ops::get(&pq_store, &meta.storage_key).await {
|
||||
Ok(bytes) => {
|
||||
let build_t = std::time::Instant::now();
|
||||
match lance_store.migrate_from_parquet_bytes(&bytes).await {
|
||||
Ok(ms) => {
|
||||
total_vectors += ms.rows_written;
|
||||
tracing::info!("lance auto-migrate '{}': {} rows in {:.2}s", meta.index_name, ms.rows_written, ms.duration_secs);
|
||||
warmed.push(WarmedIndex {
|
||||
index_name: meta.index_name.clone(),
|
||||
source: meta.source.clone(),
|
||||
vectors: ms.rows_written,
|
||||
hnsw_build_secs: build_t.elapsed().as_secs_f32(),
|
||||
});
|
||||
}
|
||||
Err(e) => failures.push(format!("{}: lance migrate: {e}", meta.index_name)),
|
||||
}
|
||||
Err(e) => failures.push(format!("{}: lance migrate: {e}", meta.index_name)),
|
||||
}
|
||||
Err(e) => failures.push(format!("{}: read parquet: {e}", meta.index_name)),
|
||||
}
|
||||
Err(e) => failures.push(format!("{}: read parquet: {e}", meta.index_name)),
|
||||
}
|
||||
} else {
|
||||
total_vectors += count;
|
||||
warmed.push(WarmedIndex {
|
||||
index_name: meta.index_name.clone(),
|
||||
source: meta.source.clone(),
|
||||
vectors: count,
|
||||
hnsw_build_secs: 0.0,
|
||||
});
|
||||
}
|
||||
// Ensure IVF_PQ vector index exists.
|
||||
if !lance_store.has_vector_index().await.unwrap_or(false) {
|
||||
match lance_store.build_index(316, 8, 48).await {
|
||||
Ok(ix) => tracing::info!(
|
||||
"lance auto-index '{}': IVF_PQ built in {:.1}s",
|
||||
meta.index_name, ix.build_time_secs,
|
||||
),
|
||||
Err(e) => failures.push(format!("{}: lance IVF_PQ build: {e}", meta.index_name)),
|
||||
}
|
||||
}
|
||||
// Ensure scalar btree on doc_id for O(log N) random fetch.
|
||||
if !lance_store.has_scalar_index("doc_id").await.unwrap_or(false) {
|
||||
match lance_store.build_scalar_index("doc_id").await {
|
||||
Ok(ix) => tracing::info!(
|
||||
"lance auto-index '{}': doc_id btree built in {:.2}s",
|
||||
meta.index_name, ix.build_time_secs,
|
||||
),
|
||||
Err(e) => failures.push(format!("{}: lance doc_id btree: {e}", meta.index_name)),
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// --- Parquet + HNSW activation path (existing) ---
|
||||
let embeddings = match state.embedding_cache.get_or_load(&meta.index_name).await {
|
||||
Ok(arc) => arc,
|
||||
Err(e) => {
|
||||
failures.push(format!("{}: load failed: {}", meta.index_name, e));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
total_vectors += embeddings.len();
|
||||
|
||||
let profile_default = trial::HnswConfig {
|
||||
ef_construction: profile.hnsw_config.ef_construction,
|
||||
ef_search: profile.hnsw_config.ef_search,
|
||||
seed: profile.hnsw_config.seed,
|
||||
};
|
||||
let cfg = state
|
||||
.promotion_registry
|
||||
.config_or(&meta.index_name, profile_default)
|
||||
.await;
|
||||
let build_t = std::time::Instant::now();
|
||||
match state
|
||||
.hnsw_store
|
||||
.build_index_with_config(&meta.index_name, (*embeddings).clone(), &cfg)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
} else {
|
||||
total_vectors += count;
|
||||
warmed.push(WarmedIndex {
|
||||
index_name: meta.index_name.clone(),
|
||||
source: meta.source.clone(),
|
||||
vectors: embeddings.len(),
|
||||
hnsw_build_secs: build_t.elapsed().as_secs_f32(),
|
||||
vectors: count,
|
||||
hnsw_build_secs: 0.0,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
failures.push(format!("{}: HNSW build failed: {}", meta.index_name, e));
|
||||
if !lance_store.has_vector_index().await.unwrap_or(false) {
|
||||
match lance_store.build_index(316, 8, 48).await {
|
||||
Ok(ix) => tracing::info!("lance auto-index '{}': IVF_PQ built in {:.1}s", meta.index_name, ix.build_time_secs),
|
||||
Err(e) => failures.push(format!("{}: lance IVF_PQ build: {e}", meta.index_name)),
|
||||
}
|
||||
}
|
||||
if !lance_store.has_scalar_index("doc_id").await.unwrap_or(false) {
|
||||
match lance_store.build_scalar_index("doc_id").await {
|
||||
Ok(ix) => tracing::info!("lance auto-index '{}': doc_id btree built in {:.2}s", meta.index_name, ix.build_time_secs),
|
||||
Err(e) => failures.push(format!("{}: lance doc_id btree: {e}", meta.index_name)),
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let embeddings = match embedding_cache.get_or_load(&meta.index_name).await {
|
||||
Ok(arc) => arc,
|
||||
Err(e) => { failures.push(format!("{}: load failed: {}", meta.index_name, e)); continue; }
|
||||
};
|
||||
total_vectors += embeddings.len();
|
||||
|
||||
let profile_default = trial::HnswConfig {
|
||||
ef_construction: profile_hnsw.ef_construction,
|
||||
ef_search: profile_hnsw.ef_search,
|
||||
seed: profile_hnsw.seed,
|
||||
};
|
||||
let cfg = promotion_registry
|
||||
.config_or(&meta.index_name, profile_default)
|
||||
.await;
|
||||
let build_t = std::time::Instant::now();
|
||||
match hnsw_store
|
||||
.build_index_with_config(&meta.index_name, (*embeddings).clone(), &cfg)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
warmed.push(WarmedIndex {
|
||||
index_name: meta.index_name.clone(),
|
||||
source: meta.source.clone(),
|
||||
vectors: embeddings.len(),
|
||||
hnsw_build_secs: build_t.elapsed().as_secs_f32(),
|
||||
});
|
||||
}
|
||||
Err(e) => failures.push(format!("{}: HNSW build failed: {}", meta.index_name, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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);
|
||||
let mut model_preloaded = false;
|
||||
match ai_client.preload_model(&profile_name).await {
|
||||
Ok(_) => {
|
||||
model_preloaded = true;
|
||||
tracing::info!("profile '{}' preloaded ollama model '{}'", profile_id_clone, profile_name);
|
||||
}
|
||||
Err(e) => failures.push(format!("preload ollama model '{}': {e}", profile_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(),
|
||||
});
|
||||
}
|
||||
{
|
||||
let mut guard = active_profile.write().await;
|
||||
*guard = Some(ActiveProfileSlot {
|
||||
profile_id: profile_id_clone.clone(),
|
||||
ollama_name: profile_name.clone(),
|
||||
activated_at: chrono::Utc::now(),
|
||||
});
|
||||
}
|
||||
|
||||
let report = ActivateReport {
|
||||
profile_id: profile.id,
|
||||
ollama_name: profile.ollama_name,
|
||||
indexes_warmed: warmed,
|
||||
failures,
|
||||
total_vectors,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
model_preloaded,
|
||||
previous_profile: previous_slot.map(|s| s.profile_id),
|
||||
};
|
||||
let result = serde_json::to_value(ActivateReport {
|
||||
profile_id: profile_id_clone,
|
||||
ollama_name: profile_name,
|
||||
indexes_warmed: warmed,
|
||||
failures,
|
||||
total_vectors,
|
||||
duration_secs: t0.elapsed().as_secs_f32(),
|
||||
model_preloaded,
|
||||
previous_profile: previous_slot.map(|s| s.profile_id),
|
||||
}).ok();
|
||||
|
||||
Ok(Json(report))
|
||||
tracker.complete(&job_id, result).await;
|
||||
});
|
||||
|
||||
Ok(Json(json!({
|
||||
"job_id": job_id_for_response,
|
||||
"message": format!("profile activation started — poll /vectors/jobs/{} for progress", job_id_for_response),
|
||||
})))
|
||||
}
|
||||
|
||||
/// Unload this profile's model and clear the active slot. No-op if the
|
||||
|
||||
Reference in New Issue
Block a user