From 6c166e09c73b6413c3c3ec014bf9da980ce032b8 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 8 Jul 2026 20:33:09 -0700 Subject: [PATCH 1/2] Add rollout support to the remote store layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the existing ContextRecord remote path for the native RolloutRecord schema across all five crates. Purely additive — no ContextRecord behavior changes. - lance-context-api: RolloutStoreApi trait (add/list/get/get_blob/ version/checkout) + rollout DTOs. - lance-context-client: RemoteRolloutStore HTTP client with a multipart upload path for large binary artifacts. - lance-context-server: rollout routes + per-name store map. A single POST /records endpoint chooses encoding by Content-Type — JSON inlines binary_payload as base64; multipart sends a `metadata` part first, then one raw binary part per record-with-a-blob, named by zero-based index. The multipart path is strict-reject: metadata-not-first, unparseable/ out-of-bounds/duplicate index, a size mismatch, or a declared-but-absent payload rejects the whole request with 400 before anything is written. - lance-context-core: RolloutRecord <-> DTO mapping + RolloutStoreApi impl. - lance-context: unified RolloutStore enum dispatching Local/Remote, mirroring ContextStore. Artifact blobs travel as raw multipart parts to avoid base64 inflation and full-blob buffering. payload_checksum is a caller-supplied opaque string, stored as-is. Rollout datasets use a `.rollout.lance` suffix so a rollout store and a context store can share a name without colliding. Tests: server 26 (8 rollout, incl. multipart strict-reject), api 6, core 102 (+1 ignored). clippy --all-features clean. Co-Authored-By: Beinan Wang --- Cargo.lock | 18 + crates/lance-context-api/src/lib.rs | 233 ++++++ crates/lance-context-client/Cargo.toml | 2 +- crates/lance-context-client/src/error.rs | 3 + crates/lance-context-client/src/lib.rs | 247 ++++++ crates/lance-context-core/src/api_impl.rs | 141 +++- crates/lance-context-server/Cargo.toml | 2 +- crates/lance-context-server/src/routes/mod.rs | 33 + .../src/routes/records.rs | 1 + .../src/routes/rollouts.rs | 730 ++++++++++++++++++ .../lance-context-server/src/routes/search.rs | 1 + crates/lance-context-server/src/state.rs | 14 +- crates/lance-context/src/lib.rs | 14 +- crates/lance-context/src/unified_rollout.rs | 128 +++ 14 files changed, 1556 insertions(+), 11 deletions(-) create mode 100644 crates/lance-context-server/src/routes/rollouts.rs create mode 100644 crates/lance-context/src/unified_rollout.rs diff --git a/Cargo.lock b/Cargo.lock index 23d3836..9ce1014 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1138,6 +1138,7 @@ dependencies = [ "matchit", "memchr", "mime", + "multer", "percent-encoding", "pin-project-lite", "serde_core", @@ -6776,6 +6777,23 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fafa6961cabd9c63bcd77a45d7e3b7f3b552b70417831fb0f56db717e72407e" +[[package]] +name = "multer" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83e87776546dc87511aa5ee218730c92b666d7264ab6ed41f9d215af9cd5224b" +dependencies = [ + "bytes", + "encoding_rs", + "futures-util", + "http 1.4.1", + "httparse", + "memchr", + "mime", + "spin", + "version_check", +] + [[package]] name = "multimap" version = "0.10.1" diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 979d07d..249f505 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -106,6 +106,43 @@ pub trait ContextStoreApi { fn compaction_stats(&self) -> impl Future> + Send; } +// --------------------------------------------------------------------------- +// Rollout trait +// --------------------------------------------------------------------------- + +/// Remote-capable surface of a rollout store — the subset of `RolloutStore` +/// operations that cross a process boundary. A rollout store is far smaller than +/// a context store: append, read, and version, with no upsert/search/compaction. +/// +/// Artifact bytes (`binary_payload`) travel inline as base64 in JSON or, for +/// large payloads on the batch endpoint, as raw multipart parts. By the time a +/// call reaches this trait the bytes are already materialized in memory, so the +/// signatures are transport-agnostic. +pub trait RolloutStoreApi { + fn add( + &mut self, + records: &[AddRolloutRequest], + ) -> impl Future> + Send; + + fn list( + &self, + limit: Option, + offset: Option, + ) -> impl Future>> + Send; + + fn get(&self, id: &str) + -> impl Future>> + Send; + + /// Materialize a single artifact row's offloaded `binary_payload` bytes. + /// Returns `None` when the row or its payload is absent. + fn get_blob(&self, id: &str) + -> impl Future>>> + Send; + + fn version(&self) -> u64; + + fn checkout(&mut self, version: u64) -> impl Future> + Send; +} + // --------------------------------------------------------------------------- // Context lifecycle // --------------------------------------------------------------------------- @@ -529,6 +566,198 @@ pub struct CompactStatsResponse { pub total_compactions: u64, } +// --------------------------------------------------------------------------- +// Rollout lifecycle +// --------------------------------------------------------------------------- + +#[derive(Debug, Serialize, Deserialize)] +pub struct CreateRolloutStoreRequest { + pub name: String, + #[serde(default)] + pub storage_options: Option>, + #[serde(default)] + pub blob_columns: Option>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RolloutStoreInfo { + pub name: String, + pub uri: String, + pub version: u64, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListRolloutStoresResponse { + pub stores: Vec, +} + +// --------------------------------------------------------------------------- +// Rollout records +// --------------------------------------------------------------------------- + +/// One rollout row to append. Unlike [`AddRecordRequest`], `id` is +/// client-supplied: rollout rows carry externally-meaningful identity (a +/// trajectory row from a harness). Because those ids are opaque and may repeat, +/// the multipart upload path matches raw binary parts to records by their +/// zero-based index in the `records` array, not by `id`. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AddRolloutRequest { + pub id: String, + pub rollout_id: String, + /// Defaults to `rollout_id` server-side when omitted (non-grouped rollouts). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub problem_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dataset: Option, + #[serde(default)] + pub sequence_order: i32, + #[serde(default = "default_rollout_role")] + pub role: String, + /// Defaults to the server's current time when omitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub created_at: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(default = "default_content_type")] + pub content_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_tokens: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_tokens: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub num_input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub num_output_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ref_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub loss_mask: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub advantage: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reward: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub raw_reward: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub grader_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub score: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub include_in_training: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub exclude_reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub policy_version: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub relationships: Vec, + /// Inline artifact bytes. On the batch endpoint, large payloads may instead + /// travel as a raw multipart part named for this record's zero-based index + /// in the `records` array, avoiding base64 inflation. + #[serde( + default, + skip_serializing_if = "Option::is_none", + serialize_with = "serialize_base64_opt", + deserialize_with = "deserialize_base64_opt" + )] + pub binary_payload: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload_size: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload_checksum: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddRolloutsRequest { + pub records: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddRolloutsResponse { + pub version: u64, + pub ids: Vec, + pub count: usize, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RolloutRecordDto { + pub id: String, + pub rollout_id: String, + pub problem_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dataset: Option, + pub sequence_order: i32, + pub role: String, + pub created_at: DateTime, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content: Option, + pub content_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_tokens: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_tokens: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub num_input_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub num_output_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ref_logprobs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub loss_mask: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub advantage: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reward: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub raw_reward: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub grader_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub score: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub include_in_training: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub exclude_reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub policy_version: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub relationships: Vec, + /// Present only when materialized on demand (`get_blob`); a plain list/get + /// scan reads the offloaded column back as `None`. + #[serde( + default, + skip_serializing_if = "Option::is_none", + serialize_with = "serialize_base64_opt", + deserialize_with = "deserialize_base64_opt" + )] + pub binary_payload: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload_size: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload_checksum: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ListRolloutsResponse { + pub records: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct GetRolloutResponse { + pub record: Option, +} + // --------------------------------------------------------------------------- // Error // --------------------------------------------------------------------------- @@ -556,6 +785,10 @@ fn default_role() -> String { "user".to_string() } +fn default_rollout_role() -> String { + "assistant".to_string() +} + fn default_upsert_key() -> String { "external_id".to_string() } diff --git a/crates/lance-context-client/Cargo.toml b/crates/lance-context-client/Cargo.toml index 471bf20..1c8912b 100644 --- a/crates/lance-context-client/Cargo.toml +++ b/crates/lance-context-client/Cargo.toml @@ -10,7 +10,7 @@ keywords = ["context", "lance", "client", "api"] [dependencies] lance-context-api = { version = "0.5.1", path = "../lance-context-api" } -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls"] } serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2" diff --git a/crates/lance-context-client/src/error.rs b/crates/lance-context-client/src/error.rs index 0ea36b1..49cd1dd 100644 --- a/crates/lance-context-client/src/error.rs +++ b/crates/lance-context-client/src/error.rs @@ -3,6 +3,9 @@ pub enum ClientError { #[error("HTTP error: {0}")] Http(#[from] reqwest::Error), + #[error("serialization error: {0}")] + Serialize(#[from] serde_json::Error), + #[error("API error ({status}): [{code}] {message}")] Api { status: u16, diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 3bbd88d..e13ea0e 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -251,6 +251,97 @@ impl ContextStoreApi for RemoteContextStore { } } +pub struct RemoteRolloutStore { + client: ContextClient, + store_name: String, + cached_version: u64, +} + +impl RemoteRolloutStore { + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let client = ContextClient::new(base_url); + let info = client.get_rollout_store(store_name).await?; + Ok(Self { + client, + store_name: store_name.to_string(), + cached_version: info.version, + }) + } + + pub async fn connect_or_create( + base_url: &str, + req: &CreateRolloutStoreRequest, + ) -> Result { + let client = ContextClient::new(base_url); + let info = match client.get_rollout_store(&req.name).await { + Ok(info) => info, + Err(ClientError::Api { status: 404, .. }) => client.create_rollout_store(req).await?, + Err(e) => return Err(e), + }; + Ok(Self { + client, + store_name: req.name.clone(), + cached_version: info.version, + }) + } +} + +impl RolloutStoreApi for RemoteRolloutStore { + async fn add(&mut self, records: &[AddRolloutRequest]) -> ContextResult { + let resp = self + .client + .add_rollouts(&self.store_name, records) + .await + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(resp) + } + + async fn list( + &self, + limit: Option, + offset: Option, + ) -> ContextResult> { + let resp = self + .client + .list_rollouts(&self.store_name, limit, offset) + .await + .map_err(to_ctx_err)?; + Ok(resp.records) + } + + async fn get(&self, id: &str) -> ContextResult> { + let resp = self + .client + .get_rollout(&self.store_name, id) + .await + .map_err(to_ctx_err)?; + Ok(resp.record) + } + + async fn get_blob(&self, id: &str) -> ContextResult>> { + self.client + .fetch_rollout_blob(&self.store_name, id) + .await + .map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + self.cached_version + } + + async fn checkout(&mut self, version: u64) -> ContextResult<()> { + let req = CheckoutRequest { version }; + let resp = self + .client + .checkout_rollout(&self.store_name, &req) + .await + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(()) + } +} + fn to_ctx_err(err: ClientError) -> ContextError { match err { ClientError::Api { @@ -276,6 +367,7 @@ fn to_ctx_err(err: ClientError) -> ContextError { } => ContextError::InvalidRequest(message), ClientError::Api { message, .. } => ContextError::Internal(message), ClientError::Http(e) => ContextError::Internal(e.to_string()), + ClientError::Serialize(e) => ContextError::InvalidRequest(e.to_string()), } } @@ -591,6 +683,161 @@ impl ContextClient { Self::handle_response(resp).await } + pub async fn create_rollout_store( + &self, + req: &CreateRolloutStoreRequest, + ) -> Result { + let resp = self + .http + .post(self.url("/rollouts")) + .json(req) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn list_rollout_stores(&self) -> Result { + let resp = self.http.get(self.url("/rollouts")).send().await?; + Self::handle_response(resp).await + } + + pub async fn get_rollout_store(&self, name: &str) -> Result { + let resp = self + .http + .get(self.url(&format!("/rollouts/{}", name))) + .send() + .await?; + Self::handle_response(resp).await + } + + pub async fn delete_rollout_store(&self, name: &str) -> Result<(), ClientError> { + let resp = self + .http + .delete(self.url(&format!("/rollouts/{}", name))) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(Self::extract_error(resp).await) + } + } + + /// Append rollout rows. When any record carries `binary_payload`, the request + /// is sent as `multipart/form-data`: the first part, `metadata`, holds the + /// records array with each `binary_payload` stripped to null; each record that + /// carries a blob then contributes one raw binary part named for that record's + /// zero-based index in the metadata array (`"0"`, `"1"`, ...). Naming by index + /// keeps part names round-trip safe (record ids may contain arbitrary bytes and + /// are not unique). The `metadata` part is sent first so the server can parse + /// the manifest before matching binary parts. When no record carries bytes, a + /// plain JSON body is sent instead. + pub async fn add_rollouts( + &self, + name: &str, + records: &[AddRolloutRequest], + ) -> Result { + let url = self.url(&format!("/rollouts/{}/records", name)); + let has_blob = records.iter().any(|r| r.binary_payload.is_some()); + + let resp = if has_blob { + let stripped: Vec = records + .iter() + .map(|r| { + let mut without_bytes = r.clone(); + without_bytes.binary_payload = None; + without_bytes + }) + .collect(); + let metadata = serde_json::to_string(&AddRolloutsRequest { records: stripped })?; + + // metadata must be the first part: multer parses parts sequentially and + // the server needs the manifest before it can match binary parts by index. + let mut form = reqwest::multipart::Form::new().text("metadata", metadata); + for (idx, r) in records.iter().enumerate() { + if let Some(bytes) = &r.binary_payload { + let part = reqwest::multipart::Part::bytes(bytes.clone()) + .mime_str("application/octet-stream")?; + form = form.part(idx.to_string(), part); + } + } + self.http.post(url).multipart(form).send().await? + } else { + let req = AddRolloutsRequest { + records: records.to_vec(), + }; + self.http.post(url).json(&req).send().await? + }; + Self::handle_response(resp).await + } + + pub async fn list_rollouts( + &self, + name: &str, + limit: Option, + offset: Option, + ) -> Result { + let mut request = self + .http + .get(self.url(&format!("/rollouts/{}/records", name))); + if let Some(limit) = limit { + request = request.query(&[("limit", limit)]); + } + if let Some(offset) = offset { + request = request.query(&[("offset", offset)]); + } + let resp = request.send().await?; + Self::handle_response(resp).await + } + + pub async fn get_rollout( + &self, + name: &str, + id: &str, + ) -> Result { + let resp = self + .http + .get(self.url(&format!("/rollouts/{}/records/{}", name, id))) + .send() + .await?; + Self::handle_response(resp).await + } + + /// Materialize a single artifact row's offloaded `binary_payload` bytes. + /// Returns `None` when the row or its payload is absent (server 404). + pub async fn fetch_rollout_blob( + &self, + name: &str, + id: &str, + ) -> Result>, ClientError> { + let resp = self + .http + .get(self.url(&format!("/rollouts/{}/records/{}/blob", name, id))) + .send() + .await?; + if resp.status().is_success() { + Ok(Some(resp.bytes().await?.to_vec())) + } else if resp.status().as_u16() == 404 { + Ok(None) + } else { + Err(Self::extract_error(resp).await) + } + } + + pub async fn checkout_rollout( + &self, + name: &str, + req: &CheckoutRequest, + ) -> Result { + let resp = self + .http + .post(self.url(&format!("/rollouts/{}/checkout", name))) + .json(req) + .send() + .await?; + Self::handle_response(resp).await + } + async fn handle_response( resp: reqwest::Response, ) -> Result { diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index be55240..5e83f32 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -3,9 +3,10 @@ use serde_json::Value; use uuid::Uuid; use lance_context_api::{ - AddRecordRequest, AddRecordsResponse, CompactRequest, CompactResponse, CompactStatsResponse, - ContextError, ContextResult, ContextStoreApi, DeleteRecordResponse, RecordDto, RecordPatchDto, - RelationshipDto, RetrieveRequest, RetrieveResultDto, SearchRequest, SearchResultDto, + AddRecordRequest, AddRecordsResponse, AddRolloutRequest, AddRolloutsResponse, CompactRequest, + CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, + DeleteRecordResponse, RecordDto, RecordPatchDto, RelationshipDto, RetrieveRequest, + RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchRequest, SearchResultDto, StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; @@ -14,6 +15,8 @@ use crate::record::{ ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, LIFECYCLE_ACTIVE, }; +use crate::rollout::RolloutRecord; +use crate::rollout_store::RolloutStore; use crate::store::{CompactionConfig, ContextStore}; impl ContextStoreApi for ContextStore { @@ -366,6 +369,138 @@ impl ContextStoreApi for ContextStore { } } +impl RolloutStoreApi for RolloutStore { + async fn add(&mut self, records: &[AddRolloutRequest]) -> ContextResult { + let mut ids = Vec::with_capacity(records.len()); + let mut core_records = Vec::with_capacity(records.len()); + for r in records { + ids.push(r.id.clone()); + core_records.push(rollout_record_from_add_request(r)); + } + + let count = core_records.len(); + let version = RolloutStore::add(self, &core_records) + .await + .map_err(to_ctx_err)?; + Ok(AddRolloutsResponse { + version, + ids, + count, + }) + } + + async fn list( + &self, + limit: Option, + offset: Option, + ) -> ContextResult> { + let records = RolloutStore::list(self, limit, offset) + .await + .map_err(to_ctx_err)?; + Ok(records.into_iter().map(rollout_record_to_dto).collect()) + } + + async fn get(&self, id: &str) -> ContextResult> { + let record = RolloutStore::get_by_id(self, id) + .await + .map_err(to_ctx_err)?; + Ok(record.map(rollout_record_to_dto)) + } + + async fn get_blob(&self, id: &str) -> ContextResult>> { + RolloutStore::get_blob(self, id).await.map_err(to_ctx_err) + } + + fn version(&self) -> u64 { + RolloutStore::version(self) + } + + async fn checkout(&mut self, version: u64) -> ContextResult<()> { + RolloutStore::checkout(self, version) + .await + .map_err(to_ctx_err) + } +} + +fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { + RolloutRecord { + id: r.id.clone(), + rollout_id: r.rollout_id.clone(), + problem_id: r.problem_id.clone().unwrap_or_else(|| r.rollout_id.clone()), + dataset: r.dataset.clone(), + sequence_order: r.sequence_order, + role: r.role.clone(), + created_at: r.created_at.unwrap_or_else(Utc::now), + content: r.content.clone(), + content_type: r.content_type.clone(), + input_tokens: r.input_tokens.clone(), + output_tokens: r.output_tokens.clone(), + num_input_tokens: r.num_input_tokens, + num_output_tokens: r.num_output_tokens, + output_logprobs: r.output_logprobs.clone(), + input_logprobs: r.input_logprobs.clone(), + ref_logprobs: r.ref_logprobs.clone(), + loss_mask: r.loss_mask.clone(), + advantage: r.advantage, + reward: r.reward, + raw_reward: r.raw_reward, + grader_id: r.grader_id.clone(), + score: r.score, + include_in_training: r.include_in_training, + exclude_reason: r.exclude_reason.clone(), + policy_version: r.policy_version.clone(), + relationships: r + .relationships + .iter() + .cloned() + .map(dto_to_relationship) + .collect(), + binary_payload: r.binary_payload.clone(), + payload_size: r.payload_size, + payload_checksum: r.payload_checksum.clone(), + metadata: r.metadata.clone(), + } +} + +fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { + RolloutRecordDto { + id: r.id, + rollout_id: r.rollout_id, + problem_id: r.problem_id, + dataset: r.dataset, + sequence_order: r.sequence_order, + role: r.role, + created_at: r.created_at, + content: r.content, + content_type: r.content_type, + input_tokens: r.input_tokens, + output_tokens: r.output_tokens, + num_input_tokens: r.num_input_tokens, + num_output_tokens: r.num_output_tokens, + output_logprobs: r.output_logprobs, + input_logprobs: r.input_logprobs, + ref_logprobs: r.ref_logprobs, + loss_mask: r.loss_mask, + advantage: r.advantage, + reward: r.reward, + raw_reward: r.raw_reward, + grader_id: r.grader_id, + score: r.score, + include_in_training: r.include_in_training, + exclude_reason: r.exclude_reason, + policy_version: r.policy_version, + relationships: r + .relationships + .into_iter() + .map(relationship_to_dto) + .collect(), + binary_payload: r.binary_payload, + payload_size: r.payload_size, + payload_checksum: r.payload_checksum, + metadata: r.metadata, + } +} + fn dto_to_relationship(r: RelationshipDto) -> Relationship { Relationship { target_id: r.target_id, diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index 647757d..7600ee9 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -15,7 +15,7 @@ path = "src/main.rs" [dependencies] 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"] } +axum = { version = "0.8", features = ["json", "multipart"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index e255c2a..818258b 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -2,11 +2,13 @@ pub mod compact; pub mod contexts; pub mod health; pub mod records; +pub mod rollouts; pub mod search; pub mod versions; use std::sync::Arc; +use axum::extract::DefaultBodyLimit; use axum::routing::{delete, get, patch, post, put}; use axum::Router; @@ -75,6 +77,37 @@ pub fn router() -> Router> { "/api/v1/contexts/{name}/compact/stats", get(compact::compact_stats), ) + .route("/api/v1/rollouts", post(rollouts::create_rollout_store)) + .route("/api/v1/rollouts", get(rollouts::list_rollout_stores)) + .route( + "/api/v1/rollouts/{name}", + get(rollouts::get_rollout_store), + ) + .route( + "/api/v1/rollouts/{name}", + delete(rollouts::delete_rollout_store), + ) + .route( + "/api/v1/rollouts/{name}/records", + post(rollouts::add_rollouts) + .layer(DefaultBodyLimit::max(rollouts::MAX_ROLLOUT_UPLOAD_BYTES)), + ) + .route( + "/api/v1/rollouts/{name}/records", + get(rollouts::list_rollouts), + ) + .route( + "/api/v1/rollouts/{name}/records/{id}", + get(rollouts::get_rollout), + ) + .route( + "/api/v1/rollouts/{name}/records/{id}/blob", + get(rollouts::fetch_rollout_blob), + ) + .route( + "/api/v1/rollouts/{name}/checkout", + post(rollouts::checkout_rollout), + ) } #[cfg(test)] diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 4037497..b7f1020 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -622,6 +622,7 @@ mod tests { stores.insert(context_name.to_string(), Arc::new(RwLock::new(store))); let state = Arc::new(AppState { stores: RwLock::new(stores), + rollout_stores: RwLock::new(HashMap::new()), base_path: dir.path().to_path_buf(), }); (state, dir) diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs new file mode 100644 index 0000000..d44849c --- /dev/null +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -0,0 +1,730 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use axum::body::Body; +use axum::extract::{FromRequest, Multipart, Path, Query, Request, State}; +use axum::http::{header, StatusCode}; +use axum::response::Response; +use axum::Json; +use chrono::Utc; +use lance_context_api::{ + AddRolloutRequest, AddRolloutsRequest, AddRolloutsResponse, CheckoutRequest, + CreateRolloutStoreRequest, GetRolloutResponse, ListRolloutStoresResponse, ListRolloutsResponse, + RelationshipDto, RolloutRecordDto, RolloutStoreInfo, VersionResponse, +}; +use lance_context_core::{Relationship, RolloutRecord, RolloutStore, RolloutStoreOptions}; +use tokio::sync::RwLock; + +use crate::error::AppError; +use crate::state::AppState; + +/// Upper bound on a single rollout append request body (JSON or multipart). +/// Artifact blobs routinely exceed axum's 2 MiB default, so the batch endpoint +/// raises the ceiling well above it while still bounding memory. +pub const MAX_ROLLOUT_UPLOAD_BYTES: usize = 1024 * 1024 * 1024; + +pub async fn create_rollout_store( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), AppError> { + let stores = state.rollout_stores.read().await; + if stores.contains_key(&req.name) { + return Err(AppError::AlreadyExists(format!( + "Rollout store '{}' already exists", + req.name + ))); + } + 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, + }; + + let store = RolloutStore::open_with_options(&uri, options) + .await + .map_err(AppError::from_lance)?; + let version = store.version(); + + let mut stores = state.rollout_stores.write().await; + stores.insert(req.name.clone(), Arc::new(RwLock::new(store))); + + Ok(( + StatusCode::CREATED, + Json(RolloutStoreInfo { + name: req.name, + uri, + version, + }), + )) +} + +pub async fn list_rollout_stores( + State(state): State>, +) -> Json { + let stores = state.rollout_stores.read().await; + let mut out = Vec::with_capacity(stores.len()); + + for (name, store_lock) in stores.iter() { + let store = store_lock.read().await; + out.push(RolloutStoreInfo { + name: name.clone(), + uri: state.rollout_uri(name), + version: store.version(), + }); + } + + Json(ListRolloutStoresResponse { stores: out }) +} + +pub async fn get_rollout_store( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))?; + let store = store_lock.read().await; + + Ok(Json(RolloutStoreInfo { + name: name.clone(), + uri: state.rollout_uri(&name), + version: store.version(), + })) +} + +pub async fn delete_rollout_store( + State(state): State>, + Path(name): Path, +) -> Result { + let mut stores = state.rollout_stores.write().await; + if stores.remove(&name).is_none() { + return Err(AppError::NotFound(format!( + "Rollout store '{}' does not exist", + name + ))); + } + + let uri = state.rollout_uri(&name); + if let Err(e) = tokio::fs::remove_dir_all(&uri).await { + tracing::warn!("Failed to remove rollout data at {}: {}", uri, e); + } + + Ok(StatusCode::NO_CONTENT) +} + +/// Append rollout rows. The single POST endpoint accepts either encoding, +/// chosen by `Content-Type`: +/// +/// - `application/json`: an [`AddRolloutsRequest`]; any `binary_payload` rides +/// inline as base64. +/// - `multipart/form-data`: a first part named `metadata` holding the records +/// array (with each `binary_payload` null), followed by one raw binary part +/// per record that carries a blob, named for that record's zero-based index. +/// +/// Both encodings funnel into a single [`RolloutStore::add`]. A rollout append +/// is atomic and feeds training, so the multipart path is validated strictly: +/// any part whose name is not a valid record index, any index with no matching +/// record, a duplicate index, a byte length disagreeing with the record's +/// `payload_size`, or a record that declares `payload_size` yet receives no +/// bytes rejects the whole request with `400` before anything is written. +pub async fn add_rollouts( + State(state): State>, + Path(name): Path, + req: Request, +) -> Result<(StatusCode, Json), AppError> { + let content_type = req + .headers() + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_string(); + + let records = if content_type_is(&content_type, "multipart/form-data") { + let multipart = Multipart::from_request(req, &state) + .await + .map_err(|e| AppError::InvalidRequest(format!("invalid multipart request: {e}")))?; + parse_multipart_rollouts(multipart).await? + } else if content_type_is(&content_type, "application/json") { + let bytes = axum::body::to_bytes(req.into_body(), MAX_ROLLOUT_UPLOAD_BYTES) + .await + .map_err(|e| AppError::InvalidRequest(format!("failed to read request body: {e}")))?; + let manifest: AddRolloutsRequest = serde_json::from_slice(&bytes) + .map_err(|e| AppError::InvalidRequest(format!("invalid JSON body: {e}")))?; + manifest.records + } else { + return Err(AppError::InvalidRequest(format!( + "unsupported Content-Type '{content_type}': expected application/json or multipart/form-data" + ))); + }; + + if records.is_empty() { + return Err(AppError::InvalidRequest( + "records array must not be empty".to_string(), + )); + } + + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? + .clone(); + drop(stores); + + let ids: Vec = records.iter().map(|r| r.id.clone()).collect(); + let core_records: Vec = + records.iter().map(rollout_record_from_add_request).collect(); + let count = core_records.len(); + + let mut store = store_lock.write().await; + let version = store + .add(&core_records) + .await + .map_err(AppError::from_lance)?; + + Ok(( + StatusCode::CREATED, + Json(AddRolloutsResponse { + version, + ids, + count, + }), + )) +} + +/// Parse a `multipart/form-data` rollout append into records with their blob +/// bytes attached. Enforces the strict-reject contract described on +/// [`add_rollouts`]. +async fn parse_multipart_rollouts( + mut multipart: Multipart, +) -> Result, AppError> { + // The manifest must arrive first so blob parts can be matched by index as + // they stream in. + let first = multipart + .next_field() + .await + .map_err(|e| AppError::InvalidRequest(format!("malformed multipart body: {e}")))? + .ok_or_else(|| AppError::InvalidRequest("multipart body is empty".to_string()))?; + if first.name() != Some("metadata") { + return Err(AppError::InvalidRequest(format!( + "first multipart part must be named 'metadata', found '{}'", + first.name().unwrap_or("") + ))); + } + let metadata_bytes = first + .bytes() + .await + .map_err(|e| AppError::InvalidRequest(format!("failed to read metadata part: {e}")))?; + let manifest: AddRolloutsRequest = serde_json::from_slice(&metadata_bytes) + .map_err(|e| AppError::InvalidRequest(format!("invalid metadata JSON: {e}")))?; + let mut records = manifest.records; + + // Collect the raw binary parts, keyed by the record index they target. + let mut blobs: HashMap> = HashMap::new(); + while let Some(field) = multipart + .next_field() + .await + .map_err(|e| AppError::InvalidRequest(format!("malformed multipart body: {e}")))? + { + let part_name = field + .name() + .ok_or_else(|| { + AppError::InvalidRequest("multipart part is missing a name".to_string()) + })? + .to_string(); + let idx: usize = part_name.parse().map_err(|_| { + AppError::InvalidRequest(format!( + "unexpected multipart part '{part_name}': expected a record index" + )) + })?; + if idx >= records.len() { + return Err(AppError::InvalidRequest(format!( + "multipart part '{idx}' has no matching record ({} records supplied)", + records.len() + ))); + } + let bytes = field.bytes().await.map_err(|e| { + AppError::InvalidRequest(format!("failed to read binary part '{idx}': {e}")) + })?; + if blobs.insert(idx, bytes.to_vec()).is_some() { + return Err(AppError::InvalidRequest(format!( + "duplicate multipart part for record index {idx}" + ))); + } + } + + // Attach each blob to its record, verifying the declared size. + for (idx, bytes) in blobs { + let record = &mut records[idx]; + if let Some(expected) = record.payload_size { + if expected != bytes.len() as i64 { + return Err(AppError::InvalidRequest(format!( + "record index {idx}: payload_size {expected} does not match received {} bytes", + bytes.len() + ))); + } + } + record.binary_payload = Some(bytes); + } + + // Symmetric strict-reject: a record that declares a payload but received no + // bytes (no inline payload and no matching part) would be written with an + // empty artifact, silently dropping training data. Reject it. + for (idx, record) in records.iter().enumerate() { + if record.payload_size.is_some() && record.binary_payload.is_none() { + return Err(AppError::InvalidRequest(format!( + "record index {idx} declares payload_size but no binary part was supplied" + ))); + } + } + + Ok(records) +} + +#[derive(Debug, Default, serde::Deserialize)] +pub struct RolloutListParams { + pub limit: Option, + pub offset: Option, +} + +pub async fn list_rollouts( + State(state): State>, + Path(name): Path, + Query(params): Query, +) -> Result, AppError> { + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? + .clone(); + drop(stores); + + let store = store_lock.read().await; + let records = store + .list(params.limit, params.offset) + .await + .map_err(AppError::from_lance)?; + + Ok(Json(ListRolloutsResponse { + records: records.into_iter().map(rollout_record_to_dto).collect(), + })) +} + +pub async fn get_rollout( + State(state): State>, + Path((name, id)): Path<(String, String)>, +) -> Result, AppError> { + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? + .clone(); + drop(stores); + + let store = store_lock.read().await; + let record = store.get_by_id(&id).await.map_err(AppError::from_lance)?; + + Ok(Json(GetRolloutResponse { + record: record.map(rollout_record_to_dto), + })) +} + +/// Materialize a rollout row's offloaded `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( + State(state): State>, + Path((name, id)): Path<(String, String)>, +) -> Result { + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? + .clone(); + drop(stores); + + let store = store_lock.read().await; + let bytes = store + .get_blob(&id) + .await + .map_err(AppError::from_lance)? + .ok_or_else(|| AppError::NotFound(format!("Rollout '{}' has no payload", id)))?; + + Response::builder() + .header(header::CONTENT_TYPE, "application/octet-stream") + .body(Body::from(bytes)) + .map_err(|err| AppError::Internal(err.to_string())) +} + +pub async fn checkout_rollout( + State(state): State>, + Path(name): Path, + Json(req): Json, +) -> Result, AppError> { + let stores = state.rollout_stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? + .clone(); + drop(stores); + + let mut store = store_lock.write().await; + store + .checkout(req.version) + .await + .map_err(AppError::from_lance)?; + + Ok(Json(VersionResponse { + version: store.version(), + })) +} + +fn content_type_is(content_type: &str, expected: &str) -> bool { + content_type + .split(';') + .next() + .map(str::trim) + .is_some_and(|mime| mime.eq_ignore_ascii_case(expected)) +} + +fn dto_to_relationship(r: RelationshipDto) -> Relationship { + Relationship { + target_id: r.target_id, + relation: r.relation, + weight: r.weight, + } +} + +fn relationship_to_dto(r: Relationship) -> RelationshipDto { + RelationshipDto { + target_id: r.target_id, + relation: r.relation, + weight: r.weight, + } +} + +fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { + RolloutRecord { + id: r.id.clone(), + rollout_id: r.rollout_id.clone(), + problem_id: r.problem_id.clone().unwrap_or_else(|| r.rollout_id.clone()), + dataset: r.dataset.clone(), + sequence_order: r.sequence_order, + role: r.role.clone(), + created_at: r.created_at.unwrap_or_else(Utc::now), + content: r.content.clone(), + content_type: r.content_type.clone(), + input_tokens: r.input_tokens.clone(), + output_tokens: r.output_tokens.clone(), + num_input_tokens: r.num_input_tokens, + num_output_tokens: r.num_output_tokens, + output_logprobs: r.output_logprobs.clone(), + input_logprobs: r.input_logprobs.clone(), + ref_logprobs: r.ref_logprobs.clone(), + loss_mask: r.loss_mask.clone(), + advantage: r.advantage, + reward: r.reward, + raw_reward: r.raw_reward, + grader_id: r.grader_id.clone(), + score: r.score, + include_in_training: r.include_in_training, + exclude_reason: r.exclude_reason.clone(), + policy_version: r.policy_version.clone(), + relationships: r + .relationships + .iter() + .cloned() + .map(dto_to_relationship) + .collect(), + binary_payload: r.binary_payload.clone(), + payload_size: r.payload_size, + payload_checksum: r.payload_checksum.clone(), + metadata: r.metadata.clone(), + } +} + +fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { + RolloutRecordDto { + id: r.id, + rollout_id: r.rollout_id, + problem_id: r.problem_id, + dataset: r.dataset, + sequence_order: r.sequence_order, + role: r.role, + created_at: r.created_at, + content: r.content, + content_type: r.content_type, + input_tokens: r.input_tokens, + output_tokens: r.output_tokens, + num_input_tokens: r.num_input_tokens, + num_output_tokens: r.num_output_tokens, + output_logprobs: r.output_logprobs, + input_logprobs: r.input_logprobs, + ref_logprobs: r.ref_logprobs, + loss_mask: r.loss_mask, + advantage: r.advantage, + reward: r.reward, + raw_reward: r.raw_reward, + grader_id: r.grader_id, + score: r.score, + include_in_training: r.include_in_training, + exclude_reason: r.exclude_reason, + policy_version: r.policy_version, + relationships: r + .relationships + .into_iter() + .map(relationship_to_dto) + .collect(), + binary_payload: r.binary_payload, + payload_size: r.payload_size, + payload_checksum: r.payload_checksum, + metadata: r.metadata, + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use lance_context_api::CreateRolloutStoreRequest; + use tempfile::TempDir; + use tokio::sync::RwLock; + + use super::*; + use crate::state::AppState; + + async fn rollout_state() -> (Arc, TempDir) { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState { + stores: RwLock::new(HashMap::new()), + rollout_stores: RwLock::new(HashMap::new()), + base_path: dir.path().to_path_buf(), + }); + let (_status, _info) = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: "rl".to_string(), + storage_options: None, + blob_columns: None, + }), + ) + .await + .expect("create rollout store"); + (state, dir) + } + + fn record_with_size(id: &str, payload_size: Option) -> AddRolloutRequest { + AddRolloutRequest { + id: id.to_string(), + rollout_id: format!("traj-{id}"), + payload_size, + ..Default::default() + } + } + + fn json_request(body: Vec) -> Request { + Request::builder() + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap() + } + + /// Assemble a `multipart/form-data` body from ordered (name, bytes) parts. + fn multipart_request(boundary: &str, parts: &[(&str, Vec)]) -> Request { + let mut body = Vec::new(); + for (name, data) in parts { + body.extend_from_slice(format!("--{boundary}\r\n").as_bytes()); + body.extend_from_slice( + format!("Content-Disposition: form-data; name=\"{name}\"\r\n\r\n").as_bytes(), + ); + body.extend_from_slice(data); + body.extend_from_slice(b"\r\n"); + } + body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes()); + Request::builder() + .header( + header::CONTENT_TYPE, + format!("multipart/form-data; boundary={boundary}"), + ) + .body(Body::from(body)) + .unwrap() + } + + async fn count_rollouts(state: &Arc) -> usize { + let Json(list) = list_rollouts( + State(state.clone()), + Path("rl".to_string()), + Query(RolloutListParams::default()), + ) + .await + .unwrap(); + list.records.len() + } + + #[tokio::test] + async fn add_rollouts_json_inlines_and_offloads_blob() { + 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)); + record.binary_payload = Some(payload.clone()); + let body = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record], + }) + .unwrap(); + + let (status, Json(resp)) = add_rollouts( + State(state.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::CREATED); + assert_eq!(resp.count, 1); + assert_eq!(resp.ids, vec!["r0".to_string()]); + + // A plain get reads the offloaded column back as None... + let Json(got) = get_rollout(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) + .await + .unwrap(); + let dto = got.record.expect("row present"); + assert_eq!(dto.payload_size, Some(payload.len() as i64)); + assert!(dto.binary_payload.is_none()); + + // ...but the blob endpoint materializes the bytes. + let resp = fetch_rollout_blob(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) + .await + .unwrap(); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(bytes.as_ref(), payload.as_slice()); + } + + #[tokio::test] + async fn add_rollouts_multipart_attaches_blob_by_index() { + let (state, _dir) = rollout_state().await; + let payload = b"raw-multipart-bytes".to_vec(); + let metadata = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", Some(payload.len() as i64))], + }) + .unwrap(); + let req = multipart_request("BOUND", &[("metadata", metadata), ("0", payload.clone())]); + + let (status, Json(resp)) = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap(); + assert_eq!(status, StatusCode::CREATED); + assert_eq!(resp.count, 1); + + let resp = fetch_rollout_blob(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) + .await + .unwrap(); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(bytes.as_ref(), payload.as_slice()); + } + + #[tokio::test] + async fn add_rollouts_multipart_rejects_metadata_not_first() { + let (state, _dir) = rollout_state().await; + let req = multipart_request("BOUND", &[("0", b"bytes".to_vec())]); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + assert_eq!(count_rollouts(&state).await, 0); + } + + #[tokio::test] + async fn add_rollouts_multipart_rejects_orphan_index() { + let (state, _dir) = rollout_state().await; + let metadata = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", Some(4))], + }) + .unwrap(); + // Part "1" has no matching record (only index 0 exists). + let req = multipart_request("BOUND", &[("metadata", metadata), ("1", b"data".to_vec())]); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + assert_eq!(count_rollouts(&state).await, 0); + } + + #[tokio::test] + async fn add_rollouts_multipart_rejects_duplicate_index() { + let (state, _dir) = rollout_state().await; + let metadata = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", Some(4))], + }) + .unwrap(); + let req = multipart_request( + "BOUND", + &[ + ("metadata", metadata), + ("0", b"data".to_vec()), + ("0", b"dupe".to_vec()), + ], + ); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + assert_eq!(count_rollouts(&state).await, 0); + } + + #[tokio::test] + async fn add_rollouts_multipart_rejects_size_mismatch() { + let (state, _dir) = rollout_state().await; + // Declared payload_size (999) disagrees with the 4 bytes received. + let metadata = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", Some(999))], + }) + .unwrap(); + let req = multipart_request("BOUND", &[("metadata", metadata), ("0", b"data".to_vec())]); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + assert_eq!(count_rollouts(&state).await, 0); + } + + #[tokio::test] + async fn add_rollouts_multipart_rejects_declared_payload_without_part() { + let (state, _dir) = rollout_state().await; + // Record declares payload_size but no binary part is attached. + let metadata = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", Some(4))], + }) + .unwrap(); + let req = multipart_request("BOUND", &[("metadata", metadata)]); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), req) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + assert_eq!(count_rollouts(&state).await, 0); + } + + #[tokio::test] + async fn add_rollouts_rejects_empty_records() { + let (state, _dir) = rollout_state().await; + let body = serde_json::to_vec(&AddRolloutsRequest { records: vec![] }).unwrap(); + let err = add_rollouts(State(state.clone()), Path("rl".to_string()), json_request(body)) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + } +} diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index a4f2d5d..0a731c0 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -155,6 +155,7 @@ mod tests { stores.insert(CTX.to_string(), Arc::new(RwLock::new(store))); let state = Arc::new(AppState { stores: RwLock::new(stores), + rollout_stores: RwLock::new(HashMap::new()), base_path: dir.path().to_path_buf(), }); (state, dir) diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 853fd3e..411953e 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -2,13 +2,14 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; -use lance_context_core::ContextStore; +use lance_context_core::{ContextStore, RolloutStore}; use tokio::sync::RwLock; use crate::config::ServerConfig; pub struct AppState { pub stores: RwLock>>>, + pub rollout_stores: RwLock>>>, pub base_path: PathBuf, } @@ -16,6 +17,7 @@ impl AppState { pub fn new(config: ServerConfig) -> Self { Self { stores: RwLock::new(HashMap::new()), + rollout_stores: RwLock::new(HashMap::new()), base_path: PathBuf::from(&config.data_dir), } } @@ -26,4 +28,14 @@ impl AppState { .to_string_lossy() .to_string() } + + /// Rollout datasets live under a distinct `.rollout.lance` suffix so a + /// rollout store and a context store may share the same name without + /// colliding on disk. + 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/src/lib.rs b/crates/lance-context/src/lib.rs index 17077d7..09a95e1 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -6,18 +6,22 @@ pub use lance_context_core::{ CompactionConfig, CompactionMetrics, CompactionStats, Context, ContextEntry, ContextNamespace, ContextRecord, ContextStoreOptions, IdIndexType, LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, PartitionSpec, RecordFilters, Relationship, RetrieveResult, - SearchResult, Snapshot, StateMetadata, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, + RolloutRecord, SearchResult, Snapshot, StateMetadata, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, }; pub use lance_context_api::{ - AddRecordRequest, AddRecordsResponse, CompactRequest, CompactResponse, CompactStatsResponse, - ContextError, ContextResult, ContextStoreApi, DeleteRecordResponse, RecordDto, RelationshipDto, - RetrieveRequest, RetrieveResponse, RetrieveResultDto, SearchResultDto, UpsertRecordRequest, + AddRecordRequest, AddRecordsResponse, AddRolloutRequest, AddRolloutsResponse, CompactRequest, + CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, + DeleteRecordResponse, RecordDto, RelationshipDto, RetrieveRequest, RetrieveResponse, + RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchResultDto, UpsertRecordRequest, UpsertRecordResponse, }; #[cfg(feature = "remote")] -pub use lance_context_client::{ClientError, RemoteContextStore}; +pub use lance_context_client::{ClientError, RemoteContextStore, RemoteRolloutStore}; mod unified; pub use unified::ContextStore; + +mod unified_rollout; +pub use unified_rollout::RolloutStore; diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs new file mode 100644 index 0000000..bd68241 --- /dev/null +++ b/crates/lance-context/src/unified_rollout.rs @@ -0,0 +1,128 @@ +use std::collections::HashSet; + +use lance_context_api::{ + AddRolloutRequest, AddRolloutsResponse, ContextError, ContextResult, RolloutRecordDto, + RolloutStoreApi, +}; +use lance_context_core::{RolloutStore as LocalStore, RolloutStoreOptions}; + +#[cfg(feature = "remote")] +use lance_context_client::RemoteRolloutStore; + +/// A rollout store that is either an in-process Lance dataset (`Local`) or a +/// handle to a remote server (`Remote`). Mirrors [`crate::ContextStore`] but for +/// the RL rollout schema. +pub enum RolloutStore { + Local(Box), + #[cfg(feature = "remote")] + Remote(RemoteRolloutStore), +} + +impl RolloutStore { + pub async fn open(uri: &str) -> Result { + let store = LocalStore::open(uri) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Local(Box::new(store))) + } + + 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, + }; + let store = LocalStore::open_with_options(uri, options) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Local(Box::new(store))) + } + + #[cfg(feature = "remote")] + pub async fn connect(base_url: &str, store_name: &str) -> Result { + let store = RemoteRolloutStore::connect(base_url, store_name) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } + + #[cfg(feature = "remote")] + pub async fn connect_or_create( + base_url: &str, + req: &lance_context_api::CreateRolloutStoreRequest, + ) -> Result { + let store = RemoteRolloutStore::connect_or_create(base_url, req) + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(Self::Remote(store)) + } +} + +macro_rules! dispatch_mut { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + RolloutStore::Local(s) => RolloutStoreApi::$method(s.as_mut() $(, $arg)*).await, + #[cfg(feature = "remote")] + RolloutStore::Remote(s) => RolloutStoreApi::$method(s $(, $arg)*).await, + } + }; +} + +macro_rules! dispatch_ref { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + RolloutStore::Local(s) => RolloutStoreApi::$method(s.as_ref() $(, $arg)*).await, + #[cfg(feature = "remote")] + RolloutStore::Remote(s) => RolloutStoreApi::$method(s $(, $arg)*).await, + } + }; +} + +macro_rules! dispatch_sync { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + RolloutStore::Local(s) => RolloutStoreApi::$method(s.as_ref() $(, $arg)*), + #[cfg(feature = "remote")] + RolloutStore::Remote(s) => RolloutStoreApi::$method(s $(, $arg)*), + } + }; +} + +impl RolloutStoreApi for RolloutStore { + async fn add(&mut self, records: &[AddRolloutRequest]) -> ContextResult { + dispatch_mut!(self, add, records) + } + + async fn list( + &self, + limit: Option, + offset: Option, + ) -> ContextResult> { + dispatch_ref!(self, list, limit, offset) + } + + async fn get(&self, id: &str) -> ContextResult> { + dispatch_ref!(self, get, id) + } + + async fn get_blob(&self, id: &str) -> ContextResult>> { + dispatch_ref!(self, get_blob, id) + } + + fn version(&self) -> u64 { + dispatch_sync!(self, version) + } + + async fn checkout(&mut self, version: u64) -> ContextResult<()> { + dispatch_mut!(self, checkout, version) + } +} From 558ab6a100db3cc1f8372e7eb30aba7e43b79cd4 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 8 Jul 2026 20:35:27 -0700 Subject: [PATCH 2/2] Apply cargo fmt to rollout remote-store sources Co-Authored-By: Beinan Wang --- crates/lance-context-api/src/lib.rs | 3 +- crates/lance-context-server/src/routes/mod.rs | 5 +-- .../src/routes/rollouts.rs | 43 +++++++++++++------ 3 files changed, 31 insertions(+), 20 deletions(-) diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 249f505..b43bbbe 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -135,8 +135,7 @@ pub trait RolloutStoreApi { /// Materialize a single artifact row's offloaded `binary_payload` bytes. /// Returns `None` when the row or its payload is absent. - fn get_blob(&self, id: &str) - -> impl Future>>> + Send; + fn get_blob(&self, id: &str) -> impl Future>>> + Send; fn version(&self) -> u64; diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index 818258b..e183919 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -79,10 +79,7 @@ pub fn router() -> Router> { ) .route("/api/v1/rollouts", post(rollouts::create_rollout_store)) .route("/api/v1/rollouts", get(rollouts::list_rollout_stores)) - .route( - "/api/v1/rollouts/{name}", - get(rollouts::get_rollout_store), - ) + .route("/api/v1/rollouts/{name}", get(rollouts::get_rollout_store)) .route( "/api/v1/rollouts/{name}", delete(rollouts::delete_rollout_store), diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index d44849c..3449d1c 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -182,8 +182,10 @@ pub async fn add_rollouts( drop(stores); let ids: Vec = records.iter().map(|r| r.id.clone()).collect(); - let core_records: Vec = - records.iter().map(rollout_record_from_add_request).collect(); + let core_records: Vec = records + .iter() + .map(rollout_record_from_add_request) + .collect(); let count = core_records.len(); let mut store = store_lock.write().await; @@ -595,17 +597,23 @@ mod tests { assert_eq!(resp.ids, vec!["r0".to_string()]); // A plain get reads the offloaded column back as None... - let Json(got) = get_rollout(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) - .await - .unwrap(); + let Json(got) = get_rollout( + State(state.clone()), + Path(("rl".to_string(), "r0".to_string())), + ) + .await + .unwrap(); let dto = got.record.expect("row present"); assert_eq!(dto.payload_size, Some(payload.len() as i64)); assert!(dto.binary_payload.is_none()); // ...but the blob endpoint materializes the bytes. - let resp = fetch_rollout_blob(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) - .await - .unwrap(); + let resp = fetch_rollout_blob( + State(state.clone()), + Path(("rl".to_string(), "r0".to_string())), + ) + .await + .unwrap(); let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) .await .unwrap(); @@ -628,9 +636,12 @@ mod tests { assert_eq!(status, StatusCode::CREATED); assert_eq!(resp.count, 1); - let resp = fetch_rollout_blob(State(state.clone()), Path(("rl".to_string(), "r0".to_string()))) - .await - .unwrap(); + let resp = fetch_rollout_blob( + State(state.clone()), + Path(("rl".to_string(), "r0".to_string())), + ) + .await + .unwrap(); let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) .await .unwrap(); @@ -722,9 +733,13 @@ mod tests { async fn add_rollouts_rejects_empty_records() { let (state, _dir) = rollout_state().await; let body = serde_json::to_vec(&AddRolloutsRequest { records: vec![] }).unwrap(); - let err = add_rollouts(State(state.clone()), Path("rl".to_string()), json_request(body)) - .await - .unwrap_err(); + let err = add_rollouts( + State(state.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .unwrap_err(); assert!(matches!(err, AppError::InvalidRequest(_))); } }