diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index 50bf5bf..36264af 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -27,6 +27,18 @@ pub struct MasterConfig { #[arg(long, env = "SCAN_CONCURRENCY", default_value_t = 8)] pub scan_concurrency: usize, + /// Total byte budget for the Lance metadata/index caches, shared across + /// every rollout store opened by this master process. + /// + /// Without an explicit shared session, Lance gives each store a fresh + /// session with caches defaulting to 6 GiB index + 1 GiB metadata. The + /// master opens stores for scans, browsing, compaction, and indexing, so + /// per-store sessions make RSS grow with the experiments touched. This + /// budget is split internally 6:1 (index:metadata), matching Lance's + /// default ratio. `0` restores Lance's per-store default sessions. + #[arg(long, env = "ROLLOUT_CACHE_BYTES", default_value = "2147483648")] + pub rollout_cache_bytes: usize, + /// Run maintenance (compaction + old-version cleanup) on the `_stats` /// dataset every Nth stats-scan round. `_stats` is written delete-then- /// append, so each scan adds versions and fragments per experiment; without @@ -153,3 +165,50 @@ pub struct MasterConfig { #[arg(long, env = "UI_DIR")] pub ui_dir: Option, } + +#[cfg(test)] +mod tests { + use super::*; + use clap::CommandFactory; + + #[test] + fn every_config_field_is_an_optional_flag() { + let command = MasterConfig::command(); + + let positionals: Vec<_> = command + .get_positionals() + .map(|arg| arg.get_id().to_string()) + .collect(); + assert!( + positionals.is_empty(), + "config fields must be flags, not positional args; found {positionals:?}" + ); + + for arg in command.get_arguments() { + if arg.get_id() == "help" || arg.get_id() == "version" { + continue; + } + assert!( + arg.get_long().is_some(), + "'{}' has no long flag", + arg.get_id() + ); + assert!( + arg.get_default_values().len() == 1 || !arg.is_required_set(), + "'{}' must have a default or be optional", + arg.get_id() + ); + } + } + + #[test] + fn rollout_cache_defaults_to_two_gib_and_can_be_disabled() { + let default = MasterConfig::try_parse_from(["lance-context-master"]).unwrap(); + assert_eq!(default.rollout_cache_bytes, 2 * 1024 * 1024 * 1024); + + let disabled = + MasterConfig::try_parse_from(["lance-context-master", "--rollout-cache-bytes", "0"]) + .unwrap(); + assert_eq!(disabled.rollout_cache_bytes, 0); + } +} diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 526189c..b04df17 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -165,7 +165,7 @@ pub async fn list_experiments( // thousands of datasets to render one page. let want = params.limit.saturating_sub(experiments.len()); for entry in matches.into_iter().take(want) { - match scanner::observe_cold(&entry.name, &entry.uri).await { + match scanner::observe_cold(&state, &entry.name, &entry.uri).await { Ok(summary) => experiments.push(summary), Err(e) => { tracing::warn!( @@ -225,7 +225,7 @@ pub async fn get_experiment( .await .map_err(MasterError::from_lance)? .ok_or_else(|| MasterError::NotFound(format!("experiment '{}' does not exist", name)))?; - let summary = scanner::observe_cold(&entry.name, &entry.uri) + let summary = scanner::observe_cold(&state, &entry.name, &entry.uri) .await .map_err(MasterError::from_lance)?; Ok(Json(ExperimentDetail { summary })) @@ -605,6 +605,7 @@ mod tests { port: 0, stats_scan_interval_secs: 0, scan_concurrency: 4, + rollout_cache_bytes: 2 * 1024 * 1024 * 1024, stats_maintenance_every_n_scans: 0, stats_history_ttl_secs: 3_600, stats_cold_retire_secs: 0, diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index be8de2b..f21c5b6 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/crates/lance-context-master/src/scanner.rs @@ -194,6 +194,7 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { .collect() }; let previous = Arc::new(previous); + let rollout_options = state.rollout_store_options(); // Observe experiments concurrently (bounded). `None` means this round could // not observe that experiment; its previous row is carried over below @@ -201,9 +202,10 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { let observed: Vec<(String, Option<(StatRow, bool)>)> = stream::iter(entries) .map(|entry| { let previous = previous.clone(); + let rollout_options = rollout_options.clone(); async move { let prev = previous.get(&entry.name); - match observe_one(&entry.name, &entry.uri, prev).await { + match observe_one(&entry.name, &entry.uri, prev, rollout_options).await { Ok(result) => (entry.name, Some(result)), Err(e) => { tracing::warn!(store = %entry.name, error = %e, "scan: observe failed"); @@ -255,8 +257,13 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { // Done before the snapshot is written so a retirement takes effect in the // same commit rather than leaving a round where the row is stale. let retire_after = Duration::from_secs(state.config.stats_cold_retire_secs); - let retired = - retire_cold_experiments(&snapshot, retire_after, Utc::now().timestamp_millis()).await; + let retired = retire_cold_experiments( + &snapshot, + retire_after, + Utc::now().timestamp_millis(), + state.rollout_store_options(), + ) + .await; if !retired.is_empty() { snapshot.retain(|row| !retired.contains(&row.name)); } @@ -329,9 +336,9 @@ async fn observe_one( name: &str, uri: &str, previous: Option<&StatRow>, + options: RolloutStoreOptions, ) -> lance::Result<(StatRow, bool)> { - let opts = RolloutStoreOptions::default(); - let open = RolloutStore::open_existing_with_options(uri, opts); + let open = RolloutStore::open_existing_with_options(uri, options); let store = match tokio::time::timeout(OBSERVE_TIMEOUT, open).await { Ok(Ok(store)) => store, Ok(Err(e)) => return Err(e), @@ -413,6 +420,7 @@ async fn retire_cold_experiments( rows: &[StatRow], retire_after: Duration, now_ms: i64, + options: RolloutStoreOptions, ) -> HashSet { if retire_after.is_zero() { return HashSet::new(); @@ -424,7 +432,7 @@ async fn retire_cold_experiments( if row.last_updated > cutoff_ms { continue; } - match prepare_for_retirement(&row.name, &row.uri).await { + match prepare_for_retirement(&row.name, &row.uri, options.clone()).await { Ok(true) => { retired.insert(row.name.clone()); metrics::counter!("master_stats_experiments_retired_total").increment(1); @@ -458,11 +466,14 @@ async fn retire_cold_experiments( /// Merge, compact, and verify one experiment is quiescent. /// /// `Ok(true)` means it is safe to drop from the stats table. -async fn prepare_for_retirement(name: &str, uri: &str) -> lance::Result { - let opts = RolloutStoreOptions::default(); +async fn prepare_for_retirement( + name: &str, + uri: &str, + options: RolloutStoreOptions, +) -> lance::Result { let mut store = match tokio::time::timeout( OBSERVE_TIMEOUT, - RolloutStore::open_existing_with_options(uri, opts), + RolloutStore::open_existing_with_options(uri, options), ) .await { @@ -502,8 +513,20 @@ async fn prepare_for_retirement(name: &str, uri: &str) -> lance::Result { /// reading about a cold experiment must not make it hot again, or browsing the /// UI would undo retirement and the table would creep back toward holding /// everything. -pub async fn observe_cold(name: &str, uri: &str) -> lance::Result { - let (row, _) = observe_one(name, uri, None).await?; +pub async fn observe_cold( + state: &MasterState, + name: &str, + uri: &str, +) -> lance::Result { + observe_cold_with_options(name, uri, state.rollout_store_options()).await +} + +async fn observe_cold_with_options( + name: &str, + uri: &str, + options: RolloutStoreOptions, +) -> lance::Result { + let (row, _) = observe_one(name, uri, None, options).await?; Ok(row.into_summary()) } @@ -515,7 +538,7 @@ pub async fn observe_cold(name: &str, uri: &str) -> lance::Result, name: &str, uri: &str) -> lance::Result<()> { let guard = state.task_store.coordination_lock("stats-writer").await?; - let result = match observe_one(name, uri, None).await { + let result = match observe_one(name, uri, None, state.rollout_store_options()).await { Ok((row, _)) => state.stats.lock().await.upsert(&row).await, Err(error) => Err(error), }; @@ -691,13 +714,18 @@ mod incremental_scan_tests { } // First pass: nothing known, so a full observation happens. - let (first, skipped) = observe_one("e", &uri, None).await.unwrap(); + let (first, skipped) = observe_one("e", &uri, None, RolloutStoreOptions::default()) + .await + .unwrap(); assert!(!skipped, "the first observation cannot be skipped"); assert_ne!(first.version, StatRow::UNKNOWN_VERSION); // Second pass with the row from the first: the version is unchanged, so // the expensive observation is skipped and the row is reused. - let (second, skipped) = observe_one("e", &uri, Some(&first)).await.unwrap(); + let (second, skipped) = + observe_one("e", &uri, Some(&first), RolloutStoreOptions::default()) + .await + .unwrap(); assert!(skipped, "an unchanged experiment must skip re-observation"); assert_eq!(second.row_count, first.row_count); assert_eq!(second.fragment_count, first.fragment_count); @@ -720,12 +748,17 @@ mod incremental_scan_tests { store.add(&[rec("a")]).await.unwrap(); store.flush().await.unwrap(); - let (first, _) = observe_one("e", &uri, None).await.unwrap(); + let (first, _) = observe_one("e", &uri, None, RolloutStoreOptions::default()) + .await + .unwrap(); // Merge the WAL into the base table: this advances the base version. store.cleanup_own_shard().await.unwrap(); - let (second, skipped) = observe_one("e", &uri, Some(&first)).await.unwrap(); + let (second, skipped) = + observe_one("e", &uri, Some(&first), RolloutStoreOptions::default()) + .await + .unwrap(); assert!( !skipped, "a changed base version must force a full observation" @@ -750,10 +783,15 @@ mod incremental_scan_tests { store.flush().await.unwrap(); } - let (mut legacy, _) = observe_one("e", &uri, None).await.unwrap(); + let (mut legacy, _) = observe_one("e", &uri, None, RolloutStoreOptions::default()) + .await + .unwrap(); legacy.version = StatRow::UNKNOWN_VERSION; - let (refreshed, skipped) = observe_one("e", &uri, Some(&legacy)).await.unwrap(); + let (refreshed, skipped) = + observe_one("e", &uri, Some(&legacy), RolloutStoreOptions::default()) + .await + .unwrap(); assert!( !skipped, "an unknown version must not be mistaken for an unchanged one" @@ -774,11 +812,15 @@ mod incremental_scan_tests { store.flush().await.unwrap(); } - let (mut prev, _) = observe_one("e", &uri, None).await.unwrap(); + let (mut prev, _) = observe_one("e", &uri, None, RolloutStoreOptions::default()) + .await + .unwrap(); prev.last_compaction = 1_700_000_000_000; prev.total_compactions = 7; - let (next, skipped) = observe_one("e", &uri, Some(&prev)).await.unwrap(); + let (next, skipped) = observe_one("e", &uri, Some(&prev), RolloutStoreOptions::default()) + .await + .unwrap(); assert!(skipped); assert_eq!(next.last_compaction, 1_700_000_000_000); assert_eq!(next.total_compactions, 7); @@ -875,9 +917,13 @@ mod retirement_tests { let now = Utc::now().timestamp_millis(); let old = now - Duration::from_secs(30 * 86_400).as_millis() as i64; - let retired = - retire_cold_experiments(&[row("e", &uri, old)], Duration::from_secs(7 * 86_400), now) - .await; + let retired = retire_cold_experiments( + &[row("e", &uri, old)], + Duration::from_secs(7 * 86_400), + now, + RolloutStoreOptions::default(), + ) + .await; assert!(retired.contains("e"), "a cold experiment should retire"); @@ -908,9 +954,13 @@ mod retirement_tests { } let now = Utc::now().timestamp_millis(); - let retired = - retire_cold_experiments(&[row("e", &uri, now)], Duration::from_secs(7 * 86_400), now) - .await; + let retired = retire_cold_experiments( + &[row("e", &uri, now)], + Duration::from_secs(7 * 86_400), + now, + RolloutStoreOptions::default(), + ) + .await; assert!( retired.is_empty(), "a hot experiment must stay in the table" @@ -921,9 +971,13 @@ mod retirement_tests { #[tokio::test] async fn zero_window_disables_retirement() { let now = Utc::now().timestamp_millis(); - let retired = - retire_cold_experiments(&[row("e", "/nonexistent", 0)], Duration::from_secs(0), now) - .await; + let retired = retire_cold_experiments( + &[row("e", "/nonexistent", 0)], + Duration::from_secs(0), + now, + RolloutStoreOptions::default(), + ) + .await; assert!(retired.is_empty()); } @@ -937,6 +991,7 @@ mod retirement_tests { &[row("gone", "/no/such/dataset.lance", old)], Duration::from_secs(7 * 86_400), now, + RolloutStoreOptions::default(), ) .await; assert!( @@ -963,17 +1018,18 @@ mod retirement_tests { store.flush().await.unwrap(); } - let summary = observe_cold("e", &uri).await.unwrap(); + let summary = observe_cold_with_options("e", &uri, RolloutStoreOptions::default()) + .await + .unwrap(); assert_eq!(summary.name, "e"); assert_eq!( summary.row_count, 1, "a cold read still reports real counts" ); - // `observe_cold` takes no `MasterState`, so it structurally cannot - // write to the stats table -- asserted here so a future refactor that - // hands it one has to justify itself. - let second = observe_cold("e", &uri).await.unwrap(); + let second = observe_cold_with_options("e", &uri, RolloutStoreOptions::default()) + .await + .unwrap(); assert_eq!(second.row_count, summary.row_count); } @@ -981,8 +1037,12 @@ mod retirement_tests { /// hit on a registry entry whose data is gone degrades to one skipped row. #[tokio::test] async fn observe_cold_errors_on_missing_dataset() { - assert!(observe_cold("gone", "/no/such/dataset.lance") - .await - .is_err()); + assert!(observe_cold_with_options( + "gone", + "/no/such/dataset.lance", + RolloutStoreOptions::default(), + ) + .await + .is_err()); } } diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 5bf2117..07e7b67 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -28,7 +28,7 @@ use std::time::Duration; use chrono::Utc; use lance_context_api::{TaskKind, TaskRecord}; -use lance_context_core::{CompactionConfig, RolloutStore, RolloutStoreOptions}; +use lance_context_core::{CompactionConfig, RolloutStore}; use tokio::sync::Semaphore; use tokio::task::JoinHandle; @@ -165,8 +165,7 @@ async fn run_index_id(state: &Arc, name: &str) -> Result, name: &str) -> Result { let uri = state.rollout_uri(name); - let opts = RolloutStoreOptions::default(); - let mut store = RolloutStore::open_existing_with_options(&uri, opts) + let mut store = RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) .await .map_err(|e| e.to_string())?; store @@ -179,9 +178,8 @@ async fn index_id_inner(state: &Arc, name: &str) -> Result, name: &str) -> Result { let uri = state.rollout_uri(name); let config = compaction_config(state); - let opts = RolloutStoreOptions::default(); - let mut store = RolloutStore::open_existing_with_options(&uri, opts) + let mut store = RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) .await .map_err(|e| e.to_string())?; let metrics = store @@ -553,6 +551,7 @@ mod tests { port: 0, stats_scan_interval_secs: 0, scan_concurrency: 4, + rollout_cache_bytes: 2 * 1024 * 1024 * 1024, stats_maintenance_every_n_scans: 0, stats_history_ttl_secs: 3_600, stats_cold_retire_secs: 0, diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index 9d8c772..bfa48ed 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -3,7 +3,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; -use lance_context_core::{join_uri, RolloutRegistry, RolloutStore, RolloutStoreOptions}; +use lance_context_core::{join_uri, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session}; use lru::LruCache; use tokio::sync::{Mutex, RwLock}; @@ -14,11 +14,24 @@ use crate::task_store::TaskStore; /// Bounded number of rollout handles retained by the master data browser. /// -/// Each handle keeps a Lance session whose fragment/file metadata caches are -/// valuable across pagination requests. The bound prevents a master managing a -/// very large registry from retaining one handle per experiment indefinitely. +/// Each handle reuses the process-wide Lance session whose fragment/file +/// metadata caches are valuable across pagination requests. The bound prevents +/// a master managing a very large registry from retaining one handle per +/// experiment indefinitely. const RECORD_STORE_CACHE_CAPACITY: usize = 128; +/// Build the process-wide rollout session from one total cache budget. +/// +/// The 6:1 index:metadata split mirrors Lance's default cache ratio. +fn build_rollout_session(cache_bytes: usize) -> Option> { + if cache_bytes == 0 { + return None; + } + let metadata_bytes = cache_bytes / 7; + let index_bytes = cache_bytes - metadata_bytes; + Some(RolloutStore::build_session(index_bytes, metadata_bytes)) +} + /// Shared state for the master process. /// /// The data-plane owns steady-state registry writes (store create/delete); the @@ -34,6 +47,9 @@ pub struct MasterState { /// Read handles used by the records browser, retained to reuse Lance /// session and fragment metadata caches across requests. record_stores: Mutex>>>, + /// Process-wide Lance cache session attached to every rollout store opened + /// by this master. `None` is only used when the configured budget is `0`. + rollout_session: Option>, /// Shared data directory / object-store prefix. pub base_uri: String, /// Effective configuration. @@ -76,6 +92,7 @@ impl MasterState { tokio::time::sleep(std::time::Duration::from_millis(200)).await; }; let base_uri = config.data_dir.clone(); + let rollout_session = build_rollout_session(config.rollout_cache_bytes); let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); let stats_uri = join_uri(&base_uri, "_stats.rollout.lance"); let mut registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; @@ -94,6 +111,7 @@ impl MasterState { record_stores: Mutex::new(LruCache::new( NonZeroUsize::new(RECORD_STORE_CACHE_CAPACITY).unwrap(), )), + rollout_session, base_uri, config, task_store, @@ -111,6 +129,14 @@ impl MasterState { join_uri(&self.base_uri, &format!("{}.rollout.lance", name)) } + /// Options for every rollout store opened by the master. + pub(crate) fn rollout_store_options(&self) -> RolloutStoreOptions { + RolloutStoreOptions { + session: self.rollout_session.clone(), + ..Default::default() + } + } + /// Return a cached rollout handle for the master records browser, opening /// it without holding the cache lock on a miss. pub async fn get_or_open_record_store( @@ -125,7 +151,7 @@ impl MasterState { metrics::counter!("master_record_store_cache_misses_total").increment(1); let opened = Arc::new(RwLock::new( - RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()).await?, + RolloutStore::open_existing_with_options(uri, self.rollout_store_options()).await?, )); let mut cache = self.record_stores.lock().await; @@ -152,6 +178,7 @@ mod tests { port: 0, stats_scan_interval_secs: 0, scan_concurrency: 4, + rollout_cache_bytes: 2 * 1024 * 1024 * 1024, stats_maintenance_every_n_scans: 0, stats_history_ttl_secs: 3_600, stats_cold_retire_secs: 0, @@ -178,6 +205,12 @@ mod tests { } } + #[test] + fn rollout_session_can_be_disabled() { + assert!(build_rollout_session(0).is_none()); + assert!(build_rollout_session(7).is_some()); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn startup_backfills_legacy_rollout_datasets() { diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 5109ece..822a250 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -910,6 +910,7 @@ mod tests { port: 0, stats_scan_interval_secs: 0, scan_concurrency: 4, + rollout_cache_bytes: 2 * 1024 * 1024 * 1024, stats_maintenance_every_n_scans: 0, stats_history_ttl_secs: 3_600, stats_cold_retire_secs: 0, diff --git a/deploy/kubernetes/master.yaml b/deploy/kubernetes/master.yaml index 8a9c3b3..9ea5175 100644 --- a/deploy/kubernetes/master.yaml +++ b/deploy/kubernetes/master.yaml @@ -40,6 +40,10 @@ spec: value: /lance-context/master - name: ETCD_LEASE_TTL_SECS value: "30" + # Shared process-wide Lance index/metadata cache budget. Keep this + # below the pod memory limit to leave room for query working sets. + - name: ROLLOUT_CACHE_BYTES + value: "536870912" - name: ETCD_USERNAME valueFrom: secretKeyRef: