diff --git a/Cargo.lock b/Cargo.lock index 008b065c..14285ffd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4527,6 +4527,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "http-range-header" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9171a2ea8a68358193d15dd5d70c1c10a2afc3e7e4c5bc92bc9f025cebd7359c" + [[package]] name = "httparse" version = "1.10.1" @@ -5423,6 +5429,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-context-master" +version = "0.6.3" +dependencies = [ + "arrow-array 58.3.0", + "arrow-schema 58.3.0", + "axum", + "chrono", + "clap", + "futures", + "lance 7.0.0", + "lance-context-api", + "lance-context-core", + "serde", + "serde_json", + "tempfile", + "tokio", + "tower-http", + "tracing", + "tracing-subscriber", +] + [[package]] name = "lance-context-python" version = "0.6.3" @@ -10312,6 +10340,11 @@ dependencies = [ "http 1.4.1", "http-body 1.0.1", "http-body-util", + "http-range-header", + "httpdate", + "mime", + "mime_guess", + "percent-encoding", "pin-project-lite", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 8a816d6f..c8133bc9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "crates/lance-context", "crates/lance-context-api", "crates/lance-context-server", + "crates/lance-context-master", "crates/lance-context-client", "python", ] diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index b9bfbbf3..03530c73 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -833,6 +833,71 @@ where } } +// --------------------------------------------------------------------------- +// Master control-plane DTOs +// --------------------------------------------------------------------------- + +/// One row of the control-plane experiment listing, backed by the master's +/// periodically-refreshed stats table. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ExperimentSummary { + /// Logical experiment / rollout-store name (unique key). + pub name: String, + /// Physical dataset URI. + pub uri: String, + /// Base-table logical row count as of the last scan. + pub row_count: i64, + /// Base-table fragment count as of the last scan. + pub fragment_count: i64, + /// Base-table manifest timestamp, Unix milliseconds. + pub last_updated: i64, + /// Flushed MemWAL generations pending merge as of the last scan. + pub pending_wal_generations: i64, + /// Timestamp of the last successful compaction, Unix ms, if any. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_compaction: Option, + /// Total compactions the master has driven for this experiment. + pub total_compactions: i64, + /// When the master last scanned this experiment, Unix milliseconds. + pub scanned_at: i64, +} + +/// Detailed view of a single experiment. Currently the same shape as +/// [`ExperimentSummary`]; a distinct type leaves room for detail-only fields. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ExperimentDetail { + #[serde(flatten)] + pub summary: ExperimentSummary, +} + +/// Paginated experiment listing. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ExperimentListResponse { + pub experiments: Vec, + /// Total number of experiments matching the (optional) search filter, + /// ignoring pagination. + pub total: i64, +} + +/// State of a manual or automatic compaction job for one experiment. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case", tag = "state")] +pub enum CompactJobStatus { + /// Queued but not yet started. + Queued, + /// Currently running. + Running, + /// Finished successfully. + Done { + fragments_removed: usize, + fragments_added: usize, + }, + /// Failed with an error message. + Failed { error: String }, + /// No compaction has been requested for this experiment. + None, +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 41114323..9e0d6559 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -34,7 +34,7 @@ pub use record::{ }; pub use registry::{RegistryEntry, RolloutRegistry}; pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; -pub use rollout_store::{rollout_schema, RolloutStore, RolloutStoreOptions}; +pub use rollout_store::{rollout_schema, RolloutObservation, RolloutStore, RolloutStoreOptions}; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, IdIndexType, ReadProjection, diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 179a0944..2abb3a15 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -85,6 +85,24 @@ use crate::store::{ /// shard state (mirrors the constant used by `ContextStore`). const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; +/// Cheap, read-only observability snapshot of a rollout store's base table. +/// +/// Produced by [`RolloutStore::observe`] from dataset metadata (no full scan). +/// Consumed by the control-plane stats scanner. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RolloutObservation { + /// Logical row count of the base table (fragment metadata, cheap). + pub row_count: i64, + /// Number of fragments in the base table. + pub fragment_count: i64, + /// Current base dataset manifest version. + pub version: u64, + /// Manifest timestamp, Unix milliseconds — when the base table last changed. + pub last_updated: i64, + /// Flushed MemWAL generations pending merge for this instance's shard. + pub pending_wal_generations: i64, +} + /// Configuration for opening a [`RolloutStore`]. #[derive(Debug, Clone, Default)] pub struct RolloutStoreOptions { @@ -632,6 +650,48 @@ impl RolloutStore { true } + /// Number of flushed MemWAL generations pending merge into the base table + /// for this instance's own shard. `0` when the shard has no manifest yet. + /// + /// Read-only: unlike [`Self::cleanup_own_shard`] it never merges. Used by the + /// control-plane stats scanner to surface read-amplification pressure. + pub async fn pending_wal_generations(&self) -> LanceResult { + let object_store = self.dataset.object_store(None).await?; + let branch_location = self.dataset.branch_location(); + let manifest_store = ShardManifestStore::new( + object_store, + &branch_location.path, + self.write_shard, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + match manifest_store.read_latest().await? { + Some(manifest) => Ok(manifest.flushed_generations.len()), + None => Ok(0), + } + } + + /// Snapshot cheap, read-only observability metrics for the base table. + /// + /// All fields come from dataset metadata (manifest + fragment counts), not a + /// full table scan: `row_count` and `fragment_count` read fragment metadata, + /// `version`/`last_updated` come from the manifest. `pending_wal_generations` + /// reads this instance's shard manifest. Intended for the control-plane stats + /// scanner, which opens each store read-only and records the result. + pub async fn observe(&self) -> LanceResult { + let row_count = self.dataset.count_rows(None).await? as i64; + let fragment_count = self.dataset.count_fragments() as i64; + let version = self.dataset.manifest.version; + let last_updated = self.dataset.manifest.timestamp().timestamp_millis(); + let pending_wal_generations = self.pending_wal_generations().await? as i64; + Ok(RolloutObservation { + row_count, + fragment_count, + version, + last_updated, + pending_wal_generations, + }) + } + /// Current compaction statistics for the base table. /// /// `is_compacting` is always `false`: compaction here runs synchronously @@ -1628,6 +1688,33 @@ mod tests { }); } + #[test] + fn observe_reports_cheap_metrics() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open(&uri).await.unwrap(); + // Empty store: no rows, no pending generations. + let obs = store.observe().await.unwrap(); + assert_eq!(obs.row_count, 0); + assert_eq!(obs.pending_wal_generations, 0); + + store + .add(&[assistant_record("a"), assistant_record("b")]) + .await + .unwrap(); + + // Two MemWAL-appended rows are visible via the LSM read path; the + // base table itself has not been merged, so they surface as a + // pending flushed generation rather than base rows. + let obs = store.observe().await.unwrap(); + assert!(obs.fragment_count >= 0); + assert!(obs.pending_wal_generations >= 1); + assert!(obs.last_updated > 0); + }); + } + #[test] fn appends_accumulate_and_are_immutable() { // MemWAL appends are append-only: successive `add`s accumulate, and a diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml new file mode 100644 index 00000000..d37f7052 --- /dev/null +++ b/crates/lance-context-master/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "lance-context-master" +version = "0.6.3" +edition = "2021" +license = "Apache-2.0" +authors = ["Lance Devs "] +repository = "https://github.com/lancedb/lance-context" +description = "Control-plane (master) process for lance-context rollout stores" +keywords = ["context", "lance", "control-plane", "compaction"] + +[[bin]] +name = "lance-context-master" +path = "src/main.rs" + +[dependencies] +lance-context-core = { version = "0.6.3", path = "../lance-context-core" } +lance-context-api = { version = "0.6.3", path = "../lance-context-api" } +arrow-array = "58" +arrow-schema = "58" +axum = { version = "0.8", features = ["json"] } +chrono = { version = "0.4", default-features = false, features = ["clock"] } +clap = { version = "4", features = ["derive", "env"] } +futures = "0.3" +lance = "7.0.0" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "sync"] } +tower-http = { version = "0.6", features = ["cors", "fs", "trace"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +[dev-dependencies] +tempfile = "3" diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs new file mode 100644 index 00000000..546966d5 --- /dev/null +++ b/crates/lance-context-master/src/config.rs @@ -0,0 +1,47 @@ +//! Command-line / environment configuration for the master control-plane. + +use clap::Parser; + +/// Control-plane (master) process for lance-context rollout stores. +#[derive(Debug, Clone, Parser)] +#[command(name = "lance-context-master", version)] +pub struct MasterConfig { + /// Data directory / object-store prefix shared with the data-plane server. + #[arg(long, env = "DATA_DIR", default_value = "./data")] + pub data_dir: String, + + /// Host to bind the admin API + UI to. + #[arg(long, env = "MASTER_HOST", default_value = "0.0.0.0")] + pub host: String, + + /// Port to bind the admin API + UI to. + #[arg(long, env = "MASTER_PORT", default_value_t = 8090)] + pub port: u16, + + /// Interval, in seconds, between full stats scans of every experiment. + /// `0` disables the background scanner. + #[arg(long, env = "STATS_SCAN_INTERVAL_SECS", default_value_t = 300)] + pub stats_scan_interval_secs: u64, + + /// Bounded concurrency when scanning experiments each round. + #[arg(long, env = "SCAN_CONCURRENCY", default_value_t = 8)] + pub scan_concurrency: usize, + + /// Interval, in seconds, between automatic compaction sweeps. `0` disables + /// automatic compaction (manual triggers still work). + #[arg(long, env = "COMPACTION_INTERVAL_SECS", default_value_t = 600)] + pub compaction_interval_secs: u64, + + /// Minimum fragment count before an experiment is auto-compacted. + #[arg(long, env = "MIN_FRAGMENTS", default_value_t = 16)] + pub min_fragments: usize, + + /// Target rows per fragment passed to compaction. + #[arg(long, env = "TARGET_ROWS_PER_FRAGMENT", default_value_t = 1_048_576)] + pub target_rows_per_fragment: usize, + + /// Directory of built UI assets to serve. When unset, only the JSON API is + /// exposed. + #[arg(long, env = "UI_DIR")] + pub ui_dir: Option, +} diff --git a/crates/lance-context-master/src/error.rs b/crates/lance-context-master/src/error.rs new file mode 100644 index 00000000..75d9ace9 --- /dev/null +++ b/crates/lance-context-master/src/error.rs @@ -0,0 +1,30 @@ +//! Admin HTTP error type. + +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::Json; +use serde_json::json; + +/// Errors surfaced by the master admin API. +#[derive(Debug)] +pub enum MasterError { + NotFound(String), + Internal(String), +} + +impl MasterError { + /// Map a Lance error to an internal server error. + pub fn from_lance(err: lance::Error) -> Self { + MasterError::Internal(err.to_string()) + } +} + +impl IntoResponse for MasterError { + fn into_response(self) -> Response { + let (status, msg) = match self { + MasterError::NotFound(m) => (StatusCode::NOT_FOUND, m), + MasterError::Internal(m) => (StatusCode::INTERNAL_SERVER_ERROR, m), + }; + (status, Json(json!({ "error": msg }))).into_response() + } +} diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs new file mode 100644 index 00000000..7efe511d --- /dev/null +++ b/crates/lance-context-master/src/lib.rs @@ -0,0 +1,8 @@ +//! Control-plane (master) library surface. + +pub mod config; +pub mod error; +pub mod routes; +pub mod scanner; +pub mod state; +pub mod stats_store; diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs new file mode 100644 index 00000000..b29a72c4 --- /dev/null +++ b/crates/lance-context-master/src/main.rs @@ -0,0 +1,73 @@ +use axum::Router; +use clap::Parser; +use tokio::net::TcpListener; +use tower_http::cors::CorsLayer; +use tower_http::services::ServeDir; +use tower_http::trace::TraceLayer; +use tracing_subscriber::EnvFilter; + +use lance_context_master::config::MasterConfig; +use lance_context_master::routes; +use lance_context_master::scanner; +use lance_context_master::state::MasterState; + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt() + .with_env_filter(EnvFilter::from_default_env().add_directive("info".parse().unwrap())) + .init(); + + let config = MasterConfig::parse(); + let addr = format!("{}:{}", config.host, config.port); + + if let Err(e) = std::fs::create_dir_all(&config.data_dir) { + tracing::error!( + "Failed to create data directory '{}': {}", + config.data_dir, + e + ); + std::process::exit(1); + } + + let ui_dir = config.ui_dir.clone(); + + let state = match MasterState::new(config).await { + Ok(state) => state, + Err(e) => { + tracing::error!("Failed to initialize master state: {:?}", e); + std::process::exit(1); + } + }; + + // Background stats scanner (detached; stops when the last Arc drops). + let _scanner = scanner::spawn_scanner(&state); + + let mut app = Router::new().nest("/api/v1", routes::api_router()); + + // Serve the built UI (SPA) as a fallback when a `--ui-dir` is configured. + if let Some(dir) = ui_dir { + let index = format!("{}/index.html", dir.trim_end_matches('/')); + let serve = ServeDir::new(&dir).fallback(tower_http::services::ServeFile::new(index)); + app = app.fallback_service(serve); + } + + let app = app + .with_state(state) + .layer(TraceLayer::new_for_http()) + .layer(CorsLayer::permissive()); + + tracing::info!("Starting lance-context-master on {}", addr); + + let listener = TcpListener::bind(&addr).await.unwrap(); + axum::serve(listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await + .unwrap(); +} + +async fn shutdown_signal() { + tokio::signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + tracing::info!("Shutting down"); +} diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs new file mode 100644 index 00000000..c4eaa35a --- /dev/null +++ b/crates/lance-context-master/src/routes.rs @@ -0,0 +1,222 @@ +//! Admin JSON API routes (PR A: read-only observability endpoints). + +use std::sync::Arc; + +use axum::extract::{Path, Query, State}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use serde::Deserialize; + +use lance_context_api::{ExperimentDetail, ExperimentListResponse, ExperimentSummary}; + +use crate::error::MasterError; +use crate::scanner; +use crate::state::MasterState; + +/// Query params for the experiment listing. +#[derive(Debug, Deserialize)] +pub struct ListParams { + #[serde(default)] + pub search: Option, + #[serde(default = "default_limit")] + pub limit: usize, + #[serde(default)] + pub offset: usize, +} + +fn default_limit() -> usize { + 50 +} + +/// Query params for the detail endpoint. +#[derive(Debug, Deserialize)] +pub struct DetailParams { + /// When true, open the dataset and recompute metrics on-demand rather than + /// returning the last scanned snapshot. + #[serde(default)] + pub fresh: bool, +} + +/// `GET /api/v1/experiments` +pub async fn list_experiments( + State(state): State>, + Query(params): Query, +) -> Result, MasterError> { + let search = params.search.as_deref().filter(|s| !s.is_empty()); + let stats = state.stats.lock().await; + let total = stats.count(search).await.map_err(MasterError::from_lance)?; + let rows = stats + .list(search, params.limit, params.offset) + .await + .map_err(MasterError::from_lance)?; + let experiments: Vec = rows.into_iter().map(|r| r.into_summary()).collect(); + Ok(Json(ExperimentListResponse { experiments, total })) +} + +/// `GET /api/v1/experiments/{name}` +pub async fn get_experiment( + State(state): State>, + Path(name): Path, + Query(params): Query, +) -> Result, MasterError> { + if params.fresh { + // Recompute this one experiment now, updating the stats table as a + // side effect, then read it back. + let entries = state + .registry + .read() + .await + .list() + .await + .map_err(MasterError::from_lance)?; + if !entries.iter().any(|e| e.name == name) { + return Err(MasterError::NotFound(format!( + "experiment '{}' does not exist", + name + ))); + } + scanner::scan_once(&state) + .await + .map_err(MasterError::from_lance)?; + } + let stats = state.stats.lock().await; + match stats.get(&name).await.map_err(MasterError::from_lance)? { + Some(row) => Ok(Json(ExperimentDetail { + summary: row.into_summary(), + })), + None => Err(MasterError::NotFound(format!( + "experiment '{}' not found in stats", + name + ))), + } +} + +/// `POST /api/v1/rescan` — trigger one immediate full scan. +pub async fn rescan( + State(state): State>, +) -> Result, MasterError> { + let n = scanner::scan_once(&state) + .await + .map_err(MasterError::from_lance)?; + Ok(Json(serde_json::json!({ "scanned": n }))) +} + +/// Build the admin API router (mounted under `/api/v1`). +pub fn api_router() -> Router> { + Router::new() + .route("/experiments", get(list_experiments)) + .route("/experiments/{name}", get(get_experiment)) + .route("/rescan", post(rescan)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::MasterConfig; + use lance_context_core::RolloutStore; + use tempfile::TempDir; + + fn test_config(dir: &TempDir) -> MasterConfig { + MasterConfig { + data_dir: dir.path().to_string_lossy().to_string(), + host: "127.0.0.1".to_string(), + port: 0, + stats_scan_interval_secs: 0, + scan_concurrency: 4, + compaction_interval_secs: 0, + min_fragments: 16, + target_rows_per_fragment: 1_048_576, + ui_dir: None, + } + } + + /// Create N experiments via core, register them, scan, and assert the stats + /// table + list endpoint reflect them. + #[tokio::test] + async fn scan_populates_and_list_returns_experiments() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(test_config(&dir)).await.unwrap(); + + for i in 0..3 { + let name = format!("exp-{i}"); + let uri = state.rollout_uri(&name); + // Creating the store materializes an (empty) base table on disk. + RolloutStore::open(&uri).await.unwrap(); + state + .registry + .write() + .await + .upsert(&name, &uri) + .await + .unwrap(); + } + + let scanned = scanner::scan_once(&state).await.unwrap(); + assert_eq!(scanned, 3); + + let Json(resp) = list_experiments( + State(state.clone()), + Query(ListParams { + search: None, + limit: 50, + offset: 0, + }), + ) + .await + .unwrap(); + assert_eq!(resp.total, 3); + assert_eq!(resp.experiments.len(), 3); + assert_eq!(resp.experiments[0].name, "exp-0"); + + // Search narrows. + let Json(one) = list_experiments( + State(state.clone()), + Query(ListParams { + search: Some("exp-1".to_string()), + limit: 50, + offset: 0, + }), + ) + .await + .unwrap(); + assert_eq!(one.total, 1); + assert_eq!(one.experiments[0].name, "exp-1"); + } + + #[tokio::test] + async fn scan_reconciles_removed_experiments() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(test_config(&dir)).await.unwrap(); + + let name = "gone"; + let uri = state.rollout_uri(name); + RolloutStore::open(&uri).await.unwrap(); + state + .registry + .write() + .await + .upsert(name, &uri) + .await + .unwrap(); + scanner::scan_once(&state).await.unwrap(); + assert!(state.stats.lock().await.get(name).await.unwrap().is_some()); + + // Remove from registry -> next scan drops the stats row. + state.registry.write().await.remove(name).await.unwrap(); + scanner::scan_once(&state).await.unwrap(); + assert!(state.stats.lock().await.get(name).await.unwrap().is_none()); + } + + #[tokio::test] + async fn get_experiment_not_found() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(test_config(&dir)).await.unwrap(); + let res = get_experiment( + State(state), + Path("missing".to_string()), + Query(DetailParams { fresh: false }), + ) + .await; + assert!(matches!(res, Err(MasterError::NotFound(_)))); + } +} diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs new file mode 100644 index 00000000..355c9623 --- /dev/null +++ b/crates/lance-context-master/src/scanner.rs @@ -0,0 +1,134 @@ +//! Background stats scanner. +//! +//! Every `stats_scan_interval_secs` the scanner enumerates all experiments from +//! the registry, opens each one read-only with bounded concurrency, records its +//! [`RolloutObservation`] into the stats table, and reconciles away rows for +//! experiments that have since been deleted from the registry. Failures on a +//! single experiment are logged and skipped; the next round retries. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use chrono::Utc; +use futures::stream::{self, StreamExt}; +use lance_context_core::{RolloutStore, RolloutStoreOptions}; +use tokio::task::JoinHandle; + +use crate::state::MasterState; +use crate::stats_store::StatRow; + +/// Per-experiment open+observe timeout so one wedged dataset cannot stall a +/// scan round. +const OBSERVE_TIMEOUT: Duration = Duration::from_secs(30); + +/// Run a single scan pass: refresh every experiment's stats row and drop rows +/// for experiments no longer in the registry. Returns the number of +/// experiments successfully observed. +pub async fn scan_once(state: &Arc) -> lance::Result { + let entries = state.registry.read().await.list().await?; + let live: HashSet = entries.iter().map(|e| e.name.clone()).collect(); + let concurrency = state.config.scan_concurrency.max(1); + + // Observe experiments concurrently (bounded), preserving prior compaction + // counters by reading the existing stats row first. + let observed: Vec = stream::iter(entries) + .map(|entry| { + let state = state.clone(); + async move { observe_one(&state, &entry.name, &entry.uri).await } + }) + .buffer_unordered(concurrency) + .filter_map(|row| async move { row }) + .collect() + .await; + + let count = observed.len(); + + // Upsert observed rows and reconcile deletions under the stats lock. + let mut stats = state.stats.lock().await; + for row in observed { + if let Err(e) = stats.upsert(&row).await { + tracing::warn!(store = %row.name, error = %e, "stats upsert failed"); + } + } + // Drop stats rows for experiments removed from the registry. + let existing = stats.list(None, usize::MAX, 0).await?; + for row in existing { + if !live.contains(&row.name) { + if let Err(e) = stats.remove(&row.name).await { + tracing::warn!(store = %row.name, error = %e, "stats reconcile-remove failed"); + } + } + } + Ok(count) +} + +/// Open one experiment read-only, observe it, and build its stats row. Preserves +/// `last_compaction`/`total_compactions` from any existing stats row. Returns +/// `None` (logged) on any error or timeout. +async fn observe_one(state: &Arc, name: &str, uri: &str) -> Option { + let opts = RolloutStoreOptions::default(); + let open = RolloutStore::open_existing_with_options(uri, opts); + let store = match tokio::time::timeout(OBSERVE_TIMEOUT, open).await { + Ok(Ok(store)) => store, + Ok(Err(e)) => { + tracing::warn!(store = %name, error = %e, "scan: open failed"); + return None; + } + Err(_) => { + tracing::warn!(store = %name, "scan: open timed out"); + return None; + } + }; + let obs = match tokio::time::timeout(OBSERVE_TIMEOUT, store.observe()).await { + Ok(Ok(obs)) => obs, + Ok(Err(e)) => { + tracing::warn!(store = %name, error = %e, "scan: observe failed"); + return None; + } + Err(_) => { + tracing::warn!(store = %name, "scan: observe timed out"); + return None; + } + }; + + // Carry compaction counters forward across scans. + let (last_compaction, total_compactions) = { + let stats = state.stats.lock().await; + match stats.get(name).await { + Ok(Some(prev)) => (prev.last_compaction, prev.total_compactions), + _ => (StatRow::NO_COMPACTION, 0), + } + }; + + Some(StatRow { + name: name.to_string(), + uri: uri.to_string(), + row_count: obs.row_count, + fragment_count: obs.fragment_count, + last_updated: obs.last_updated, + pending_wal_generations: obs.pending_wal_generations, + last_compaction, + total_compactions, + scanned_at: Utc::now().timestamp_millis(), + }) +} + +/// Spawn the periodic scanner. Returns `None` when the interval is `0`. +pub fn spawn_scanner(state: &Arc) -> Option> { + let interval_secs = state.config.stats_scan_interval_secs; + if interval_secs == 0 { + return None; + } + let state = state.clone(); + Some(tokio::spawn(async move { + let mut ticker = tokio::time::interval(Duration::from_secs(interval_secs)); + loop { + ticker.tick().await; + match scan_once(&state).await { + Ok(n) => tracing::info!(experiments = n, "stats scan complete"), + Err(e) => tracing::warn!(error = %e, "stats scan round failed"), + } + } + })) +} diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs new file mode 100644 index 00000000..56ccf0e9 --- /dev/null +++ b/crates/lance-context-master/src/state.rs @@ -0,0 +1,59 @@ +//! Shared master state: durable registry (read-only) + stats table (read/write). + +use std::path::PathBuf; +use std::sync::Arc; + +use lance_context_core::RolloutRegistry; +use tokio::sync::{Mutex, RwLock}; + +use crate::config::MasterConfig; +use crate::stats_store::StatsStore; + +/// Shared state for the master process. +/// +/// The `registry` is written only by the data-plane (store create/delete); the +/// master reads it to enumerate experiments. The `stats_store` is owned and +/// written exclusively by the master. Both are wrapped in locks because their +/// mutating methods take `&mut self`. +pub struct MasterState { + /// Durable directory of which rollout stores exist (data-plane-owned). + pub registry: RwLock, + /// Periodically-refreshed per-experiment metrics (master-owned). + pub stats: Mutex, + /// Shared data directory / object-store prefix. + pub base_path: PathBuf, + /// Effective configuration. + pub config: MasterConfig, +} + +impl MasterState { + /// Open (or create) the registry and stats datasets under `data_dir`. + pub async fn new(config: MasterConfig) -> lance::Result> { + let base_path = PathBuf::from(&config.data_dir); + let registry_uri = base_path + .join("_registry.rollout.lance") + .to_string_lossy() + .to_string(); + let stats_uri = base_path + .join("_stats.rollout.lance") + .to_string_lossy() + .to_string(); + let registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; + let stats = StatsStore::open_or_create(&stats_uri, None).await?; + Ok(Arc::new(Self { + registry: RwLock::new(registry), + stats: Mutex::new(stats), + base_path, + config, + })) + } + + /// Physical rollout dataset URI for `name`, matching the data-plane's + /// `rollout_uri` convention (`{name}.rollout.lance`). + pub fn rollout_uri(&self, name: &str) -> String { + self.base_path + .join(format!("{}.rollout.lance", name)) + .to_string_lossy() + .to_string() + } +} diff --git a/crates/lance-context-master/src/stats_store.rs b/crates/lance-context-master/src/stats_store.rs new file mode 100644 index 00000000..a1238bf2 --- /dev/null +++ b/crates/lance-context-master/src/stats_store.rs @@ -0,0 +1,399 @@ +//! Persistent stats table for the control-plane master. +//! +//! Mirrors [`lance_context_core::RolloutRegistry`] in shape and idioms: a single +//! small Lance dataset (`_stats.rollout.lance`), one row per experiment, written +//! delete-then-append for idempotency. Where the registry answers *which* stores +//! exist, this table caches the *periodically-scanned metrics* of each store so +//! the UI list endpoint can render row/fragment counts without opening every +//! dataset on each request. +//! +//! The master is the sole writer of this table; the data-plane never touches it. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator, StringArray}; +use arrow_schema::{ArrowError, DataType, Field, Schema}; +use futures::TryStreamExt; +use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; +use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; +use lance::{Error as LanceError, Result as LanceResult}; + +use lance_context_api::ExperimentSummary; + +/// One row of scanned experiment stats. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StatRow { + pub name: String, + pub uri: String, + pub row_count: i64, + pub fragment_count: i64, + pub last_updated: i64, + pub pending_wal_generations: i64, + /// Unix ms of the last successful compaction, or a negative sentinel + /// (`-1`) meaning "never", since the column is stored non-nullable. + pub last_compaction: i64, + pub total_compactions: i64, + pub scanned_at: i64, +} + +impl StatRow { + /// Sentinel stored in `last_compaction` when no compaction has run. + pub const NO_COMPACTION: i64 = -1; + + /// Convert to the wire DTO, mapping the `-1` sentinel back to `None`. + #[must_use] + pub fn into_summary(self) -> ExperimentSummary { + ExperimentSummary { + name: self.name, + uri: self.uri, + row_count: self.row_count, + fragment_count: self.fragment_count, + last_updated: self.last_updated, + pending_wal_generations: self.pending_wal_generations, + last_compaction: (self.last_compaction >= 0).then_some(self.last_compaction), + total_compactions: self.total_compactions, + scanned_at: self.scanned_at, + } + } +} + +/// Durable directory of per-experiment stats, backed by a single Lance dataset. +/// +/// Mutations (`upsert`, `remove`) take `&mut self` and must be serialized by the +/// caller (the master wraps this in a lock). Reads take `&self`. +pub struct StatsStore { + dataset: Dataset, + uri: String, + storage_options: Option>, +} + +fn stats_schema() -> Schema { + Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("uri", DataType::Utf8, false), + Field::new("row_count", DataType::Int64, false), + Field::new("fragment_count", DataType::Int64, false), + Field::new("last_updated", DataType::Int64, false), + Field::new("pending_wal_generations", DataType::Int64, false), + Field::new("last_compaction", DataType::Int64, false), + Field::new("total_compactions", DataType::Int64, false), + Field::new("scanned_at", DataType::Int64, false), + ]) +} + +impl StatsStore { + /// Open the stats dataset at `uri`, creating an empty one if absent. + pub async fn open_or_create( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + let dataset = match Self::load(uri, storage_options.clone()).await { + Ok(dataset) => dataset, + Err(LanceError::DatasetNotFound { .. }) => { + Self::create(uri, storage_options.clone()).await? + } + Err(err) => return Err(err), + }; + Ok(Self { + dataset, + uri: uri.to_string(), + storage_options, + }) + } + + async fn load( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + if let Some(options) = storage_options { + DatasetBuilder::from_uri(uri) + .with_storage_options(options) + .load() + .await + } else { + Dataset::open(uri).await + } + } + + async fn create( + uri: &str, + storage_options: Option>, + ) -> LanceResult { + let schema = Arc::new(stats_schema()); + let empty = RecordBatch::new_empty(schema.clone()); + let batches = RecordBatchIterator::new( + vec![Ok::(empty)].into_iter(), + schema.clone(), + ); + let params = Self::write_params(WriteMode::Create, storage_options); + Dataset::write(batches, uri, Some(params)).await + } + + fn write_params( + mode: WriteMode, + storage_options: Option>, + ) -> WriteParams { + let mut params = WriteParams { + mode, + ..Default::default() + }; + if let Some(options) = storage_options { + params.store_params = Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_static_options(options), + )), + ..Default::default() + }); + } + params + } + + fn row_to_batch(row: &StatRow) -> LanceResult { + let schema = Arc::new(stats_schema()); + Ok(RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec![row.name.as_str()])), + Arc::new(StringArray::from(vec![row.uri.as_str()])), + Arc::new(Int64Array::from(vec![row.row_count])), + Arc::new(Int64Array::from(vec![row.fragment_count])), + Arc::new(Int64Array::from(vec![row.last_updated])), + Arc::new(Int64Array::from(vec![row.pending_wal_generations])), + Arc::new(Int64Array::from(vec![row.last_compaction])), + Arc::new(Int64Array::from(vec![row.total_compactions])), + Arc::new(Int64Array::from(vec![row.scanned_at])), + ], + )?) + } + + /// Insert or replace the stats row for `row.name` (delete-then-append, + /// idempotent). Callers must serialize mutations. + pub async fn upsert(&mut self, row: &StatRow) -> LanceResult<()> { + self.delete_row(&row.name).await?; + let batch = Self::row_to_batch(row)?; + let schema = Arc::new(stats_schema()); + let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + let params = Self::write_params(WriteMode::Append, self.storage_options.clone()); + self.dataset.append(reader, Some(params)).await?; + Ok(()) + } + + /// Remove the stats row for `name`, if present. No-op when absent. + pub async fn remove(&mut self, name: &str) -> LanceResult<()> { + self.delete_row(name).await + } + + async fn delete_row(&mut self, name: &str) -> LanceResult<()> { + let escaped = name.replace('\'', "''"); + self.dataset + .delete(&format!("name = '{}'", escaped)) + .await?; + Ok(()) + } + + /// Fetch a single experiment's stats row by name. + pub async fn get(&self, name: &str) -> LanceResult> { + let escaped = name.replace('\'', "''"); + let mut scanner = self.dataset.scan(); + scanner.filter(&format!("name = '{}'", escaped))?; + scanner.limit(Some(1), None)?; + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + let rows = Self::batch_to_rows(&batch)?; + if let Some(first) = rows.into_iter().next() { + return Ok(Some(first)); + } + } + Ok(None) + } + + /// Total number of rows matching an optional case-sensitive substring on + /// `name`, ignoring pagination. + pub async fn count(&self, search: Option<&str>) -> LanceResult { + let mut scanner = self.dataset.scan(); + scanner.project(&["name"])?; + if let Some(q) = search { + scanner.filter(&Self::like_filter(q))?; + } + let mut stream = scanner.try_into_stream().await?; + let mut total = 0i64; + while let Some(batch) = stream.try_next().await? { + total += batch.num_rows() as i64; + } + Ok(total) + } + + /// List stats rows, optionally filtered by a `name` substring, ordered by + /// name, with `limit`/`offset` pagination applied in memory. + pub async fn list( + &self, + search: Option<&str>, + limit: usize, + offset: usize, + ) -> LanceResult> { + let mut scanner = self.dataset.scan(); + if let Some(q) = search { + scanner.filter(&Self::like_filter(q))?; + } + let mut stream = scanner.try_into_stream().await?; + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(Self::batch_to_rows(&batch)?); + } + rows.sort_by(|a, b| a.name.cmp(&b.name)); + Ok(rows.into_iter().skip(offset).take(limit).collect()) + } + + fn like_filter(search: &str) -> String { + let escaped = search.replace('\'', "''"); + format!("name LIKE '%{}%'", escaped) + } + + fn batch_to_rows(batch: &RecordBatch) -> LanceResult> { + let names = Self::str_col(batch, "name")?; + let uris = Self::str_col(batch, "uri")?; + let row_count = Self::i64_col(batch, "row_count")?; + let fragment_count = Self::i64_col(batch, "fragment_count")?; + let last_updated = Self::i64_col(batch, "last_updated")?; + let pending = Self::i64_col(batch, "pending_wal_generations")?; + let last_compaction = Self::i64_col(batch, "last_compaction")?; + let total_compactions = Self::i64_col(batch, "total_compactions")?; + let scanned_at = Self::i64_col(batch, "scanned_at")?; + let mut out = Vec::with_capacity(batch.num_rows()); + for i in 0..batch.num_rows() { + out.push(StatRow { + name: names.value(i).to_string(), + uri: uris.value(i).to_string(), + row_count: row_count.value(i), + fragment_count: fragment_count.value(i), + last_updated: last_updated.value(i), + pending_wal_generations: pending.value(i), + last_compaction: last_compaction.value(i), + total_compactions: total_compactions.value(i), + scanned_at: scanned_at.value(i), + }); + } + Ok(out) + } + + fn str_col<'a>(batch: &'a RecordBatch, name: &str) -> LanceResult<&'a StringArray> { + batch + .column_by_name(name) + .and_then(|c| c.as_any().downcast_ref::()) + .ok_or_else(|| { + LanceError::from(ArrowError::InvalidArgumentError(format!( + "stats column '{}' missing or not Utf8", + name + ))) + }) + } + + fn i64_col<'a>(batch: &'a RecordBatch, name: &str) -> LanceResult<&'a Int64Array> { + batch + .column_by_name(name) + .and_then(|c| c.as_any().downcast_ref::()) + .ok_or_else(|| { + LanceError::from(ArrowError::InvalidArgumentError(format!( + "stats column '{}' missing or not Int64", + name + ))) + }) + } + + /// The stats dataset URI. + #[must_use] + pub fn uri(&self) -> &str { + &self.uri + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + fn sample(name: &str, rows: i64) -> StatRow { + StatRow { + name: name.to_string(), + uri: format!("/data/{name}.rollout.lance"), + row_count: rows, + fragment_count: 3, + last_updated: 1_700_000_000_000, + pending_wal_generations: 0, + last_compaction: StatRow::NO_COMPACTION, + total_compactions: 0, + scanned_at: 1_700_000_001_000, + } + } + + async fn new_store(dir: &TempDir) -> StatsStore { + let uri = dir.path().join("_stats.rollout.lance"); + StatsStore::open_or_create(uri.to_str().unwrap(), None) + .await + .unwrap() + } + + #[tokio::test] + async fn upsert_is_idempotent() { + let dir = TempDir::new().unwrap(); + let mut s = new_store(&dir).await; + s.upsert(&sample("exp-a", 10)).await.unwrap(); + s.upsert(&sample("exp-a", 20)).await.unwrap(); + let rows = s.list(None, 100, 0).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].row_count, 20); + } + + #[tokio::test] + async fn get_and_remove() { + let dir = TempDir::new().unwrap(); + let mut s = new_store(&dir).await; + assert!(s.get("missing").await.unwrap().is_none()); + s.upsert(&sample("exp-b", 5)).await.unwrap(); + assert_eq!(s.get("exp-b").await.unwrap().unwrap().row_count, 5); + s.remove("exp-b").await.unwrap(); + assert!(s.get("exp-b").await.unwrap().is_none()); + } + + #[tokio::test] + async fn search_like_and_pagination() { + let dir = TempDir::new().unwrap(); + let mut s = new_store(&dir).await; + for i in 0..5 { + s.upsert(&sample(&format!("alpha-{i}"), i)).await.unwrap(); + } + s.upsert(&sample("beta-0", 99)).await.unwrap(); + + // LIKE substring narrows to the alpha family. + assert_eq!(s.count(Some("alpha")).await.unwrap(), 5); + let page = s.list(Some("alpha"), 2, 0).await.unwrap(); + assert_eq!(page.len(), 2); + assert_eq!(page[0].name, "alpha-0"); + let page2 = s.list(Some("alpha"), 2, 2).await.unwrap(); + assert_eq!(page2[0].name, "alpha-2"); + + // Full listing plus the beta row. + assert_eq!(s.count(None).await.unwrap(), 6); + } + + #[tokio::test] + async fn survives_reopen() { + let dir = TempDir::new().unwrap(); + { + let mut s = new_store(&dir).await; + s.upsert(&sample("persist", 7)).await.unwrap(); + } + let s = new_store(&dir).await; + assert_eq!(s.get("persist").await.unwrap().unwrap().row_count, 7); + } + + #[tokio::test] + async fn no_compaction_sentinel_maps_to_none() { + let row = sample("x", 1); + assert_eq!(row.into_summary().last_compaction, None); + let mut with = sample("y", 1); + with.last_compaction = 123; + assert_eq!(with.into_summary().last_compaction, Some(123)); + } +}