diff --git a/README.md b/README.md index 1918324d..ae16d04a 100644 --- a/README.md +++ b/README.md @@ -98,7 +98,7 @@ store.add({ "content": "The answer is 42.", "reward": 1.0, "policy_version": "ckpt-100", -}) +}, flush=True) # make this row readable before add returns for row in store.list(): print(row["rollout_id"], row["reward"]) @@ -119,9 +119,19 @@ your event loop: from lance_context import AsyncRolloutStore store = await AsyncRolloutStore.connect_or_create("http://localhost:8080", "rl-run-1") -await store.add({"rollout_id": "traj-1", "role": "assistant", "reward": 1.0}) +await store.add( + {"id": "step-1", "rollout_id": "traj-1", "role": "assistant", "reward": 1.0}, + flush=True, +) +rows = await store.list() # includes step-1 without polling ``` +By default, `add` returns after the write is durable in the WAL. Remote servers +make it readable later on a timer; embedded stores need `flush=True` or +`store.flush()`. Use `flush=True` only when the next operation must read the new +rows because it adds write latency. If flushing fails, the append may still be +durable. Reuse stable `id` values when retrying. + ## Rust The engine is written in Rust and usable directly: diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 8c876aa5..81f13440 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -285,17 +285,40 @@ impl RemoteRolloutStore { cached_version: info.version.unwrap_or(0), }) } + + /// Append rollout rows and wait until they are visible to subsequent reads. + /// + /// Requires server v0.6.5 or newer; older servers ignore the flush option. + /// If flushing fails, the append may still be durable. Retries should reuse + /// the same record ids. + pub async fn add_with_flush( + &mut self, + records: &[AddRolloutRequest], + ) -> ContextResult { + self.add_inner(records, true).await + } + + async fn add_inner( + &mut self, + records: &[AddRolloutRequest], + flush: bool, + ) -> ContextResult { + let resp = if flush { + self.client + .add_rollouts_with_flush(&self.store_name, records) + .await + } else { + self.client.add_rollouts(&self.store_name, records).await + } + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(resp) + } } 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) + self.add_inner(records, false).await } async fn list( @@ -1020,12 +1043,41 @@ impl ContextClient { /// 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. + /// + /// This method does not wait for read visibility. Use + /// [`Self::add_rollouts_with_flush`] when the next operation must read the rows. pub async fn add_rollouts( &self, name: &str, records: &[AddRolloutRequest], + ) -> Result { + self.add_rollouts_inner(name, records, false).await + } + + /// Append rollout rows and wait until the server makes them readable. + /// + /// Requires server v0.6.5 or newer; older servers ignore the flush option. + /// If flushing fails, the append may still be durable. Retries should reuse + /// the same record ids. + pub async fn add_rollouts_with_flush( + &self, + name: &str, + records: &[AddRolloutRequest], + ) -> Result { + self.add_rollouts_inner(name, records, true).await + } + + async fn add_rollouts_inner( + &self, + name: &str, + records: &[AddRolloutRequest], + flush: bool, ) -> Result { let url = self.url(&format!("/rollouts/{}/records", name)); + let mut request = self.http.post(url); + if flush { + request = request.query(&[("flush", true)]); + } let has_blob = records.iter().any(|r| r.binary_payload.is_some()); let resp = if has_blob { @@ -1049,12 +1101,12 @@ impl ContextClient { form = form.part(idx.to_string(), part); } } - self.http.post(url).multipart(form).send().await? + request.multipart(form).send().await? } else { let req = AddRolloutsRequest { records: records.to_vec(), }; - self.http.post(url).json(&req).send().await? + request.json(&req).send().await? }; Self::handle_response(resp).await } diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs index 8d8d445b..ccc36d98 100644 --- a/crates/lance-context/src/unified_rollout.rs +++ b/crates/lance-context/src/unified_rollout.rs @@ -65,6 +65,36 @@ impl RolloutStore { Ok(Self::Remote(store)) } + /// Append rollout rows and wait until they are visible to subsequent reads. + /// + /// Remote stores require server v0.6.5 or newer; older servers ignore the + /// flush option. If flushing fails, the append may still be durable. Retries + /// should reuse the same record ids. A flush may also publish other pending + /// appends. + pub async fn add_with_flush( + &mut self, + records: &[AddRolloutRequest], + ) -> ContextResult { + if records.is_empty() { + return Err(ContextError::InvalidRequest( + "records must not be empty".to_string(), + )); + } + + match self { + Self::Local(store) => { + let response = RolloutStoreApi::add(store.as_mut(), records).await?; + store + .flush() + .await + .map_err(|e| ContextError::Internal(e.to_string()))?; + Ok(response) + } + #[cfg(feature = "remote")] + Self::Remote(store) => store.add_with_flush(records).await, + } + } + /// Seal the local MemWAL memtable so rows written by [`RolloutStoreApi::add`] /// become visible to subsequent reads on this handle. /// @@ -73,8 +103,8 @@ impl RolloutStore { /// sweeper, so without an explicit `flush` a write-then-read sequence /// returns nothing. Call this before reading back rows you just wrote. /// - /// A no-op for `Remote` stores, where the server owns flush scheduling and - /// per-request `?flush=true` provides read-your-write. + /// A no-op for `Remote` stores. Use [`Self::add_with_flush`] when a remote + /// append needs read-your-write visibility. pub async fn flush(&self) -> Result<(), ContextError> { match self { RolloutStore::Local(s) => s diff --git a/docs/src/guide/rollouts.md b/docs/src/guide/rollouts.md index f1f93ec2..14266109 100644 --- a/docs/src/guide/rollouts.md +++ b/docs/src/guide/rollouts.md @@ -23,7 +23,7 @@ store.add({ "content": "The answer is 42.", "reward": 1.0, "policy_version": "ckpt-100", -}) +}, flush=True) # make this row readable before add returns for row in store.list(): print(row["rollout_id"], row["reward"]) @@ -46,9 +46,19 @@ your event loop: from lance_context import AsyncRolloutStore store = await AsyncRolloutStore.connect_or_create("http://localhost:8080", "rl-run-1") -await store.add({"rollout_id": "traj-1", "role": "assistant", "reward": 1.0}) +await store.add( + {"id": "step-1", "rollout_id": "traj-1", "role": "assistant", "reward": 1.0}, + flush=True, +) +rows = await store.list() # includes step-1 without polling ``` +By default, `add` returns after the write is durable in the WAL. Remote servers +make it readable later on a timer; embedded stores need `flush=True` or +`store.flush()`. Use `flush=True` only when the next operation must read the new +rows because it adds write latency. If flushing fails, the append may still be +durable. Reuse stable `id` values when retrying. + ## Further reading - [Rollout schema design](../specs/rollout-schema-design.md) — the full column diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 144f7811..5e543b15 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -2636,13 +2636,22 @@ def version(self) -> int: def add( self, records: Mapping[str, Any] | Iterable[Mapping[str, Any]], + *, + flush: bool = False, ) -> dict[str, Any]: - """Append one record (dict) or many (iterable of dicts).""" - return self._sync.add(_rollout_records_to_json(records)) + """Append one record (dict) or many (iterable of dicts). + + By default, the append is durable but may not yet be readable. Set + ``flush=True`` to wait for visibility. Remote stores require server + 0.6.5 or newer; older servers ignore this option. If flushing fails, + retry with the same stable ``id`` values because the append may already + be durable. + """ + return self._sync.add(_rollout_records_to_json(records), flush=flush) - def add_one(self, **fields: Any) -> dict[str, Any]: + def add_one(self, *, flush: bool = False, **fields: Any) -> dict[str, Any]: """Append a single record given as keyword arguments.""" - return self.add(fields) + return self.add(fields, flush=flush) def flush(self) -> None: """Make previously added rows visible to subsequent reads. @@ -2653,7 +2662,9 @@ def flush(self) -> None: store has no such timer, so a write-then-read sequence returns nothing until you call this. - No-op for stores connected to a remote server. + For stores connected to a remote server this is a no-op. Pass + ``flush=True`` to :meth:`add` when a remote write must be immediately + readable. """ self._sync.flush() @@ -2741,23 +2752,33 @@ def version(self) -> int: async def add( self, records: Mapping[str, Any] | Iterable[Mapping[str, Any]], + *, + flush: bool = False, ) -> dict[str, Any]: - """Append one record (dict) or many (iterable of dicts).""" + """Append one record (dict) or many (iterable of dicts). + + By default, the append is durable but may not yet be readable. Set + ``flush=True`` to wait for visibility. Remote stores require server + 0.6.5 or newer; older servers ignore this option. If flushing fails, + retry with the same stable ``id`` values because the append may already + be durable. + """ payload = _rollout_records_to_json(records) loop = asyncio.get_running_loop() - return await loop.run_in_executor(None, lambda: self._sync.add(payload)) + return await loop.run_in_executor( + None, lambda: self._sync.add(payload, flush=flush) + ) - async def add_one(self, **fields: Any) -> dict[str, Any]: + async def add_one(self, *, flush: bool = False, **fields: Any) -> dict[str, Any]: """Append a single record given as keyword arguments.""" - return await self.add(fields) + return await self.add(fields, flush=flush) async def flush(self) -> None: """Make previously added rows visible to subsequent reads. For an embedded store this seals the local memtable. For a store - connected to a remote server this is a no-op: the server flushes on its - own interval (``ROLLOUT_FLUSH_INTERVAL_SECS``), which bounds how long a - just-written row stays invisible. + connected to a remote server this is a no-op. Pass ``flush=True`` to + :meth:`add` when a remote write must be immediately readable. """ loop = asyncio.get_running_loop() await loop.run_in_executor(None, self._sync.flush) diff --git a/python/src/lib.rs b/python/src/lib.rs index 8f1b8f7a..682cd707 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -2541,7 +2541,8 @@ impl RolloutStore { /// Append rollout rows given as a JSON array of `AddRolloutRequest` objects. /// Returns a dict `{version, ids, count}`. - fn add(&mut self, py: Python<'_>, records_json: &str) -> PyResult { + #[pyo3(signature = (records_json, *, flush = false))] + fn add(&mut self, py: Python<'_>, records_json: &str, flush: bool) -> PyResult { let records: Vec = serde_json::from_str(records_json) .map_err(|e| PyRuntimeError::new_err(format!("invalid records JSON: {e}")))?; if records.is_empty() { @@ -2549,9 +2550,14 @@ impl RolloutStore { "records must not be empty".to_string(), )); } - let resp = py - .allow_threads(|| self.runtime.block_on(self.store.add(&records))) - .map_err(to_py_err)?; + let resp = py.allow_threads(|| { + if flush { + self.runtime.block_on(self.store.add_with_flush(&records)) + } else { + self.runtime.block_on(self.store.add(&records)) + } + }); + let resp = resp.map_err(to_py_err)?; let dict = PyDict::new(py); dict.set_item("version", resp.version)?; dict.set_item("ids", resp.ids)?; diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 6f6f8261..4643dc1f 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -7,6 +7,7 @@ import time import urllib.error import urllib.request +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING @@ -37,8 +38,8 @@ def _wait_for_health(base_url: str, timeout: float = 30.0) -> None: raise RuntimeError(f"server did not become healthy at {url}") -@pytest.fixture() -def server() -> Iterator[str]: +@contextmanager +def _run_server(flush_interval_seconds: int) -> Iterator[str]: """Run a real HTTP server with an isolated dataset directory.""" if not _SERVER_BIN.exists(): # Locally a missing binary is a fair reason to skip. In CI it is not: @@ -50,11 +51,12 @@ def server() -> Iterator[str]: pytest.skip(msg) port = _free_port() with tempfile.TemporaryDirectory() as data_dir: - # Rows are durable on `add` but only become visible when the server's - # sweeper seals the memtable. The 30s production default would make - # every write-then-assert below hang; 1s keeps the tests honest about - # the async-visibility contract without waiting on it. - env = {**os.environ, "ROLLOUT_FLUSH_INTERVAL_SECS": "1"} + # Disable cleanup so another sweeper cannot make flush-on-add tests pass. + env = { + **os.environ, + "ROLLOUT_CLEANUP_INTERVAL_SECS": "0", + "ROLLOUT_FLUSH_INTERVAL_SECS": str(flush_interval_seconds), + } with tempfile.TemporaryFile() as log: proc = subprocess.Popen( [ @@ -90,3 +92,15 @@ def server() -> Iterator[str]: "server did not handle SIGTERM gracefully " f"(exit code {proc.returncode})" ) + + +@pytest.fixture() +def server() -> Iterator[str]: + with _run_server(flush_interval_seconds=1) as base_url: + yield base_url + + +@pytest.fixture() +def server_without_sweeps() -> Iterator[str]: + with _run_server(flush_interval_seconds=0) as base_url: + yield base_url diff --git a/python/tests/test_rollout.py b/python/tests/test_rollout.py index c640fcfa..795d8b3f 100644 --- a/python/tests/test_rollout.py +++ b/python/tests/test_rollout.py @@ -69,6 +69,24 @@ def test_add_many_and_add_one(store_uri): assert {r["id"] for r in rows} == {"row-0", "row-1", "row-2", "row-3"} +def test_add_with_flush_is_immediately_visible(store_uri): + store = RolloutStore.open(store_uri) + + store.add({"id": "deferred", "rollout_id": "traj-1"}) + assert store.list() == [] + + store.add({"id": "visible", "rollout_id": "traj-1"}, flush=True) + assert {row["id"] for row in store.list()} == {"deferred", "visible"} + + store.add_one(id="one", rollout_id="traj-1", flush=True) + assert {row["id"] for row in store.list()} == {"deferred", "visible", "one"} + + with pytest.raises(TypeError): + store.add( # type: ignore[call-arg] + {"id": "positional", "rollout_id": "traj-1"}, True + ) + + def test_id_autogenerated_when_omitted(store_uri): import uuid diff --git a/python/tests/test_rollout_remote.py b/python/tests/test_rollout_remote.py index 9d5b7f2b..a92b6340 100644 --- a/python/tests/test_rollout_remote.py +++ b/python/tests/test_rollout_remote.py @@ -33,9 +33,12 @@ async def _eventually(fn, predicate, timeout: float = 15.0): raise AssertionError(f"condition not met within {timeout}s; last value: {last!r}") -def test_remote_roundtrip(server): +def test_remote_roundtrip(server_without_sweeps): async def run(): - store = await AsyncRolloutStore.connect_or_create(server, "rl-run-1") + store = await AsyncRolloutStore.connect_or_create( + server_without_sweeps, + "rl-run-1", + ) resp = await store.add( [ @@ -56,11 +59,12 @@ async def run(): "binary_payload": b"\x00\x01\x02trace", "payload_size": 8, }, - ] + ], + flush=True, ) assert resp["count"] == 2 - rows = await _eventually(store.list, lambda r: len(r) == 2) + rows = await store.list() assert {r["id"] for r in rows} == {"row-0", "row-1"} one = await store.get("row-0") @@ -76,15 +80,17 @@ async def run(): asyncio.run(run()) -def test_remote_add_one_and_connect(server): +def test_remote_add_one_and_connect(server_without_sweeps): async def run(): - store = await AsyncRolloutStore.connect_or_create(server, "rl-run-2") - await store.add_one(id="only", rollout_id="traj-9", reward=0.5) + store = await AsyncRolloutStore.connect_or_create( + server_without_sweeps, + "rl-run-2", + ) + await store.add_one(id="only", rollout_id="traj-9", reward=0.5, flush=True) - # A second connection sees the first's flushed write (durable, no - # read affinity). - reader = await AsyncRolloutStore.connect(server, "rl-run-2") - rows = await _eventually(reader.list, lambda r: len(r) == 1) + # A separately connected client sees the flushed row immediately. + reader = await AsyncRolloutStore.connect(server_without_sweeps, "rl-run-2") + rows = await reader.list() assert [r["id"] for r in rows] == ["only"] asyncio.run(run())