From 7a274d735605795ed23bc81fff8d320afb18a247 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Thu, 9 Jul 2026 08:08:31 +0000 Subject: [PATCH 1/2] feat(rollout): MemWAL ingest with server-id sharding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Convert RolloutStore from atomic Dataset::append to Lance MemWAL, sharded by stable server instance id. Each REST server instance owns exactly one MemWAL shard (shard_uuid = uuid_v5(NAMESPACE_OID, instance_id)), so the single-active-writer + epoch-fencing invariant holds by construction and a plain round-robin load balancer is correct — no consistent hashing, lock service, or client changes. close-per-append flushes each write to object storage before returning, and the LSM read path rebuilds from every shard's manifests, so any instance sees every instance's writes (read-your-writes with no affinity). Store binary_payload inline instead of blob-v2 offloaded. The MemWAL LSM read path is the only read path for rollouts and has no blob-materialization step, so an offloaded (lance-encoding:blob) column reads back as None and get_blob / the blob endpoint would return nothing. Inline LargeBinary round-trips correctly; list/get_by_id still project the column out so hot-path scans never materialize artifact bytes (learner doesn't pay for artifacts), and get_blob projects it in on demand. Removes the broken blob_columns knob from the rollout options, server handler, unified facade, and CreateRolloutStoreRequest. Thread instance id through the server: --instance-id / INSTANCE_ID, falling back to HOSTNAME (StatefulSet pod ordinal), into AppState and the rollout shard key. Reproducibility is now a filter over immutable append-only rows (e.g. policy_version), not a per-add checkout snapshot — MemWAL appends do not advance the base dataset version. Add specs/rollout-deployment.md covering the K8s StatefulSet deployment, the sharding model, and the single-writer/read-your-writes/append-only rationale. Co-Authored-By: Claude Opus 4 --- crates/lance-context-api/src/lib.rs | 2 - .../lance-context-core/src/rollout_store.rs | 387 ++++++++++++------ crates/lance-context-server/Cargo.toml | 2 +- crates/lance-context-server/src/config.rs | 24 ++ .../src/routes/records.rs | 1 + .../src/routes/rollouts.rs | 20 +- .../lance-context-server/src/routes/search.rs | 1 + crates/lance-context-server/src/state.rs | 6 + crates/lance-context/src/unified_rollout.rs | 14 +- specs/rollout-deployment.md | 196 +++++++++ 10 files changed, 512 insertions(+), 141 deletions(-) create mode 100644 specs/rollout-deployment.md diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index b43bbbe..158e6f4 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -574,8 +574,6 @@ pub struct CreateRolloutStoreRequest { pub name: String, #[serde(default)] pub storage_options: Option>, - #[serde(default)] - pub blob_columns: Option>, } #[derive(Debug, Serialize, Deserialize)] diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index ff37943..b96ac42 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -3,16 +3,48 @@ //! [`RolloutStore`] is a second, independent first-class store alongside //! [`crate::store::ContextStore`]. It owns its own Lance dataset and Arrow //! schema (see [`crate::rollout::RolloutRecord`] and spec §5) but reuses the -//! schema-agnostic infrastructure: native dataset versioning, blob v2 payload -//! offload, and the relationship graph. +//! schema-agnostic infrastructure: MemWAL ingest and the relationship graph. //! -//! Unlike `ContextStore`, the v1 write path is a plain atomic -//! [`Dataset::append`] — no MemWAL region grouping (that machinery is specific -//! to conversation `bot_id`/`session_id` sharding). Rollout writes still get -//! native versioning (`checkout`), blob offload, and relationships, all at the -//! schema level. - -use std::collections::{HashMap, HashSet}; +//! # Artifact bytes are stored inline, not blob-v2 offloaded +//! +//! `binary_payload` holds artifact bytes (spec §6) as a plain inline +//! `LargeBinary` column, *not* a blob-v2 offloaded column. Rollout reads go +//! exclusively through the MemWAL LSM scanner +//! ([`RolloutStore::lsm_scanner`]), which has no blob-materialization step: a +//! blob-v2 (`lance-encoding:blob`) column reads back as `None` through it, so +//! [`RolloutStore::get_blob`] could never return the bytes. Inline storage is +//! therefore the only encoding that round-trips. To keep the "learner doesn't +//! pay for artifacts" property (spec §2), list-style scans project the column +//! out (see [`RolloutStore::list`]) rather than relying on physical offload. +//! +//! # Distributed writes: server-id sharding +//! +//! High fan-in rollout ingest (spec §2: thousands of workers) is served by +//! writing through Lance's **MemWAL** rather than a plain [`Dataset::append`]. +//! Each REST server instance owns exactly one MemWAL shard, keyed by its +//! server/instance id (see [`RolloutStoreOptions::shard_id`]). Because a shard +//! has a single active writer per instance and no two instances share a shard, +//! the MemWAL epoch-fencing invariant holds *by construction* — there is never +//! a write war, no matter how a load balancer spreads worker requests across +//! instances. See `specs/rollout-deployment.md`. +//! +//! `MemWAL close-per-append` makes each write durable on object storage before +//! `add` returns, and the read path ([`RolloutStore::lsm_scanner`]) rebuilds +//! purely from object storage (base table ∪ every shard's flushed +//! generations). So any instance reads every instance's writes — reads are not +//! pinned to the writer node. +//! +//! # Reproducibility without `checkout` +//! +//! MemWAL writes land in a separate `_mem_wal/` manifest namespace and do not +//! bump the base dataset version, so per-`add` [`RolloutStore::checkout`] no +//! longer isolates a single append (it did under the old atomic-append path). +//! This is intentional: rollout rows are **append-only and immutable**, so "the +//! exact rollouts that trained checkpoint N" is a *filter over immutable rows* +//! (e.g. by `policy_version`), not a table snapshot — reproducible because the +//! rows never change. `checkout` remains available for base-table time-travel. + +use std::collections::HashMap; use std::sync::Arc; use arrow_array::builder::{ @@ -28,10 +60,15 @@ use arrow_array::{ }; use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit}; use futures::TryStreamExt; +use lance::dataset::mem_wal::{ + DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, +}; use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; -use lance::datatypes::BlobHandling; +use lance::index::DatasetIndexExt; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use lance::{Error as LanceError, Result as LanceResult}; +use lance_index::mem_wal::MEM_WAL_INDEX_NAME; +use uuid::Uuid; use crate::rollout::RolloutRecord; use crate::store::{ @@ -40,68 +77,57 @@ use crate::store::{ RELATIONSHIPS_COLUMN, }; -/// Blob v2 columns permitted on a rollout dataset. `binary_payload` holds -/// artifact bytes offloaded to independent files (spec §6). -const VALID_ROLLOUT_BLOB_COLUMNS: &[&str] = &["binary_payload"]; +/// Number of shard manifest files to scan per batch when discovering the latest +/// shard state (mirrors the constant used by `ContextStore`). +const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; /// Configuration for opening a [`RolloutStore`]. #[derive(Debug, Clone, Default)] pub struct RolloutStoreOptions { /// Object-store credentials/config (e.g. S3), forwarded to Lance. pub storage_options: Option>, - /// Columns physically offloaded via blob v2. Only `binary_payload` is valid. - /// Empty means `binary_payload` is stored inline (no offload). - pub blob_columns: HashSet, + /// Stable identity of the writing server instance (e.g. a StatefulSet + /// ordinal hostname like `rollout-0`). Rollout writes go to the MemWAL + /// shard derived from this id, so each instance owns exactly one shard and + /// no two instances ever contend for the same shard. `None` falls back to a + /// single fixed shard (`"default"`), which is correct for single-instance + /// deployments but must be set per-instance when running multiple writers. + /// See `specs/rollout-deployment.md`. + pub shard_id: Option, } /// A Lance-backed store for RL rollout trajectories. pub struct RolloutStore { dataset: Dataset, - storage_options: Option>, - blob_columns: HashSet, + /// MemWAL shard this instance writes to (derived from `shard_id`). + write_shard: Uuid, } impl RolloutStore { - /// Open an existing rollout dataset or create a new one, defaulting - /// `binary_payload` to a blob v2 (offloaded) column. + /// Open an existing rollout dataset or create a new one with default + /// options (`binary_payload` stored inline; see [`RolloutStoreOptions`]). + /// + /// Uses the fallback single-shard identity; for a multi-instance deployment + /// open with [`RolloutStoreOptions::shard_id`] set per instance. pub async fn open(uri: &str) -> LanceResult { - let mut blob_columns = HashSet::new(); - blob_columns.insert("binary_payload".to_string()); - Self::open_with_options( - uri, - RolloutStoreOptions { - storage_options: None, - blob_columns, - }, - ) - .await + Self::open_with_options(uri, RolloutStoreOptions::default()).await } - /// Open a rollout dataset with explicit storage and blob configuration. + /// Open a rollout dataset with explicit storage and shard configuration. pub async fn open_with_options(uri: &str, options: RolloutStoreOptions) -> LanceResult { - for col in &options.blob_columns { - if !VALID_ROLLOUT_BLOB_COLUMNS.contains(&col.as_str()) { - return Err(LanceError::from(ArrowError::InvalidArgumentError(format!( - "invalid rollout blob column '{}': valid columns are {:?}", - col, VALID_ROLLOUT_BLOB_COLUMNS - )))); - } - } - let storage_options = options.storage_options.clone(); - let blob_columns = options.blob_columns.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 { .. }) => { - Self::create_with_options(uri, storage_options.clone(), &blob_columns).await? + Self::create_with_options(uri, storage_options.clone()).await? } Err(err) => return Err(err), }; Ok(Self { dataset, - storage_options, - blob_columns, + write_shard, }) } @@ -124,43 +150,65 @@ impl RolloutStore { Ok(()) } - /// Append rollout rows in one atomic write; returns the new dataset version. + /// Append rollout rows through this instance's MemWAL shard; returns the + /// current base dataset version. + /// + /// The write is routed to the shard derived from the configured + /// `shard_id`, so concurrent appends from other server instances (each + /// owning a distinct shard) never contend. `close`-per-append flushes the + /// rows to object storage before returning, so they are immediately visible + /// to reads on any instance (see [`Self::lsm_scanner`]). + /// + /// The returned value is the base dataset version, which MemWAL appends do + /// **not** advance; it is retained for API compatibility, not as a per-append + /// snapshot handle (see the module docs on reproducibility). pub async fn add(&mut self, records: &[RolloutRecord]) -> LanceResult { if records.is_empty() { return Ok(self.dataset.manifest.version); } let batch = self.records_to_batch(records)?; - let schema = batch.schema(); - let batches = RecordBatchIterator::new( - vec![Ok::(batch)].into_iter(), - schema, - ); - let mut params = WriteParams { - mode: WriteMode::Append, + self.ensure_mem_wal().await?; + + let config = ShardWriterConfig { + shard_id: self.write_shard, ..Default::default() }; - if let Some(options) = &self.storage_options { - params.store_params = Some(ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_static_options(options.clone()), - )), - ..Default::default() - }); - } + let writer = self.dataset.mem_wal_writer(self.write_shard, config).await?; + writer.put(vec![batch]).await?; + writer.close().await?; - self.dataset.append(batches, Some(params)).await?; Ok(self.dataset.manifest.version) } - /// List rollout rows in storage order. + /// Initialize the (unsharded) MemWAL index on first write, exactly once. + /// Subsequent writes see the index already present and skip this. The shard + /// a write targets is chosen by the writer (`shard_id`), independent of the + /// index's declared sharding strategy. + async fn ensure_mem_wal(&mut self) -> LanceResult<()> { + let indices = self.dataset.load_indices().await?; + let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME); + if !has_mem_wal { + self.dataset.initialize_mem_wal().unsharded().execute().await?; + } + Ok(()) + } + + /// List rollout rows (base table ∪ every instance's flushed MemWAL rows). + /// + /// `binary_payload` is projected out so artifact bytes are never + /// materialized on a list scan (spec §2); fetch them on demand via + /// [`Self::get_blob`]. Ordering is not guaranteed to match append order + /// because the LSM read path dedups by `id` across generations. pub async fn list( &self, limit: Option, offset: Option, ) -> LanceResult> { - let scanner = self.dataset.scan(); + let columns = self.non_blob_columns(); + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let scanner = self.lsm_scanner().await?.project(&refs); let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { @@ -176,37 +224,41 @@ impl RolloutStore { Ok(results) } - /// Retrieve a single rollout row by its unique id. + /// Retrieve a single rollout row by its unique id, including any freshly + /// appended (MemWAL-flushed) row on any instance. `binary_payload` is + /// projected out (fetch bytes via [`Self::get_blob`]). pub async fn get_by_id(&self, id: &str) -> LanceResult> { let escaped_id = id.replace('\'', "''"); - let mut scanner = self.dataset.scan(); - scanner.filter(&format!("id = '{}'", escaped_id))?; - scanner.limit(Some(1), None)?; - + let columns = self.non_blob_columns(); + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let scanner = self + .lsm_scanner() + .await? + .project(&refs) + .filter(&format!("id = '{}'", escaped_id))?; let mut stream = scanner.try_into_stream().await?; - if let Some(batch) = stream.try_next().await? { - return Ok(batch_to_rollout_records(&batch)?.into_iter().next()); + while let Some(batch) = stream.try_next().await? { + for record in batch_to_rollout_records(&batch)? { + if record.id == id { + return Ok(Some(record)); + } + } } Ok(None) } /// Fetch a single artifact row's `binary_payload` bytes on demand. /// - /// When `binary_payload` is offloaded via blob v2 (the default, see - /// [`RolloutStore::open`]), a normal scan returns the column as a - /// description and [`RolloutRecord::binary_payload`] reads back as `None`. - /// This method forces byte materialization with - /// [`BlobHandling::AllBinary`]. Returns `None` if the row or its payload is - /// absent. + /// `list`/`get_by_id` project `binary_payload` out, so it reads back as + /// `None` there. This method projects it in and reads the inline bytes + /// directly. Returns `None` if the row or its payload is absent. pub async fn get_blob(&self, id: &str) -> LanceResult>> { let escaped_id = id.replace('\'', "''"); - let mut scanner = self.dataset.scan(); - scanner.project(&["id", "binary_payload"])?; - scanner.filter(&format!("id = '{}'", escaped_id))?; - if self.blob_columns.contains("binary_payload") { - scanner.blob_handling(BlobHandling::AllBinary); - } - + let scanner = self + .lsm_scanner() + .await? + .project(&["id", "binary_payload"]) + .filter(&format!("id = '{}'", escaped_id))?; let mut stream = scanner.try_into_stream().await?; while let Some(batch) = stream.try_next().await? { let id_array = column_as::(&batch, "id")?; @@ -223,6 +275,56 @@ impl RolloutStore { Ok(None) } + /// Top-level column names excluding `binary_payload`, so list-style scans + /// never materialize artifact bytes. + fn non_blob_columns(&self) -> Vec { + self.dataset + .schema() + .fields + .iter() + .map(|field| field.name.clone()) + .filter(|name| name != "binary_payload") + .collect() + } + + /// Build an LSM scanner over the base table unioned with every shard's + /// flushed MemWAL generations, discovered from object storage. Because the + /// snapshot is rebuilt from shard manifests on each call, one instance sees + /// every other instance's flushed appends — reads are not pinned to the + /// writing instance. Deduplicates by `id`. + async fn lsm_scanner(&self) -> LanceResult { + let object_store = self.dataset.object_store(None).await?; + let branch_location = self.dataset.branch_location(); + let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; + + let mut shard_snapshots = Vec::with_capacity(shard_ids.len()); + for shard_id in shard_ids { + let manifest_store = ShardManifestStore::new( + object_store.clone(), + &branch_location.path, + shard_id, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + let Some(manifest) = manifest_store.read_latest().await? else { + continue; + }; + + let mut snapshot = ShardSnapshot::new(shard_id) + .with_spec_id(manifest.shard_spec_id) + .with_current_generation(manifest.current_generation); + for flushed in manifest.flushed_generations { + snapshot = snapshot.with_flushed_generation(flushed.generation, flushed.path); + } + shard_snapshots.push(snapshot); + } + + Ok(LsmScanner::new( + Arc::new(self.dataset.clone()), + shard_snapshots, + vec!["id".to_string()], + )) + } + async fn load_with_options( uri: &str, storage_options: Option>, @@ -240,9 +342,8 @@ impl RolloutStore { async fn create_with_options( uri: &str, storage_options: Option>, - blob_columns: &HashSet, ) -> LanceResult { - let schema = Arc::new(rollout_schema(blob_columns)); + let schema = Arc::new(rollout_schema()); let empty_batch = RecordBatch::new_empty(schema.clone()); let batches = RecordBatchIterator::new( vec![Ok::(empty_batch)].into_iter(), @@ -503,24 +604,31 @@ impl RolloutStore { } } -/// Arrow schema for a rollout dataset (spec §5). `binary_payload` is marked as a -/// blob v2 column when present in `blob_columns`, so its bytes sink to an -/// independent file and column scans skip them. +/// Derive the MemWAL shard UUID a server instance writes to from its stable +/// instance id. Deterministic (UUID v5), so the same instance id always maps to +/// the same shard across restarts and reopens — the losing/gaining semantics of +/// StatefulSet rescheduling therefore keep writing the same physical shard. +/// `None` maps to a single fixed `"default"` shard for single-instance use. +#[must_use] +pub fn derive_shard_id(instance_id: Option<&str>) -> Uuid { + let input = instance_id.unwrap_or("default"); + Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes()) +} + +/// Arrow schema for a rollout dataset (spec §5). `binary_payload` is a plain +/// inline `LargeBinary` column: rollout reads go through the MemWAL LSM scanner, +/// which has no blob-materialization step, so a blob-v2 offloaded column would +/// read back as `None`. List-style scans project the column out instead (see +/// [`RolloutStore::list`]). #[must_use] -pub fn rollout_schema(blob_columns: &HashSet) -> Schema { +pub fn rollout_schema() -> Schema { let mut id_metadata = HashMap::new(); id_metadata.insert( "lance-schema:unenforced-primary-key".to_string(), "true".to_string(), ); - let binary_field = if blob_columns.contains("binary_payload") { - let mut metadata = HashMap::new(); - metadata.insert("lance-encoding:blob".to_string(), "true".to_string()); - Field::new("binary_payload", DataType::LargeBinary, true).with_metadata(metadata) - } else { - Field::new("binary_payload", DataType::LargeBinary, true) - }; + let binary_field = Field::new("binary_payload", DataType::LargeBinary, true); let fields = vec![ // Identity & grouping. @@ -953,8 +1061,8 @@ mod tests { assert_eq!(a.relation, e.relation); assert_eq!(a.weight, e.weight); } - // `binary_payload` is offloaded via blob v2, so `list`/`get_by_id` - // scans read it back as `None` (byte materialization is verified + // `binary_payload` is projected out of `list`/`get_by_id` scans, so + // they read it back as `None` (byte materialization is verified // separately through `get_blob`). The inline sidecar columns still // round-trip. assert_eq!(actual.payload_size, expected.payload_size); @@ -973,18 +1081,24 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { let mut store = RolloutStore::open(&uri).await.unwrap(); - let version = store + // MemWAL appends land in the `_mem_wal` namespace and do not advance + // the base dataset version; `add` returns it unchanged. + store .add(&[assistant.clone(), artifact.clone()]) .await .unwrap(); - assert!(version > 1, "append should advance the dataset version"); + // The LSM read path dedups by `id` across generations and does not + // guarantee append order, so look rows up by id rather than by + // position. let listed = store.list(None, None).await.unwrap(); assert_eq!(listed.len(), 2); - assert_records_eq(&listed[0], &assistant); - assert_records_eq(&listed[1], &artifact); - // Offloaded blob column is not materialized by a plain scan. - assert_eq!(listed[1].binary_payload, None); + let listed_assistant = listed.iter().find(|r| r.id == "row-0").unwrap(); + let listed_artifact = listed.iter().find(|r| r.id == "row-1").unwrap(); + assert_records_eq(listed_assistant, &assistant); + assert_records_eq(listed_artifact, &artifact); + // Offloaded blob column is projected out of a list scan. + assert_eq!(listed_artifact.binary_payload, None); let fetched = store.get_by_id("row-1").await.unwrap().unwrap(); assert_records_eq(&fetched, &artifact); @@ -1002,7 +1116,13 @@ mod tests { } #[test] - fn checkout_recovers_earlier_version() { + fn appends_accumulate_and_are_immutable() { + // MemWAL appends are append-only: successive `add`s accumulate, and a + // row is never mutated. Reproducing "the rollouts that trained + // checkpoint N" is therefore a filter over immutable rows (here, by id), + // not a per-append table snapshot — so this replaces the old + // `checkout_recovers_earlier_version` test, whose per-add snapshot + // semantics MemWAL intentionally drops. let dir = TempDir::new().unwrap(); let uri = dir.path().to_string_lossy().to_string(); @@ -1012,20 +1132,59 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { let mut store = RolloutStore::open(&uri).await.unwrap(); - let first_version = store.add(&[artifact]).await.unwrap(); + store.add(std::slice::from_ref(&artifact)).await.unwrap(); + + // A second append accumulates rather than replacing the first. store.add(&[assistant_record("row-1")]).await.unwrap(); assert_eq!(store.list(None, None).await.unwrap().len(), 2); - store.checkout(first_version).await.unwrap(); - let recovered = store.list(None, None).await.unwrap(); - assert_eq!(recovered.len(), 1); - assert_eq!(recovered[0].id, "row-0"); - - // Blob offload survives version checkout: the artifact bytes - // written in the earlier version are still materializable after - // rewinding the dataset to it. + // The first-appended row is still present and unchanged after the + // later append (immutability), and its inline artifact bytes remain + // materializable. + let recovered = store.get_by_id("row-0").await.unwrap().unwrap(); + assert_records_eq(&recovered, &artifact); let blob = store.get_blob("row-0").await.unwrap(); assert_eq!(blob.as_deref(), Some(&artifact_bytes[..])); }); } + + #[test] + fn distinct_shards_share_one_dataset() { + // Two instances writing distinct shards of the same dataset both + // contribute to reads, and neither fences the other. This models the + // server-id sharding deployment: a load balancer may route appends to + // either instance, and every reader sees the union. + 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 options = |shard: &str| RolloutStoreOptions { + storage_options: None, + shard_id: Some(shard.to_string()), + }; + + let mut instance_a = RolloutStore::open_with_options(&uri, options("rollout-0")) + .await + .unwrap(); + instance_a.add(&[assistant_record("a-0")]).await.unwrap(); + + let mut instance_b = RolloutStore::open_with_options(&uri, options("rollout-1")) + .await + .unwrap(); + instance_b.add(&[assistant_record("b-0")]).await.unwrap(); + + // Distinct instance ids derive distinct shards. + assert_ne!( + derive_shard_id(Some("rollout-0")), + derive_shard_id(Some("rollout-1")) + ); + + // Either instance's reader sees both shards' rows. + let seen = instance_b.list(None, None).await.unwrap(); + assert_eq!(seen.len(), 2); + assert!(seen.iter().any(|r| r.id == "a-0")); + assert!(seen.iter().any(|r| r.id == "b-0")); + }); + } } diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index 7600ee9..94c649d 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -17,7 +17,7 @@ lance-context-core = { version = "0.5.1", path = "../lance-context-core" } lance-context-api = { version = "0.5.1", path = "../lance-context-api" } axum = { version = "0.8", features = ["json", "multipart"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } -clap = { version = "4", features = ["derive"] } +clap = { version = "4", features = ["derive", "env"] } serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index e645a5d..4ebe2c1 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -12,4 +12,28 @@ pub struct ServerConfig { #[arg(long, default_value = "./lance-data")] pub data_dir: String, + + /// Stable identity of this server instance, used as the MemWAL shard key for + /// rollout writes. Each instance must present a distinct, stable value so it + /// owns exactly one shard and never contends with peers (see + /// `specs/rollout-deployment.md`). In Kubernetes, set this to the + /// StatefulSet pod ordinal hostname (e.g. `rollout-0`). Defaults to the + /// `INSTANCE_ID` env var, then the pod/host `HOSTNAME`; if neither is set, + /// rollout writes fall back to a single shared `default` shard, which is + /// only safe for a single-instance deployment. + #[arg(long, env = "INSTANCE_ID")] + pub instance_id: Option, +} + +impl ServerConfig { + /// Resolve the instance id used for rollout MemWAL sharding: the explicit + /// `--instance-id`/`INSTANCE_ID` if provided, otherwise the `HOSTNAME` + /// environment variable (stable per-pod under a StatefulSet). + #[must_use] + pub fn resolved_instance_id(&self) -> Option { + self.instance_id + .clone() + .or_else(|| std::env::var("HOSTNAME").ok()) + .filter(|value| !value.is_empty()) + } } diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index b7f1020..b051f69 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -624,6 +624,7 @@ mod tests { stores: RwLock::new(stores), rollout_stores: RwLock::new(HashMap::new()), base_path: dir.path().to_path_buf(), + instance_id: None, }); (state, dir) } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 3449d1c..1f67919 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use axum::body::Body; @@ -36,18 +36,10 @@ pub async fn create_rollout_store( } drop(stores); - // `None` keeps the default blob offload of `binary_payload`; an explicit - // empty list stores it inline. - let blob_columns: HashSet = req - .blob_columns - .unwrap_or_else(|| vec!["binary_payload".to_string()]) - .into_iter() - .collect(); - let uri = state.rollout_uri(&req.name); let options = RolloutStoreOptions { storage_options: req.storage_options, - blob_columns, + shard_id: state.instance_id.clone(), }; let store = RolloutStore::open_with_options(&uri, options) @@ -341,7 +333,7 @@ pub async fn get_rollout( })) } -/// Materialize a rollout row's offloaded `binary_payload` bytes. The artifact +/// Materialize a rollout row's `binary_payload` bytes. The artifact /// bytes are opaque, so they stream as `application/octet-stream`. `404` when /// the row is absent or carries no payload. pub async fn fetch_rollout_blob( @@ -512,13 +504,13 @@ mod tests { stores: RwLock::new(HashMap::new()), rollout_stores: RwLock::new(HashMap::new()), base_path: dir.path().to_path_buf(), + instance_id: None, }); let (_status, _info) = create_rollout_store( State(state.clone()), Json(CreateRolloutStoreRequest { name: "rl".to_string(), storage_options: None, - blob_columns: None, }), ) .await @@ -575,7 +567,7 @@ mod tests { } #[tokio::test] - async fn add_rollouts_json_inlines_and_offloads_blob() { + async fn add_rollouts_json_projects_blob_out_but_fetch_materializes() { let (state, _dir) = rollout_state().await; let payload = b"artifact-bytes".to_vec(); let mut record = record_with_size("r0", Some(payload.len() as i64)); @@ -596,7 +588,7 @@ mod tests { assert_eq!(resp.count, 1); assert_eq!(resp.ids, vec!["r0".to_string()]); - // A plain get reads the offloaded column back as None... + // A plain get projects the binary column out, reading it back as None... let Json(got) = get_rollout( State(state.clone()), Path(("rl".to_string(), "r0".to_string())), diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index 0a731c0..eea4895 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -157,6 +157,7 @@ mod tests { stores: RwLock::new(stores), rollout_stores: RwLock::new(HashMap::new()), base_path: dir.path().to_path_buf(), + instance_id: None, }); (state, dir) } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 411953e..efb36a6 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -11,14 +11,20 @@ pub struct AppState { pub stores: RwLock>>>, pub rollout_stores: RwLock>>>, pub base_path: PathBuf, + /// Stable identity of this server instance, used as the MemWAL shard key for + /// rollout writes so each instance owns exactly one shard. `None` falls back + /// to a single shared shard (single-instance deployments only). + pub instance_id: Option, } impl AppState { pub fn new(config: ServerConfig) -> Self { + let instance_id = config.resolved_instance_id(); Self { stores: RwLock::new(HashMap::new()), rollout_stores: RwLock::new(HashMap::new()), base_path: PathBuf::from(&config.data_dir), + instance_id, } } diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs index bd68241..66cf750 100644 --- a/crates/lance-context/src/unified_rollout.rs +++ b/crates/lance-context/src/unified_rollout.rs @@ -1,5 +1,3 @@ -use std::collections::HashSet; - use lance_context_api::{ AddRolloutRequest, AddRolloutsResponse, ContextError, ContextResult, RolloutRecordDto, RolloutStoreApi, @@ -29,17 +27,13 @@ impl RolloutStore { pub async fn open_with_options( uri: &str, storage_options: Option>, - blob_columns: Option>, ) -> Result { - // `None` keeps the default blob offload of `binary_payload`; an explicit - // list (including empty) is taken verbatim. - let blob_columns: HashSet = match blob_columns { - Some(cols) => cols.into_iter().collect(), - None => std::iter::once("binary_payload".to_string()).collect(), - }; let options = RolloutStoreOptions { storage_options, - blob_columns, + // Embedded single-process use writes to the fallback shard. A + // multi-writer embedded deployment should thread a per-writer id + // through here (see `RolloutStoreOptions::shard_id`). + shard_id: None, }; let store = LocalStore::open_with_options(uri, options) .await diff --git a/specs/rollout-deployment.md b/specs/rollout-deployment.md new file mode 100644 index 0000000..90665ec --- /dev/null +++ b/specs/rollout-deployment.md @@ -0,0 +1,196 @@ +# Deployment: Distributed Rollout Ingest with Server-ID Sharding + +**Status:** Draft +**Author:** Beinan Wang +**Date:** 2026-07-09 +**Scope:** How to deploy the `lance-context-server` rollout store as a multi-instance service behind a plain load balancer, and how server-id MemWAL sharding makes high fan-in ingest correct by construction. + +--- + +## 1. Summary + +Rollout ingest is high fan-in: thousands of generation workers append trajectories concurrently (schema-design §2). To serve that from more than one server instance without a coordinator, each instance writes through Lance's **MemWAL** to a shard it owns exclusively, keyed by its **stable server/instance id**. + +The result is a service you can scale horizontally behind an ordinary round-robin load balancer: + +- **No write coordination.** Each instance owns one shard; no two instances share a shard, so MemWAL's single-writer + epoch-fencing invariant holds *by construction* — there is never a write war, no matter how the load balancer spreads requests. +- **No read affinity.** Every append is flushed to object storage before the request returns, and the read path rebuilds from object storage across all shards. Any instance sees every instance's writes. +- **No client changes.** Sharding is entirely server-side. Workers just POST to a load-balancer VIP. + +The one hard rule: **each instance must present a distinct, stable id.** A Kubernetes StatefulSet gives you exactly that for free. + +--- + +## 2. The sharding model + +```text + ┌──────── clients: generation workers (append) + learner (read) ────────┐ + │ POST /rollouts/{name} GET /rollouts/{name}/{id} │ + └───────────────────────────────────┬───────────────────────────────────┘ + │ round-robin, no session affinity + ┌────────▼────────┐ + │ Service (LB) │ + └───┬─────────┬───┘ + ┌─────────────────────┘ └─────────────────────┐ + ┌─────▼─────┐ ┌─────▼─────┐ + │ rollout-0 │ owns shard = uuid5("rollout-0") │ rollout-1 │ owns shard = uuid5("rollout-1") + │ (Pod) │ single active writer │ (Pod) │ single active writer + └─────┬─────┘ └─────┬─────┘ + │ close-per-append → flush to storage │ + └───────────────────────┬──────────────────────────────┘ + ┌────────▼────────┐ + │ object storage │ ONE Lance dataset: + │ (S3 / GCS) │ base ∪ _mem_wal/{shard}/generations + └─────────────────┘ +``` + +- **One instance = one shard.** The write shard is derived deterministically from the instance id: `shard_uuid = uuid_v5(NAMESPACE_OID, instance_id)` (`derive_shard_id` in `rollout_store.rs`). The same id always maps to the same shard across restarts and reschedules. +- **One dataset, many shards.** All instances open the *same* Lance dataset (same URI on shared object storage) and write disjoint shards of it under the `_mem_wal/{shard}/` namespace. They do not each get a private dataset. +- **Single active writer per shard, two ways:** across instances by distinct shard ids; within an instance because `RolloutStore::add` takes `&mut self` and the server holds a per-store write lock, serializing that instance's own appends. + +### Instance-id resolution + +The server resolves its id in this order (`ServerConfig::resolved_instance_id` in `config.rs`): + +1. `--instance-id` flag or `INSTANCE_ID` env var, if set; +2. otherwise the `HOSTNAME` env var (a StatefulSet pod's hostname is its stable ordinal name, e.g. `rollout-0`); +3. otherwise a single fixed `"default"` shard. + +Because Kubernetes sets `HOSTNAME` to the pod name automatically, a StatefulSet needs **no explicit configuration** to shard correctly — though setting `INSTANCE_ID` explicitly (§5) makes the intent obvious and is recommended. + +> **The `default` fallback is single-instance only.** If you run several instances that all resolve to `default` (or you set the *same* id on two of them), they contend for one shard and MemWAL fences all but one writer — appends start failing. Multi-instance deployments **must** give each instance a distinct id. + +--- + +## 3. Why this is correct: single-writer + epoch fencing + +MemWAL is a log-structured writer. Each shard has a manifest chain in object storage, and advancing it is a compare-and-set (`PUT-IF-NOT-EXISTS` with epoch fencing): if two writers try to extend the same shard concurrently, exactly one wins and the other is fenced and must fail or retry. This is the invariant that makes MemWAL safe without a lock service — **but it assumes a single active writer per shard.** + +Consistent hashing, a write-forwarding leader, or a distributed lock are the usual ways to guarantee that. Server-id sharding gets it for free: since the shard id is a pure function of the instance id, and instance ids are distinct and stable, **no two instances ever target the same shard.** The fencing path is never exercised in normal operation; it exists only as a backstop against the misconfiguration in §2 (two writers, one id). + +This is why a *dumb* load balancer is sufficient. Correctness does not depend on *which* instance a request lands on — only on each instance owning its own shard. Round-robin, least-connections, random: all fine. + +--- + +## 4. Read-your-writes and cross-instance visibility + +Writes are made durable synchronously, so reads need no affinity: + +- **Write (`RolloutStore::add`):** append the batch to the shard's memtable (`put`) and then `close()` the writer. `close()` freezes the memtable and blocks until the flush produces a Lance fragment in `_mem_wal/{shard}/` on object storage. When `add` returns, the rows are durable — not merely buffered in the writing instance's memory. +- **Read (`RolloutStore::lsm_scanner`):** on every call, list the dataset's shard ids (`list_mem_wal_latest_shard_ids`), read each shard's latest manifest from object storage (`ShardManifestStore::read_latest`), and scan the union of the base table and every shard's flushed generations, de-duplicated by `id`. + +Because the read snapshot is rebuilt from object storage — not from any instance's local state — an instance sees **every** instance's flushed appends, including its own most recent one. A worker can `POST` a rollout to `rollout-0` and the learner can immediately `GET` it via `rollout-1`; the load balancer may route the two requests anywhere. + +Trade-off: this is the durability-favoring end of the dial. `close`-per-append flushes each batch to object storage, which bounds throughput per shard by flush latency. High aggregate fan-in is absorbed by **adding shards (instances)**, not by batching many appends into one flush. If you later need higher per-shard throughput, group multiple trajectories into a single `add` call at the client. + +--- + +## 5. Kubernetes deployment + +Two pieces: a **StatefulSet** (stable per-pod ids) and a **Service** (the load-balancer VIP clients talk to). A headless Service backs the StatefulSet's stable DNS; a normal Service fronts it for load balancing. + +```yaml +# Headless service: stable network identity for the StatefulSet pods. +apiVersion: v1 +kind: Service +metadata: + name: rollout-headless +spec: + clusterIP: None + selector: { app: rollout } + ports: [{ name: http, port: 3000 }] +--- +# Load-balancer service: the single VIP clients POST/GET against. +apiVersion: v1 +kind: Service +metadata: + name: rollout +spec: + type: LoadBalancer # or ClusterIP behind an Ingress + selector: { app: rollout } + ports: [{ name: http, port: 80, targetPort: 3000 }] +--- +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: rollout +spec: + serviceName: rollout-headless + replicas: 4 # rollout-0 .. rollout-3, each owns one shard + selector: { matchLabels: { app: rollout } } + template: + metadata: { labels: { app: rollout } } + spec: + containers: + - name: server + image: lance-context-server:latest + args: ["--host", "0.0.0.0", "--port", "3000", + "--data-dir", "s3://my-bucket/rollouts"] + env: + # Stable per-pod id → stable shard. metadata.name is rollout-0, + # rollout-1, ... (This is also what HOSTNAME resolves to, so this + # block is belt-and-suspenders — but explicit is clearer.) + - name: INSTANCE_ID + valueFrom: { fieldRef: { fieldPath: metadata.name } } + # Object-store credentials (example: S3). + - name: AWS_REGION + value: us-east-1 + # AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY from a Secret, etc. + ports: [{ containerPort: 3000 }] +``` + +Key points: + +- **Shared dataset on object storage.** `--data-dir` must point at shared object storage (`s3://…`, `gs://…`, …) that **all** pods open at the **same** URI. Rollout datasets live at `{data_dir}/{name}.rollout.lance`. Do **not** give each pod a private `PersistentVolumeClaim` for the data dir — that would create N separate datasets, and reads would not union across them. (PVCs for logs/scratch are fine.) +- **StatefulSet, not Deployment.** A Deployment assigns random pod names, so ids would not be stable across rollouts. A StatefulSet's ordinal names (`rollout-0…`) are stable and reused on reschedule. +- **Replicas = shard count.** Scaling `replicas` changes how many shards absorb the ingest fan-in. + +--- + +## 6. Scaling, rescheduling, and failure + +Every append is flushed before `add` returns and every flushed generation is recorded in the shard manifest on object storage, so instance lifecycle is safe: + +| Event | What happens | Safe? | +|---|---|---| +| **Pod reschedule** (`rollout-2` dies, K8s recreates it) | Same pod name → same instance id → same shard uuid. The new pod resumes writing the same physical shard. Nothing was buffered in memory (close-per-append), so nothing is lost. | ✅ | +| **Scale up** (`replicas: 4 → 6`) | New pods `rollout-4`, `rollout-5` derive fresh distinct shards and begin writing immediately; their rows are visible to all readers on first flush. | ✅ | +| **Scale down** (`replicas: 6 → 4`) | Retired pods stop writing. Their already-flushed shard generations remain in object storage and are still read (the shard manifest persists). The shard simply receives no new appends. | ✅ | +| **Two writers, one id** (misconfig: same `INSTANCE_ID` on two pods, or several `default`s) | Both target one shard; MemWAL epoch-fencing lets one win and fails the other. Appends error. | ❌ Avoid — see §2. | + +There is no compaction/merge step to coordinate: rollout rows are append-only (§7), so shards accumulate independently and are unioned at read time. + +--- + +## 7. Reproducibility is a filter, not a `checkout` + +Rollout rows are **append-only and immutable** — no updates, no deletes. This changes how "exactly which rollouts trained checkpoint N" is answered, versus the earlier atomic-append design: + +- MemWAL appends land in the `_mem_wal/` namespace and **do not advance the base dataset version.** So a per-`add` `checkout(version)` no longer isolates a single append the way plain `Dataset::append` did. `RolloutStore::add` still returns the base version, but for API compatibility only — not as a per-append snapshot handle. +- Reproducibility is instead a **filter over immutable rows** — e.g. `policy_version = 'ckpt-N'`, or an explicit set of `rollout_id`s recorded by the trainer. This is correct precisely *because* the rows never change: the same filter always selects the same bytes. + +`checkout` remains available for base-table time travel, but is not the mechanism for per-checkpoint rollout reproducibility. (See schema-design §3 and §7 — and note that `learner_iteration` is likewise a column, not a dataset version.) + +--- + +## 8. Artifacts are stored inline, not blob-offloaded + +Schema-design §6 anticipated storing `binary_payload` as a **blob v2 offloaded** column. That is **incompatible with the MemWAL read path**, which is the only read path for rollouts: the LSM scanner has no blob-materialization step, so a blob-v2 (`lance-encoding:blob`) column reads back as `None`, and `get_blob` / the `…/blob` endpoint would return nothing. This resolves the schema-design §9 open item ("confirm the blob v2 encoding marker against the pinned Lance version"): **do not** use blob-v2 offload for rollout artifacts. + +Artifacts are therefore stored as a plain inline `LargeBinary` column. The schema-design §2 property — *the learner does not pay for artifacts it does not read* — is preserved **columnar-ly** rather than by physical offload: + +- `list` and `get_by_id` **project `binary_payload` out**, so bulk/hot-path scans never materialize artifact bytes. +- `get_blob` (HTTP `GET /rollouts/{name}/{id}/blob`) projects the column **in** and returns the bytes on demand. + +From a client's perspective the behavior is unchanged: metadata scans are cheap, and artifact bytes are fetched explicitly by id. + +--- + +## 9. Operational checklist + +- [ ] Deploy as a **StatefulSet** (stable ordinal pod names), not a Deployment. +- [ ] Point `--data-dir` at **shared object storage** at the **same URI for every pod**; no per-pod PVC for the dataset. +- [ ] Ensure each instance has a **distinct stable id** — `INSTANCE_ID` from `metadata.name`, or rely on the pod `HOSTNAME`. Never reuse an id across live instances. +- [ ] Front the pods with an ordinary round-robin **Service / LB**; no session affinity or consistent hashing required. +- [ ] Size **replicas** to the ingest fan-in; scale up/down freely (§6). +- [ ] Reproduce training sets by **filtering immutable rows** (e.g. `policy_version`), not by `checkout` (§7). From d7084e829e23e15e276235123b2537c111d8cff8 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Thu, 9 Jul 2026 08:41:53 +0000 Subject: [PATCH 2/2] style: rustfmt rollout_store MemWAL call chains Apply cargo fmt to the MemWAL writer/index-init call chains added in the previous commit; satisfies the `cargo fmt --all -- --check` CI gate. Co-Authored-By: Claude Opus 4 --- crates/lance-context-core/src/rollout_store.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index b96ac42..9a41a54 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -175,7 +175,10 @@ impl RolloutStore { shard_id: self.write_shard, ..Default::default() }; - let writer = self.dataset.mem_wal_writer(self.write_shard, config).await?; + let writer = self + .dataset + .mem_wal_writer(self.write_shard, config) + .await?; writer.put(vec![batch]).await?; writer.close().await?; @@ -190,7 +193,11 @@ impl RolloutStore { let indices = self.dataset.load_indices().await?; let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME); if !has_mem_wal { - self.dataset.initialize_mem_wal().unsharded().execute().await?; + self.dataset + .initialize_mem_wal() + .unsharded() + .execute() + .await?; } Ok(()) }