diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 716e350..904e632 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -61,3 +61,7 @@ pub use lance::dataset::optimize::CompactionMetrics; // Re-export the Lance error type so downstream crates (e.g. the server) can // match on its typed variants instead of string-matching `Display` output. pub use lance::Error as LanceError; + +// Re-export the Lance session type so the server can build one shared, +// capacity-bounded cache session across all resident rollout stores. +pub use lance::session::Session; diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 4ee1e50..22d09f4 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -74,6 +74,7 @@ use lance::dataset::{ }; use lance::index::DatasetIndexExt; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; +use lance::session::Session; use lance::{Error as LanceError, Result as LanceResult}; use lance_index::mem_wal::{ShardManifest, MEM_WAL_INDEX_NAME}; use lance_index::scalar::ScalarIndexParams; @@ -341,6 +342,21 @@ pub struct RolloutStoreOptions { /// `None` or `0` disables self-merge (the 0.6.0 behavior: generations /// accumulate and are unioned at read time). pub merge_after_generations: Option, + /// Shared Lance [`Session`] used to open this store's base dataset (and, + /// transitively, every flushed MemWAL generation it reads — those inherit + /// the base dataset's session). + /// + /// When `None`, Lance builds a *fresh* per-store session whose index and + /// metadata caches default to **6 GiB / 1 GiB** and are keyed by dataset + /// URI. Each flushed generation is a distinct URI, so a busy store's read + /// path (`observe`, list, stats) feeds an ever-growing set of keys into that + /// cache until it approaches the 6 GiB cap — worker RSS then grows linearly + /// with cumulative appends and is never released across merge/compact cycles + /// (the session outlives the merged generations). Passing a single shared, + /// capacity-bounded session across all resident stores bounds this: the + /// deployment's total Lance cache is the session's capacity, shared, rather + /// than 6 GiB *per store*. Build one with [`RolloutStore::build_session`]. + pub session: Option>, } /// A Lance-backed store for RL rollout trajectories. @@ -370,6 +386,11 @@ pub struct RolloutStore { /// `add` transparently reopens with the fresh epoch. [`ShardWriter`] has no /// `Drop`, so it must be closed explicitly to drain its background tasks; /// see [`Self::close`] and the `Drop` impl. + /// Shared Lance session used to open the base dataset and every reload + /// (compact, zonemap, generation reads), so all resident stores share one + /// capacity-bounded cache instead of each allocating Lance's 6 GiB default. + /// `None` preserves Lance's per-open default session. + session: Option>, /// Resident MemWAL writer for this instance's shard, wrapped for `&self` /// concurrent access. The [`tokio::sync::Mutex`] is held only to fetch-or-open /// and clone the `Arc` (see [`Self::resident_writer`]) and to invalidate a @@ -380,6 +401,25 @@ pub struct RolloutStore { } impl RolloutStore { + /// Build a shared, capacity-bounded Lance [`Session`] for opening rollout + /// stores. + /// + /// Split the total cache budget across Lance's two caches: `index_cache_bytes` + /// bounds opened-index data and `metadata_cache_bytes` bounds file/dataset + /// metadata. Both are byte-weighted LRUs keyed by dataset URI, so a single + /// session shared across every resident [`RolloutStore`] caps the process's + /// *total* Lance cache at this budget — instead of Lance's default of a fresh + /// 6 GiB + 1 GiB session per store (the source of the per-append RSS growth; + /// see [`RolloutStoreOptions::session`]). + #[must_use] + pub fn build_session(index_cache_bytes: usize, metadata_cache_bytes: usize) -> Arc { + Arc::new(Session::new( + index_cache_bytes, + metadata_cache_bytes, + Arc::default(), + )) + } + /// Open an existing rollout dataset or create a new one with default /// options (`binary_payload` stored inline; see [`RolloutStoreOptions`]). /// @@ -417,14 +457,16 @@ impl RolloutStore { create_if_missing: bool, ) -> LanceResult { let storage_options = options.storage_options.clone(); + let session = options.session.clone(); let write_shard = derive_shard_id(options.shard_id.as_deref()); - let dataset = match Self::load_with_options(uri, storage_options.clone()).await { - Ok(dataset) => dataset, - Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { - Self::create_with_options(uri, storage_options.clone()).await? - } - Err(err) => return Err(err), - }; + let dataset = + match Self::load_with_options(uri, storage_options.clone(), session.clone()).await { + Ok(dataset) => dataset, + Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { + Self::create_with_options(uri, storage_options.clone(), session.clone()).await? + } + Err(err) => return Err(err), + }; let mut store = Self { dataset, @@ -434,6 +476,7 @@ impl RolloutStore { last_compaction: None, total_compactions: 0, last_compaction_error: None, + session, write_writer: tokio::sync::Mutex::new(None), }; // Initialize the MemWAL index up front (idempotent, cheap when already @@ -755,8 +798,12 @@ impl RolloutStore { "{}/_mem_wal/{}/{}", base_uri, self.write_shard, flushed.path ); - let gen_dataset = - Self::load_with_options(&gen_uri, self.storage_options.clone()).await?; + let gen_dataset = Self::load_with_options( + &gen_uri, + self.storage_options.clone(), + self.session.clone(), + ) + .await?; let mut stream = gen_dataset.scan().try_into_stream().await?; while let Some(batch) = stream.try_next().await? { if batch.num_rows() > 0 { @@ -926,7 +973,12 @@ impl RolloutStore { // Reload the handle so the caller (and subsequent reads on this // instance) observe the compacted version. let uri = self.dataset.uri().to_string(); - self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; + self.dataset = Self::load_with_options( + &uri, + self.storage_options.clone(), + self.session.clone(), + ) + .await?; self.last_compaction = Some(Utc::now()); self.total_compactions += 1; self.last_compaction_error = None; @@ -972,7 +1024,9 @@ impl RolloutStore { // Reload the handle so subsequent reads on this instance observe the // new index (mirrors the reload done after `compact`). let uri = self.dataset.uri().to_string(); - self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; + self.dataset = + Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) + .await?; Ok(()) } @@ -1089,7 +1143,12 @@ impl RolloutStore { // A concurrent first-writer may have created the index between // our check and our commit. Reload and accept it if so. let uri = self.dataset.uri().to_string(); - self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; + self.dataset = Self::load_with_options( + &uri, + self.storage_options.clone(), + self.session.clone(), + ) + .await?; if self.mem_wal_index_present().await? { Ok(()) } else { @@ -1749,20 +1808,25 @@ impl RolloutStore { async fn load_with_options( uri: &str, storage_options: Option>, + session: Option>, ) -> LanceResult { + // A plain `Dataset::open` builds a fresh per-store session (Lance's 6 GiB + // index + 1 GiB metadata default). Route every open through the builder so + // an explicit shared, capacity-bounded session can be attached. + let mut builder = DatasetBuilder::from_uri(uri); if let Some(options) = storage_options { - DatasetBuilder::from_uri(uri) - .with_storage_options(options) - .load() - .await - } else { - Dataset::open(uri).await + builder = builder.with_storage_options(options); + } + if let Some(session) = session { + builder = builder.with_session(session); } + builder.load().await } async fn create_with_options( uri: &str, storage_options: Option>, + session: Option>, ) -> LanceResult { let schema = Arc::new(rollout_schema()); let empty_batch = RecordBatch::new_empty(schema.clone()); @@ -1783,6 +1847,7 @@ impl RolloutStore { ..Default::default() }); } + params.session = session; Dataset::write(batches, uri, Some(params)).await } @@ -3010,6 +3075,7 @@ mod tests { let shard_b_blob = b"shard-b-blob"; let options = |shard: &str| RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some(shard.to_string()), merge_after_generations: None, }; @@ -3289,6 +3355,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // no merge → epoch never reclaimed }, @@ -3342,6 +3409,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), // Merge on every append → epoch is reclaimed each time, so // the following append always hits the reopen path. @@ -3380,6 +3448,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, }, @@ -3512,6 +3581,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), // Merge every append into base so each forms its own fragment. merge_after_generations: Some(1), @@ -3581,6 +3651,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), }, @@ -3625,6 +3696,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), }, @@ -3642,6 +3714,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-1".to_string()), ..Default::default() }, @@ -3690,6 +3763,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), }, @@ -3724,6 +3798,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), }, @@ -3779,6 +3854,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(2), }, @@ -3817,6 +3893,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // count trigger off }, @@ -4019,6 +4096,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some(shard), ..Default::default() }, @@ -4091,6 +4169,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), ..Default::default() }, @@ -4140,6 +4219,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), ..Default::default() }, @@ -4349,6 +4429,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // disabled }, @@ -4378,6 +4459,7 @@ mod tests { &uri_m, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), }, @@ -4429,6 +4511,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), ..Default::default() }, @@ -4511,6 +4594,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), ..Default::default() }, @@ -4549,6 +4633,7 @@ mod tests { &uri, RolloutStoreOptions { storage_options: None, + session: None, shard_id: Some("rollout-0".to_string()), ..Default::default() }, diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 10e6bd3..4078139 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -82,6 +82,23 @@ pub struct ServerConfig { /// `0` (the default) disables the budget. #[arg(long, env = "ROLLOUT_MAX_INFLIGHT_BLOB_BYTES", default_value = "0")] pub rollout_max_inflight_blob_bytes: usize, + + /// Total byte budget for the Lance metadata/index caches, shared across + /// **all** resident rollout stores on this instance. + /// + /// Every rollout store opens a Lance dataset; without an explicit session + /// Lance gives each store its own caches defaulting to **6 GiB index + 1 GiB + /// metadata**, keyed by dataset URI. Because each flushed MemWAL generation + /// is a distinct URI, a busy store's read path feeds an ever-growing key set + /// into that per-store cache until it approaches 6 GiB — worker RSS then + /// grows linearly with cumulative appends and never releases across + /// merge/compact cycles. Attaching one shared, capacity-bounded session caps + /// the process's *total* Lance cache at this budget instead of 6 GiB per + /// store. Split internally 6:1 (index:metadata), matching Lance's own + /// default ratio. `0` disables sharing and restores Lance's per-store + /// default session (the pre-fix, leak-prone behavior). + #[arg(long, env = "ROLLOUT_CACHE_BYTES", default_value = "2147483648")] + pub rollout_cache_bytes: usize, } impl ServerConfig { diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 849fdd6..0becf3a 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -125,6 +125,7 @@ pub async fn create_rollout_store( shard_id: state.instance_id.clone(), merge_after_generations: (state.rollout_merge_after_generations > 0) .then_some(state.rollout_merge_after_generations), + session: state.rollout_session.clone(), }; let store = RolloutStore::open_with_options(&uri, options) diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 272b4c4..e844159 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -6,7 +6,7 @@ use std::time::Duration; use lance_context_core::{ join_uri, validate_store_name, ContextStore, ContextStoreOptions, RolloutRegistry, - RolloutStore, RolloutStoreOptions, + RolloutStore, RolloutStoreOptions, Session, }; use lru::LruCache; use tokio::sync::{Mutex, RwLock}; @@ -55,6 +55,11 @@ pub struct AppState { /// uploads/downloads. `None` disables the budget (unbounded). See /// [`BlobBudget`]. pub blob_budget: Option>, + /// Shared Lance cache session attached to every resident rollout store, so + /// the process's total metadata/index cache is bounded by one budget rather + /// than Lance's default 6 GiB *per store*. `None` restores the per-store + /// default session (leak-prone; only when the budget is configured to `0`). + pub rollout_session: Option>, } /// Process-wide admission control for the total artifact-blob payload held in @@ -136,6 +141,21 @@ impl Drop for BlobReservation { } } +/// Build the process-wide shared Lance cache session for rollout stores from a +/// single total byte budget, or `None` when the budget is `0` (which restores +/// Lance's per-store default session). +/// +/// The budget is split 6:1 index:metadata to mirror Lance's own default ratio +/// (`DEFAULT_INDEX_CACHE_SIZE` = 6 GiB, `DEFAULT_METADATA_CACHE_SIZE` = 1 GiB). +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)) +} + impl AppState { /// Build the shared server state, opening (or creating) the rollout registry /// under `data_dir`. Async because opening the registry touches storage. @@ -150,6 +170,7 @@ impl AppState { .unwrap_or_else(|| NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap()); let blob_budget = (config.rollout_max_inflight_blob_bytes > 0) .then(|| BlobBudget::new(config.rollout_max_inflight_blob_bytes)); + let rollout_session = build_rollout_session(config.rollout_cache_bytes); Ok(Self { stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), @@ -160,6 +181,7 @@ impl AppState { rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, rollout_flush_interval_secs: config.rollout_flush_interval_secs, blob_budget, + rollout_session, }) } @@ -198,6 +220,7 @@ impl AppState { rollout_cleanup_interval_secs: 0, rollout_flush_interval_secs: 0, blob_budget: None, + rollout_session: build_rollout_session(2 * 1024 * 1024 * 1024), } } @@ -216,6 +239,7 @@ impl AppState { shard_id: self.instance_id.clone(), merge_after_generations: (self.rollout_merge_after_generations > 0) .then_some(self.rollout_merge_after_generations), + session: self.rollout_session.clone(), } } diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs index 553e4b5..a8bcc4a 100644 --- a/crates/lance-context/src/unified_rollout.rs +++ b/crates/lance-context/src/unified_rollout.rs @@ -38,6 +38,7 @@ impl RolloutStore { // caller that wants count-triggered self-merge opens the core // `RolloutStore` directly with `merge_after_generations` set. merge_after_generations: None, + ..Default::default() }; let store = LocalStore::open_with_options(uri, options) .await