diff --git a/README.md b/README.md index d0d5dc0..de5a75c 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,13 @@ store.add({ for row in store.list(): print(row["rollout_id"], row["reward"]) + +# Exact-match filters are combined with AND and applied before pagination. +training_rows = store.list(filters={ + "policy_version": "ckpt-100", + "include_in_training": True, + "role": "assistant", +}) ``` In a real training run, generation workers and the learner talk to a shared diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 24301a9..78d13ea 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -130,6 +130,13 @@ pub trait RolloutStoreApi { offset: Option, ) -> impl Future>> + Send; + fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> impl Future>> + Send; + fn get(&self, id: &str) -> impl Future>> + Send; diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 1ece4c3..71c4e39 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -310,6 +310,25 @@ impl RolloutStoreApi for RemoteRolloutStore { Ok(resp.records) } + async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> ContextResult> { + let filters = filters + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| ContextError::InvalidRequest(err.to_string()))?; + let resp = self + .client + .list_rollouts_filtered(&self.store_name, limit, offset, filters.as_deref()) + .await + .map_err(to_ctx_err)?; + Ok(resp.records) + } + async fn get(&self, id: &str) -> ContextResult> { let resp = self .client @@ -776,6 +795,16 @@ impl ContextClient { name: &str, limit: Option, offset: Option, + ) -> Result { + self.list_rollouts_filtered(name, limit, offset, None).await + } + + pub async fn list_rollouts_filtered( + &self, + name: &str, + limit: Option, + offset: Option, + filters: Option<&str>, ) -> Result { let mut request = self .http @@ -786,6 +815,9 @@ impl ContextClient { if let Some(offset) = offset { request = request.query(&[("offset", offset)]); } + if let Some(filters) = filters { + request = request.query(&[("filters", filters)]); + } let resp = request.send().await?; Self::handle_response(resp).await } diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 328cc64..6031569 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -16,7 +16,7 @@ use crate::record::{ LIFECYCLE_ACTIVE, }; use crate::rollout::RolloutRecord; -use crate::rollout_store::RolloutStore; +use crate::rollout_store::{RolloutFilters, RolloutStore}; use crate::store::{CompactionConfig, ContextStore}; impl ContextStoreApi for ContextStore { @@ -400,6 +400,22 @@ impl RolloutStoreApi for RolloutStore { Ok(records.into_iter().map(rollout_record_to_dto).collect()) } + async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> ContextResult> { + let filters = filters + .map(RolloutFilters::from_json_value) + .transpose() + .map_err(ContextError::InvalidRequest)?; + let records = RolloutStore::list_with_filters(self, limit, offset, filters.as_ref()) + .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 diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index e59d986..bae712d 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -72,6 +72,7 @@ use lance::{Error as LanceError, Result as LanceResult}; use lance_index::mem_wal::{ShardManifest, MEM_WAL_INDEX_NAME}; use lance_index::scalar::ScalarIndexParams; use lance_index::IndexType; +use serde_json::Value; use tokio::task::JoinHandle; use tracing::{info, warn}; use uuid::Uuid; @@ -128,6 +129,36 @@ pub struct RolloutFilters { } impl RolloutFilters { + pub fn from_json_value(value: Value) -> Result { + let Value::Object(object) = value else { + return Err("rollout filters must be a JSON object".to_string()); + }; + + let mut filters = Self::default(); + for (key, value) in object { + let string_value = || { + value + .as_str() + .map(str::to_string) + .ok_or_else(|| format!("rollout filter '{key}' must be a string")) + }; + match key.as_str() { + "rollout_id" => filters.rollout_id = Some(string_value()?), + "problem_id" => filters.problem_id = Some(string_value()?), + "policy_version" => filters.policy_version = Some(string_value()?), + "role" => filters.role = Some(string_value()?), + "include_in_training" => { + filters.include_in_training = Some(value.as_bool().ok_or_else(|| { + "rollout filter 'include_in_training' must be a boolean".to_string() + })?); + } + "artifact_type" => filters.artifact_type = Some(string_value()?), + _ => return Err(format!("unsupported rollout filter '{key}'")), + } + } + Ok(filters) + } + fn expression(&self) -> Option { let mut clauses = Vec::new(); for (column, value) in [ @@ -896,22 +927,41 @@ impl RolloutStore { &self, limit: Option, offset: Option, + ) -> LanceResult> { + self.list_with_filters(limit, offset, None).await + } + + /// List rollout rows matching exact field filters. + /// + /// The filter is pushed into every LSM source before merge and + /// deduplication. Rollout rows are immutable by id, so filtering each + /// generation cannot reveal an older state of a mutable record. + pub async fn list_with_filters( + &self, + limit: Option, + offset: Option, + filters: Option<&RolloutFilters>, ) -> LanceResult> { 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 scanner = self.lsm_scanner().await?.project(&refs); + if let Some(predicate) = filters.and_then(RolloutFilters::expression) { + scanner = scanner.filter(&predicate)?; + } + let post_scan_offset = if limit.is_none() { offset } else { None }; + if let Some(limit) = limit { + scanner = scanner.limit(limit, offset); + } + let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { results.extend(batch_to_rollout_records(&batch)?); } - if let Some(offset) = offset { + if let Some(offset) = post_scan_offset { results = results.into_iter().skip(offset).collect(); } - if let Some(limit) = limit { - results.truncate(limit); - } Ok(results) } @@ -1746,6 +1796,51 @@ mod tests { use serde_json::json; use tempfile::TempDir; + #[test] + fn rollout_filters_parse_supported_fields() { + let filters = RolloutFilters::from_json_value(json!({ + "rollout_id": "traj-1", + "problem_id": "problem-7", + "policy_version": "ckpt-42", + "role": "assistant", + "include_in_training": false, + "artifact_type": "screenshot" + })) + .unwrap(); + + assert_eq!(filters.rollout_id.as_deref(), Some("traj-1")); + assert_eq!(filters.problem_id.as_deref(), Some("problem-7")); + assert_eq!(filters.policy_version.as_deref(), Some("ckpt-42")); + assert_eq!(filters.role.as_deref(), Some("assistant")); + assert_eq!(filters.include_in_training, Some(false)); + assert_eq!(filters.artifact_type.as_deref(), Some("screenshot")); + } + + #[test] + fn rollout_filters_reject_unknown_and_wrong_types() { + assert!(RolloutFilters::from_json_value(json!({"reward": 1.0})).is_err()); + assert!(RolloutFilters::from_json_value(json!({"policy_version": 42})).is_err()); + assert!(RolloutFilters::from_json_value(json!({"include_in_training": "yes"})).is_err()); + assert!(RolloutFilters::from_json_value(json!([])).is_err()); + } + + #[test] + fn rollout_filter_expression_escapes_strings_and_fields() { + let filters = RolloutFilters { + policy_version: Some("worker's-ckpt".to_string()), + role: Some("assistant".to_string()), + include_in_training: Some(true), + ..Default::default() + }; + + assert_eq!( + filters.expression().as_deref(), + Some( + "role = 'assistant' AND policy_version = 'worker''s-ckpt' AND include_in_training = true" + ) + ); + } + fn assistant_record(id: &str) -> RolloutRecord { RolloutRecord { id: id.to_string(), @@ -1937,6 +2032,88 @@ mod tests { }); } + #[test] + fn filtered_list_matches_fields_before_pagination() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open(&uri).await.unwrap(); + + let mut first = assistant_record("row-a"); + first.rollout_id = "traj-a".to_string(); + first.problem_id = "problem-a".to_string(); + first.policy_version = Some("ckpt-1".to_string()); + first.include_in_training = Some(true); + + let mut second = assistant_record("row-b"); + second.rollout_id = "traj-b".to_string(); + second.problem_id = "problem-b".to_string(); + second.policy_version = Some("ckpt-2".to_string()); + second.include_in_training = Some(false); + + let mut artifact = artifact_record("row-c", b"bytes"); + artifact.rollout_id = "traj-b".to_string(); + artifact.problem_id = "problem-b".to_string(); + artifact.policy_version = Some("ckpt-2".to_string()); + artifact.include_in_training = Some(false); + + store.add(&[first, second, artifact]).await.unwrap(); + + let filters = RolloutFilters { + policy_version: Some("ckpt-2".to_string()), + include_in_training: Some(false), + ..Default::default() + }; + let all = store + .list_with_filters(None, None, Some(&filters)) + .await + .unwrap(); + assert_eq!(all.len(), 2); + assert!(all + .iter() + .all(|row| row.policy_version.as_deref() == Some("ckpt-2"))); + assert!(all.iter().all(|row| row.include_in_training == Some(false))); + + let page = store + .list_with_filters(Some(1), Some(1), Some(&filters)) + .await + .unwrap(); + assert_eq!(page.len(), 1); + assert_eq!(page[0].policy_version.as_deref(), Some("ckpt-2")); + }); + } + + #[test] + fn filtered_list_supports_dictionary_and_nullable_string_columns() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open(&uri).await.unwrap(); + store + .add(&[ + assistant_record("assistant"), + artifact_record("artifact", b"bytes"), + ]) + .await + .unwrap(); + + let filters = RolloutFilters { + role: Some(ROLE_ARTIFACT.to_string()), + artifact_type: Some("excel_grade_screenshot".to_string()), + ..Default::default() + }; + let rows = store + .list_with_filters(None, None, Some(&filters)) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, "artifact"); + assert!(rows[0].binary_payload.is_none()); + }); + } + #[test] fn appends_accumulate_and_are_immutable() { // MemWAL appends are append-only: successive `add`s accumulate, and a diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 6884df9..3ffd4e5 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -14,7 +14,8 @@ use lance_context_api::{ RolloutStoreInfo, VersionResponse, }; use lance_context_core::{ - CompactionConfig, Relationship, RolloutRecord, RolloutStore, RolloutStoreOptions, + CompactionConfig, Relationship, RolloutFilters, RolloutRecord, RolloutStore, + RolloutStoreOptions, }; use tokio::sync::RwLock; @@ -316,6 +317,7 @@ async fn parse_multipart_rollouts( pub struct RolloutListParams { pub limit: Option, pub offset: Option, + pub filters: Option, } pub async fn list_rollouts( @@ -323,11 +325,22 @@ pub async fn list_rollouts( Path(name): Path, Query(params): Query, ) -> Result, AppError> { + let filters = params + .filters + .as_deref() + .map(|raw| { + let value = serde_json::from_str(raw).map_err(|err| { + AppError::InvalidRequest(format!("invalid rollout filters: {err}")) + })?; + RolloutFilters::from_json_value(value).map_err(AppError::InvalidRequest) + }) + .transpose()?; + let store_lock = state.get_or_open_rollout_store(&name).await?; let store = store_lock.read().await; let records = store - .list(params.limit, params.offset) + .list_with_filters(params.limit, params.offset, filters.as_ref()) .await .map_err(AppError::from_lance)?; @@ -669,6 +682,18 @@ mod tests { } } + fn filtered_record(id: &str, policy_version: &str, training: bool) -> AddRolloutRequest { + AddRolloutRequest { + id: id.to_string(), + rollout_id: format!("traj-{id}"), + problem_id: Some(format!("problem-{id}")), + role: "assistant".to_string(), + policy_version: Some(policy_version.to_string()), + include_in_training: Some(training), + ..Default::default() + } + } + fn json_request(body: Vec) -> Request { Request::builder() .header(header::CONTENT_TYPE, "application/json") @@ -1064,4 +1089,71 @@ mod tests { let Json(listed) = list_rollout_stores(State(state.clone())).await.unwrap(); assert!(listed.stores.is_empty()); } + + #[tokio::test] + async fn list_rollouts_applies_json_filters() { + let (state, _dir) = rollout_state().await; + let body = serde_json::to_vec(&AddRolloutsRequest { + records: vec![ + filtered_record("a", "ckpt-1", true), + filtered_record("b", "ckpt-2", false), + ], + }) + .unwrap(); + let _ = add_rollouts( + State(state.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .unwrap(); + + let Json(response) = list_rollouts( + State(state), + Path("rl".to_string()), + Query(RolloutListParams { + filters: Some( + serde_json::json!({ + "policy_version": "ckpt-2", + "include_in_training": false + }) + .to_string(), + ), + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!(response.records.len(), 1); + assert_eq!(response.records[0].id, "b"); + } + + #[tokio::test] + async fn list_rollouts_rejects_invalid_filters() { + let (state, _dir) = rollout_state().await; + let invalid_json = list_rollouts( + State(state.clone()), + Path("rl".to_string()), + Query(RolloutListParams { + filters: Some("{not-json}".to_string()), + ..Default::default() + }), + ) + .await + .unwrap_err(); + assert!(matches!(invalid_json, AppError::InvalidRequest(_))); + + let unsupported = list_rollouts( + State(state), + Path("rl".to_string()), + Query(RolloutListParams { + filters: Some(serde_json::json!({"reward": 1.0}).to_string()), + ..Default::default() + }), + ) + .await + .unwrap_err(); + assert!(matches!(unsupported, AppError::InvalidRequest(_))); + } } diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index 6059fbd..d63f33c 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -6,7 +6,8 @@ pub use lance_context_core::{ CompactionConfig, CompactionMetrics, CompactionStats, Context, ContextEntry, ContextNamespace, ContextRecord, ContextStoreOptions, IdIndexType, LifecycleQueryOptions, MetadataFilter, PartitionInfo, PartitionSelector, PartitionSpec, RecordFilters, Relationship, RetrieveResult, - RolloutRecord, SearchResult, Snapshot, StateMetadata, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED, + RolloutFilters, RolloutRecord, SearchResult, Snapshot, StateMetadata, LIFECYCLE_ACTIVE, + LIFECYCLE_CONTRADICTED, }; pub use lance_context_api::{ diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs index 06717f8..17c0f6c 100644 --- a/crates/lance-context/src/unified_rollout.rs +++ b/crates/lance-context/src/unified_rollout.rs @@ -113,6 +113,15 @@ impl RolloutStoreApi for RolloutStore { dispatch_ref!(self, list, limit, offset) } + async fn list_filtered( + &self, + limit: Option, + offset: Option, + filters: Option, + ) -> ContextResult> { + dispatch_ref!(self, list_filtered, limit, offset, filters) + } + async fn get(&self, id: &str) -> ContextResult> { dispatch_ref!(self, get, id) } diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index c9514fd..fe2d2f6 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -2583,10 +2583,15 @@ def list( self, limit: int | None = None, offset: int | None = None, + filters: Mapping[str, Any] | None = None, ) -> list[dict[str, Any]]: - """List rollout rows (artifact bytes projected out; fetch via - :meth:`get_blob`).""" - raw = self._sync.list(limit, offset) + """List rollout rows matching exact field filters. + + Supported filter keys are ``rollout_id``, ``problem_id``, + ``policy_version``, ``role``, ``include_in_training``, and + ``artifact_type``. Artifact bytes remain projected out. + """ + raw = self._sync.list(limit, offset, _json_dumps(filters, "filters")) return [_rollout_record_from_json(r) for r in json.loads(raw)] def get(self, id: str) -> dict[str, Any] | None: @@ -2666,10 +2671,14 @@ async def list( self, limit: int | None = None, offset: int | None = None, + filters: Mapping[str, Any] | None = None, ) -> list[dict[str, Any]]: - """List rollout rows (artifact bytes projected out).""" + """List rollout rows matching exact field filters.""" + filters_json = _json_dumps(filters, "filters") loop = asyncio.get_running_loop() - raw = await loop.run_in_executor(None, lambda: self._sync.list(limit, offset)) + raw = await loop.run_in_executor( + None, lambda: self._sync.list(limit, offset, filters_json) + ) return [_rollout_record_from_json(r) for r in json.loads(raw)] async def get(self, id: str) -> dict[str, Any] | None: diff --git a/python/src/lib.rs b/python/src/lib.rs index eeb204e..9fa075d 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -2475,15 +2475,25 @@ impl RolloutStore { /// List rollout rows (artifact bytes projected out). Returns a JSON array /// string of records for the Python wrapper to parse into dicts. - #[pyo3(signature = (limit = None, offset = None))] + #[pyo3(signature = (limit = None, offset = None, filters_json = None))] fn list( &self, py: Python<'_>, limit: Option, offset: Option, + filters_json: Option, ) -> PyResult { + let filters = filters_json + .map(|raw| { + serde_json::from_str(&raw) + .map_err(|err| PyRuntimeError::new_err(format!("invalid filters JSON: {err}"))) + }) + .transpose()?; let records = py - .allow_threads(|| self.runtime.block_on(self.store.list(limit, offset))) + .allow_threads(|| { + self.runtime + .block_on(self.store.list_filtered(limit, offset, filters)) + }) .map_err(to_py_err)?; rollout_records_to_json(&records) } diff --git a/python/tests/test_rollout.py b/python/tests/test_rollout.py index 028b3e8..175a31f 100644 --- a/python/tests/test_rollout.py +++ b/python/tests/test_rollout.py @@ -127,3 +127,55 @@ def test_reopen_sees_prior_rows(store_uri): rows = reopened.list() assert len(rows) == 1 assert rows[0]["id"] == "row-0" + + +def test_list_filters_before_pagination(store_uri): + store = RolloutStore.open(store_uri) + store.add( + [ + { + "id": "row-a", + "rollout_id": "traj-a", + "problem_id": "problem-a", + "policy_version": "ckpt-1", + "role": "assistant", + "include_in_training": True, + }, + { + "id": "row-b", + "rollout_id": "traj-b", + "problem_id": "problem-b", + "policy_version": "ckpt-2", + "role": "assistant", + "include_in_training": False, + }, + { + "id": "row-c", + "rollout_id": "traj-b", + "problem_id": "problem-b", + "policy_version": "ckpt-2", + "role": "artifact", + "artifact_type": "screenshot", + "include_in_training": False, + }, + ] + ) + + rows = store.list( + filters={"policy_version": "ckpt-2", "include_in_training": False} + ) + assert {row["id"] for row in rows} == {"row-b", "row-c"} + + page = store.list( + limit=1, + offset=1, + filters={"policy_version": "ckpt-2", "include_in_training": False}, + ) + assert len(page) == 1 + assert page[0]["policy_version"] == "ckpt-2" + + +def test_list_rejects_invalid_filter_fields(store_uri): + store = RolloutStore.open(store_uri) + with pytest.raises(RuntimeError, match="unsupported rollout filter"): + store.list(filters={"reward": 1.0}) diff --git a/python/tests/test_rollout_remote.py b/python/tests/test_rollout_remote.py index 00a342b..8ee8576 100644 --- a/python/tests/test_rollout_remote.py +++ b/python/tests/test_rollout_remote.py @@ -129,3 +129,33 @@ async def run(): assert [r["id"] for r in rows] == ["only"] asyncio.run(run()) + + +def test_remote_filtered_list(server): + async def run(): + store = await AsyncRolloutStore.connect_or_create(server, "rl-filtered") + await store.add( + [ + { + "id": "row-7", + "rollout_id": "traj-7", + "role": "assistant", + "policy_version": "ckpt-7", + "include_in_training": True, + }, + { + "id": "row-8", + "rollout_id": "traj-8", + "role": "assistant", + "policy_version": "ckpt-8", + "include_in_training": True, + }, + ] + ) + + rows = await store.list( + filters={"policy_version": "ckpt-7", "include_in_training": True} + ) + assert [row["id"] for row in rows] == ["row-7"] + + asyncio.run(run()) diff --git a/specs/rollout-deployment.md b/specs/rollout-deployment.md index 7c92640..7f5bdec 100644 --- a/specs/rollout-deployment.md +++ b/specs/rollout-deployment.md @@ -196,6 +196,13 @@ Rollout rows are **append-only and immutable** — no updates, no deletes. This - 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. +```python +rows = await store.list(filters={ + "policy_version": "ckpt-N", + "include_in_training": True, +}) +``` + `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.) ---