From b8e15919d8bfcd87a673d8bdee4272e847483dba Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Fri, 10 Jul 2026 17:56:32 +0000 Subject: [PATCH 1/2] feat(python): rollout store bindings (local + remote) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rollout store had no Python surface — only Rust core/client/server. This adds `RolloutStore` (sync, embedded via open() or remote via connect()) and `RemoteRolloutStore` (async HTTP wrapper), mirroring the existing Context / RemoteContext bindings. Both dispatch through the unified `lance-context` RolloutStore enum, so one pyclass covers embedded and remote. Rollout rows carry ~35 fields, so records cross FFI as JSON (serde on the Rust side), with `binary_payload` base64 per the DTO contract; the Python wrapper accepts raw bytes and dict/list[dict] as well as add_one(**kwargs). Methods: open / connect / connect_or_create, add, add_one, list, get, get_blob, version, checkout. Exported from lance_context. Tests: test_rollout.py (embedded add/list/get/get_blob/reopen, dict+list+ kwargs, blob roundtrip) and test_rollout_remote.py (spawns a real lance-context-server subprocess, async connect_or_create + roundtrip + cross-connection read-your-writes; skips if the server binary isn't built). ruff + pyright clean. Co-Authored-By: Claude Opus 4 --- Cargo.lock | 1 + crates/lance-context/src/lib.rs | 6 +- python/Cargo.toml | 1 + python/python/lance_context/__init__.py | 4 + python/python/lance_context/api.py | 231 ++++++++++++++++++++++++ python/src/lib.rs | 159 ++++++++++++++++ python/tests/test_rollout.py | 107 +++++++++++ python/tests/test_rollout_remote.py | 131 ++++++++++++++ 8 files changed, 637 insertions(+), 3 deletions(-) create mode 100644 python/tests/test_rollout.py create mode 100644 python/tests/test_rollout_remote.py diff --git a/Cargo.lock b/Cargo.lock index e8970e9..cb61283 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5428,6 +5428,7 @@ name = "lance-context-python" version = "0.6.0" dependencies = [ "chrono", + "lance-context", "lance-context-api", "lance-context-client", "lance-context-core", diff --git a/crates/lance-context/src/lib.rs b/crates/lance-context/src/lib.rs index 09a95e1..6059fbd 100644 --- a/crates/lance-context/src/lib.rs +++ b/crates/lance-context/src/lib.rs @@ -12,9 +12,9 @@ pub use lance_context_core::{ pub use lance_context_api::{ AddRecordRequest, AddRecordsResponse, AddRolloutRequest, AddRolloutsResponse, CompactRequest, CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, - DeleteRecordResponse, RecordDto, RelationshipDto, RetrieveRequest, RetrieveResponse, - RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchResultDto, UpsertRecordRequest, - UpsertRecordResponse, + CreateRolloutStoreRequest, DeleteRecordResponse, RecordDto, RelationshipDto, RetrieveRequest, + RetrieveResponse, RetrieveResultDto, RolloutRecordDto, RolloutStoreApi, SearchResultDto, + UpsertRecordRequest, UpsertRecordResponse, }; #[cfg(feature = "remote")] diff --git a/python/Cargo.toml b/python/Cargo.toml index d105e5a..e9d46d0 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -11,6 +11,7 @@ crate-type = ["cdylib"] [dependencies] chrono = { version = "0.4", default-features = false, features = ["clock"] } +lance-context = { path = "../crates/lance-context", features = ["remote"] } lance-context-api = { path = "../crates/lance-context-api" } lance-context-client = { path = "../crates/lance-context-client" } lance-context-core = { path = "../crates/lance-context-core" } diff --git a/python/python/lance_context/__init__.py b/python/python/lance_context/__init__.py index da6b8fd..b6c6f1b 100644 --- a/python/python/lance_context/__init__.py +++ b/python/python/lance_context/__init__.py @@ -6,6 +6,8 @@ ContextNamespace, EmbeddingProvider, RemoteContext, + RemoteRolloutStore, + RolloutStore, __version__, ) from .embeddings import ( # pyright: ignore[reportMissingImports] @@ -19,5 +21,7 @@ "EmbeddingProvider", "MultiModalEmbeddingProvider", "RemoteContext", + "RemoteRolloutStore", + "RolloutStore", "__version__", ] diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 31a9f33..ea8d395 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -18,6 +18,9 @@ from ._internal import ( # pyright: ignore[reportMissingImports] RemoteContext as _RemoteContext, ) +from ._internal import ( # pyright: ignore[reportMissingImports] + RolloutStore as _RolloutStore, +) from ._internal import version as _version # pyright: ignore[reportMissingImports] from .embeddings import EmbeddingProvider, _build_provider, supports_media @@ -27,6 +30,8 @@ "ContextNamespace", "EmbeddingProvider", "RemoteContext", + "RemoteRolloutStore", + "RolloutStore", "__version__", ] @@ -2438,3 +2443,229 @@ async def compaction_stats(self) -> dict[str, Any]: def __repr__(self) -> str: return f"RemoteContext(version={self._sync.version()})" + + +# --------------------------------------------------------------------------- +# Rollout stores +# --------------------------------------------------------------------------- + +# Fields that carry raw artifact bytes; the native layer marshals records as +# JSON with these base64-encoded, so the wrapper encodes on the way in. +_ROLLOUT_BLOB_FIELD = "binary_payload" + + +def _rollout_record_to_native(record: Mapping[str, Any]) -> dict[str, Any]: + """Normalize one rollout record dict for the JSON FFI boundary. + + Accepts ``binary_payload`` as raw ``bytes``/``bytearray`` and base64-encodes + it (the DTO's JSON representation), leaving everything else untouched. + """ + out = dict(record) + blob = out.get(_ROLLOUT_BLOB_FIELD) + if isinstance(blob, (bytes, bytearray)): + import base64 + + out[_ROLLOUT_BLOB_FIELD] = base64.b64encode(bytes(blob)).decode("ascii") + return out + + +def _rollout_records_to_json( + records: Mapping[str, Any] | Iterable[Mapping[str, Any]], +) -> str: + # `records` is either a single mapping or an iterable of mappings. Because a + # Mapping is itself Iterable, static narrowing of the union is ambiguous, so + # branch on the concrete runtime type via a local typed as Any. + raw: Any = records + if isinstance(raw, Mapping): + payload = [_rollout_record_to_native(raw)] + else: + payload = [_rollout_record_to_native(r) for r in raw] + if not payload: + raise ValueError("records must not be empty") + return json.dumps(payload) + + +def _rollout_record_from_json(record: Mapping[str, Any]) -> dict[str, Any]: + """Decode a rollout record dict coming back from the native layer. + + ``binary_payload``, when present, arrives base64-encoded; decode it to raw + ``bytes`` so callers get bytes back from ``get``/``list`` if it was + materialized. + """ + out = dict(record) + blob = out.get(_ROLLOUT_BLOB_FIELD) + if isinstance(blob, str): + import base64 + + out[_ROLLOUT_BLOB_FIELD] = base64.b64decode(blob) + return out + + +class RolloutStore: + """Synchronous store for RL rollout trajectories. + + Open an embedded dataset with :meth:`open`, or talk to a remote + ``lance-context-server`` with :meth:`connect` / :meth:`connect_or_create`. + + Records are plain dicts matching the rollout schema (``id``, ``rollout_id``, + ``reward``, ``binary_payload`` as raw ``bytes``, ...). Only ``id`` and + ``rollout_id`` are required; every other field is optional. + """ + + def __init__(self, sync_store: _RolloutStore) -> None: + self._sync = sync_store + + @classmethod + def open( + cls, + uri: str, + *, + storage_options: Mapping[str, str] | None = None, + ) -> "RolloutStore": + """Open (or create) an embedded rollout dataset at ``uri``.""" + opts = dict(storage_options) if storage_options else None + return cls(_RolloutStore.open(uri, opts)) + + @classmethod + def connect(cls, base_url: str, name: str) -> "RolloutStore": + """Connect to an existing rollout store on a remote server.""" + return cls(_RolloutStore.connect(base_url, name)) + + @classmethod + def connect_or_create( + cls, + base_url: str, + name: str, + *, + storage_options: Mapping[str, str] | None = None, + ) -> "RolloutStore": + """Connect to a remote rollout store, creating it if absent.""" + opts = dict(storage_options) if storage_options else None + return cls(_RolloutStore.connect_or_create(base_url, name, opts)) + + def version(self) -> int: + """Return the current base-table version.""" + return self._sync.version() + + def add( + self, + records: Mapping[str, Any] | Iterable[Mapping[str, Any]], + ) -> dict[str, Any]: + """Append one record (dict) or many (iterable of dicts).""" + return self._sync.add(_rollout_records_to_json(records)) + + def add_one(self, **fields: Any) -> dict[str, Any]: + """Append a single record given as keyword arguments.""" + return self.add(fields) + + def list( + self, + limit: int | None = None, + offset: int | None = None, + ) -> list[dict[str, Any]]: + """List rollout rows (artifact bytes projected out; fetch via + :meth:`get_blob`).""" + raw = self._sync.list(limit, offset) + return [_rollout_record_from_json(r) for r in json.loads(raw)] + + def get(self, id: str) -> dict[str, Any] | None: + """Fetch a single rollout row by id, or ``None``.""" + raw = self._sync.get(id) + if raw is None: + return None + return _rollout_record_from_json(json.loads(raw)) + + def get_blob(self, id: str) -> bytes | None: + """Fetch a single artifact row's inline bytes on demand, or ``None``.""" + return self._sync.get_blob(id) + + def checkout(self, version: int) -> None: + """Checkout a base-table version (base-table time travel only).""" + self._sync.checkout(version) + + def __repr__(self) -> str: + return f"RolloutStore(version={self._sync.version()})" + + +class RemoteRolloutStore: + """Async wrapper around a remote rollout store over HTTP. + + Mirrors :class:`RolloutStore` but runs the blocking calls in an executor so + they can be awaited. Every method matches its sync counterpart. + """ + + def __init__(self, sync_store: _RolloutStore) -> None: + self._sync = sync_store + + @classmethod + async def connect(cls, base_url: str, name: str) -> "RemoteRolloutStore": + """Connect to an existing remote rollout store.""" + loop = asyncio.get_running_loop() + sync_store = await loop.run_in_executor( + None, lambda: _RolloutStore.connect(base_url, name) + ) + return cls(sync_store) + + @classmethod + async def connect_or_create( + cls, + base_url: str, + name: str, + *, + storage_options: Mapping[str, str] | None = None, + ) -> "RemoteRolloutStore": + """Connect to a remote rollout store, creating it if absent.""" + opts = dict(storage_options) if storage_options else None + loop = asyncio.get_running_loop() + sync_store = await loop.run_in_executor( + None, lambda: _RolloutStore.connect_or_create(base_url, name, opts) + ) + return cls(sync_store) + + def version(self) -> int: + """Return the current version (cached, sync).""" + return self._sync.version() + + async def add( + self, + records: Mapping[str, Any] | Iterable[Mapping[str, Any]], + ) -> dict[str, Any]: + """Append one record (dict) or many (iterable of dicts).""" + payload = _rollout_records_to_json(records) + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, lambda: self._sync.add(payload)) + + async def add_one(self, **fields: Any) -> dict[str, Any]: + """Append a single record given as keyword arguments.""" + return await self.add(fields) + + async def list( + self, + limit: int | None = None, + offset: int | None = None, + ) -> list[dict[str, Any]]: + """List rollout rows (artifact bytes projected out).""" + loop = asyncio.get_running_loop() + raw = await loop.run_in_executor(None, lambda: self._sync.list(limit, offset)) + return [_rollout_record_from_json(r) for r in json.loads(raw)] + + async def get(self, id: str) -> dict[str, Any] | None: + """Fetch a single rollout row by id, or ``None``.""" + loop = asyncio.get_running_loop() + raw = await loop.run_in_executor(None, lambda: self._sync.get(id)) + if raw is None: + return None + return _rollout_record_from_json(json.loads(raw)) + + async def get_blob(self, id: str) -> bytes | None: + """Fetch a single artifact row's bytes on demand, or ``None``.""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, lambda: self._sync.get_blob(id)) + + async def checkout(self, version: int) -> None: + """Checkout a base-table version (base-table time travel only).""" + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lambda: self._sync.checkout(version)) + + def __repr__(self) -> str: + return f"RemoteRolloutStore(version={self._sync.version()})" diff --git a/python/src/lib.rs b/python/src/lib.rs index 73cb468..fa9eaed 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -11,6 +11,10 @@ use pyo3::IntoPyObject; use serde_json::Value; use tokio::runtime::Runtime; +use lance_context::{ + AddRolloutRequest, CreateRolloutStoreRequest, RolloutRecordDto, + RolloutStore as UnifiedRolloutStore, RolloutStoreApi, +}; use lance_context_api::{ AddRecordRequest, CompactRequest, CompactResponse, CompactStatsResponse, ContextStoreApi, RecordDto, RecordPatchDto, RelationshipDto, RetrieveRequest, RetrieveResultDto, SearchRequest, @@ -2356,11 +2360,166 @@ fn dto_compact_stats_to_py(py: Python<'_>, stats: CompactStatsResponse) -> PyRes Ok(dict.into_pyobject(py)?.unbind().into()) } +// --------------------------------------------------------------------------- +// Rollout store binding (local + remote via the unified enum) +// --------------------------------------------------------------------------- + +/// A rollout store, either an embedded Lance dataset (`open`) or a handle to a +/// remote `lance-context-server` (`connect` / `connect_or_create`). +/// +/// Rollout rows carry ~35 fields, so records cross the FFI boundary as JSON +/// (a `records_json` array of objects matching `AddRolloutRequest`), and read +/// results come back as a JSON array string. The Python wrapper in `api.py` +/// hides this and exposes dict/list ergonomics. +#[pyclass] +struct RolloutStore { + store: UnifiedRolloutStore, + runtime: Arc, +} + +impl RolloutStore { + fn from_store(store: UnifiedRolloutStore, runtime: Arc) -> Self { + Self { store, runtime } + } +} + +#[pymethods] +impl RolloutStore { + /// Open (or create) an embedded rollout dataset at `uri`. + #[classmethod] + #[pyo3(signature = (uri, storage_options = None))] + fn open( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + uri: &str, + storage_options: Option>, + ) -> PyResult { + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedRolloutStore::open_with_options(uri, storage_options)) + }); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Connect to an existing rollout store on a remote server. + #[classmethod] + fn connect( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + ) -> PyResult { + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = + py.allow_threads(|| runtime.block_on(UnifiedRolloutStore::connect(base_url, name))); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Connect to a remote rollout store, creating it if it does not exist. + #[classmethod] + #[pyo3(signature = (base_url, name, storage_options = None))] + fn connect_or_create( + _cls: &Bound<'_, PyType>, + py: Python<'_>, + base_url: &str, + name: &str, + storage_options: Option>, + ) -> PyResult { + let req = CreateRolloutStoreRequest { + name: name.to_string(), + storage_options, + }; + let runtime = Arc::new(Runtime::new().map_err(to_py_err)?); + let store_res = py.allow_threads(|| { + runtime.block_on(UnifiedRolloutStore::connect_or_create(base_url, &req)) + }); + let store = store_res.map_err(to_py_err)?; + Ok(Self::from_store(store, runtime)) + } + + /// Current store version (base dataset version). + fn version(&self) -> u64 { + self.store.version() + } + + /// 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 { + let records: Vec = serde_json::from_str(records_json) + .map_err(|e| PyRuntimeError::new_err(format!("invalid records JSON: {e}")))?; + if records.is_empty() { + return Err(PyRuntimeError::new_err( + "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 dict = PyDict::new(py); + dict.set_item("version", resp.version)?; + dict.set_item("ids", resp.ids)?; + dict.set_item("count", resp.count)?; + Ok(dict.into_pyobject(py)?.unbind().into()) + } + + /// 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))] + fn list( + &self, + py: Python<'_>, + limit: Option, + offset: Option, + ) -> PyResult { + let records = py + .allow_threads(|| self.runtime.block_on(self.store.list(limit, offset))) + .map_err(to_py_err)?; + rollout_records_to_json(&records) + } + + /// Fetch a single rollout row by id, or `None`. Returns a JSON object + /// string (artifact bytes projected out). + fn get(&self, py: Python<'_>, id: &str) -> PyResult> { + let record = py + .allow_threads(|| self.runtime.block_on(self.store.get(id))) + .map_err(to_py_err)?; + match record { + Some(r) => Ok(Some(rollout_record_to_json(&r)?)), + None => Ok(None), + } + } + + /// Fetch a single artifact row's inline bytes on demand, or `None`. + fn get_blob(&self, py: Python<'_>, id: &str) -> PyResult>> { + let bytes = py + .allow_threads(|| self.runtime.block_on(self.store.get_blob(id))) + .map_err(to_py_err)?; + Ok(bytes.map(|b| PyBytes::new(py, &b).unbind())) + } + + /// Checkout a base-table version (time travel over the base table only). + fn checkout(&mut self, py: Python<'_>, version: u64) -> PyResult<()> { + py.allow_threads(|| self.runtime.block_on(self.store.checkout(version))) + .map_err(to_py_err) + } +} + +fn rollout_records_to_json(records: &[RolloutRecordDto]) -> PyResult { + serde_json::to_string(records).map_err(to_py_err) +} + +fn rollout_record_to_json(record: &RolloutRecordDto) -> PyResult { + serde_json::to_string(record).map_err(to_py_err) +} + #[pymodule] fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(version, m)?)?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/python/tests/test_rollout.py b/python/tests/test_rollout.py new file mode 100644 index 0000000..f196f46 --- /dev/null +++ b/python/tests/test_rollout.py @@ -0,0 +1,107 @@ +"""Tests for the rollout store Python bindings (local/embedded path). + +The embedded `RolloutStore.open` path exercises the exact same binding surface +(JSON marshalling of records/DTOs, base64 blob handling, projection) that the +remote path uses, since both dispatch through the unified Rust `RolloutStore` +enum. Remote-specific wiring (HTTP) is covered by the Rust client tests. +""" + +from __future__ import annotations + +import tempfile + +import pytest +from lance_context import RolloutStore + + +@pytest.fixture() +def store_uri(): + with tempfile.TemporaryDirectory() as d: + yield d + + +def test_add_dict_list_get(store_uri): + store = RolloutStore.open(store_uri) + assert store.version() >= 0 + + resp = store.add( + { + "id": "row-0", + "rollout_id": "traj-1", + "problem_id": "p-1", + "role": "assistant", + "content": "the answer is 42", + "reward": 1.0, + "policy_version": "ckpt-7", + } + ) + assert resp["count"] == 1 + assert resp["ids"] == ["row-0"] + + rows = store.list() + assert len(rows) == 1 + assert rows[0]["id"] == "row-0" + assert rows[0]["rollout_id"] == "traj-1" + assert rows[0]["reward"] == 1.0 + assert rows[0]["policy_version"] == "ckpt-7" + + fetched = store.get("row-0") + assert fetched is not None + assert fetched["id"] == "row-0" + assert store.get("missing") is None + + +def test_add_many_and_add_one(store_uri): + store = RolloutStore.open(store_uri) + store.add( + [ + {"id": f"row-{i}", "rollout_id": "traj-1", "reward": float(i)} + for i in range(3) + ] + ) + store.add_one(id="row-3", rollout_id="traj-1", reward=3.0) + + rows = store.list() + assert {r["id"] for r in rows} == {"row-0", "row-1", "row-2", "row-3"} + + +def test_binary_payload_roundtrip(store_uri): + store = RolloutStore.open(store_uri) + blob = b"\x00\x01\x02trace-bytes" + store.add( + { + "id": "art-0", + "rollout_id": "traj-1", + "role": "artifact", + "content_type": "application/octet-stream", + "binary_payload": blob, + "payload_size": len(blob), + } + ) + + # list/get project the blob column out (cheap metadata scans). + row = store.get("art-0") + assert row is not None + assert row.get("binary_payload") is None + assert row["payload_size"] == len(blob) + + # get_blob materializes the bytes on demand. + assert store.get_blob("art-0") == blob + assert store.get_blob("missing") is None + + +def test_empty_add_rejected(store_uri): + store = RolloutStore.open(store_uri) + with pytest.raises(ValueError): + store.add([]) + + +def test_reopen_sees_prior_rows(store_uri): + store = RolloutStore.open(store_uri) + store.add({"id": "row-0", "rollout_id": "traj-1"}) + del store + + reopened = RolloutStore.open(store_uri) + rows = reopened.list() + assert len(rows) == 1 + assert rows[0]["id"] == "row-0" diff --git a/python/tests/test_rollout_remote.py b/python/tests/test_rollout_remote.py new file mode 100644 index 0000000..2ae4b50 --- /dev/null +++ b/python/tests/test_rollout_remote.py @@ -0,0 +1,131 @@ +"""Remote (HTTP) rollout store tests. + +Spins up a real `lance-context-server` subprocess and drives it through the +async `RemoteRolloutStore` wrapper — the path RL generation workers and the +learner use in a deployed setup. Skips gracefully if the server binary has not +been built (e.g. a pure-Python CI job). +""" + +from __future__ import annotations + +import asyncio +import socket +import subprocess +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path + +import pytest +from lance_context import RemoteRolloutStore + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_SERVER_BIN = _REPO_ROOT / "target" / "debug" / "lance-context-server" + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _wait_for_health(base_url: str, timeout: float = 30.0) -> None: + deadline = time.time() + timeout + url = f"{base_url}/api/v1/health" + while time.time() < deadline: + try: + with urllib.request.urlopen(url, timeout=1) as resp: # noqa: S310 + if resp.status == 200: + return + except (urllib.error.URLError, ConnectionError, OSError): + time.sleep(0.2) + raise RuntimeError(f"server did not become healthy at {url}") + + +@pytest.fixture() +def server(): + if not _SERVER_BIN.exists(): + pytest.skip(f"server binary not built at {_SERVER_BIN}") + port = _free_port() + with tempfile.TemporaryDirectory() as data_dir: + proc = subprocess.Popen( + [ + str(_SERVER_BIN), + "--host", + "127.0.0.1", + "--port", + str(port), + "--data-dir", + data_dir, + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + base_url = f"http://127.0.0.1:{port}" + try: + _wait_for_health(base_url) + yield base_url + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +def test_remote_roundtrip(server): + async def run(): + store = await RemoteRolloutStore.connect_or_create(server, "rl-run-1") + + resp = await store.add( + [ + { + "id": "row-0", + "rollout_id": "traj-1", + "problem_id": "p-1", + "role": "assistant", + "content": "the answer is 42", + "reward": 1.0, + "policy_version": "ckpt-7", + }, + { + "id": "row-1", + "rollout_id": "traj-1", + "role": "artifact", + "content_type": "application/octet-stream", + "binary_payload": b"\x00\x01\x02trace", + "payload_size": 8, + }, + ] + ) + assert resp["count"] == 2 + + rows = await store.list() + assert {r["id"] for r in rows} == {"row-0", "row-1"} + + one = await store.get("row-0") + assert one is not None + assert one["reward"] == 1.0 + # blob projected out of get, materialized on demand + assert one.get("binary_payload") is None + + blob = await store.get_blob("row-1") + assert blob == b"\x00\x01\x02trace" + assert await store.get("missing") is None + + asyncio.run(run()) + + +def test_remote_add_one_and_connect(server): + async def run(): + store = await RemoteRolloutStore.connect_or_create(server, "rl-run-2") + await store.add_one(id="only", rollout_id="traj-9", reward=0.5) + + # A second connection sees the first's flushed write (durable, no + # read affinity). + reader = await RemoteRolloutStore.connect(server, "rl-run-2") + rows = await reader.list() + assert [r["id"] for r in rows] == ["only"] + + asyncio.run(run()) From 6e4e62a7156fba01fa4ae838b83d30697cd8d661 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Fri, 10 Jul 2026 20:32:45 +0000 Subject: [PATCH 2/2] feat(rollout): auto-generate id + first-class artifact_type column - Mint a UUID4 for a rollout record when `id` is omitted/empty, so callers never have to hand-roll ids (applies to both add and add_one, local + remote). - Add `artifact_type` as a first-class nullable Utf8 column, orthogonal to `content_type`: content_type is the transport/media type (e.g. image/png), artifact_type is the user-defined semantic category (e.g. excel_grade_screenshot). A real column so it can be filtered / grouped-by / projected without materializing the free-form metadata JSON. Wired through RolloutRecord, Arrow schema (read+write), the API DTOs (AddRolloutRequest, RolloutRecordDto), the server routes, and the Python bindings (which pass dicts through generically, so no field allowlist change). Co-Authored-By: Claude Opus 4 --- crates/lance-context-api/src/lib.rs | 6 +++++ crates/lance-context-core/src/api_impl.rs | 2 ++ crates/lance-context-core/src/rollout.rs | 10 +++++++- .../lance-context-core/src/rollout_store.rs | 14 ++++++++++- .../src/routes/rollouts.rs | 2 ++ python/python/lance_context/api.py | 9 ++++++- python/tests/test_rollout.py | 24 ++++++++++++++++++- 7 files changed, 63 insertions(+), 4 deletions(-) diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 158e6f4..b57d632 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -666,6 +666,8 @@ pub struct AddRolloutRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub payload_checksum: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub artifact_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub metadata: Option, } @@ -741,6 +743,10 @@ pub struct RolloutRecordDto { pub payload_size: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub payload_checksum: Option, + /// User-defined semantic artifact category (e.g. `"excel_grade_screenshot"`), + /// orthogonal to `content_type`. A first-class, filterable column. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub artifact_type: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub metadata: Option, } diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 5e83f32..96084a6 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -458,6 +458,7 @@ fn rollout_record_from_add_request(r: &AddRolloutRequest) -> RolloutRecord { binary_payload: r.binary_payload.clone(), payload_size: r.payload_size, payload_checksum: r.payload_checksum.clone(), + artifact_type: r.artifact_type.clone(), metadata: r.metadata.clone(), } } @@ -497,6 +498,7 @@ fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { binary_payload: r.binary_payload, payload_size: r.payload_size, payload_checksum: r.payload_checksum, + artifact_type: r.artifact_type, metadata: r.metadata, } } diff --git a/crates/lance-context-core/src/rollout.rs b/crates/lance-context-core/src/rollout.rs index 12370d9..a6db415 100644 --- a/crates/lance-context-core/src/rollout.rs +++ b/crates/lance-context-core/src/rollout.rs @@ -84,8 +84,16 @@ pub struct RolloutRecord { pub binary_payload: Option>, pub payload_size: Option, pub payload_checksum: Option, + /// User-defined semantic category of an artifact, e.g. + /// `"excel_grade_screenshot"`. Orthogonal to `content_type`, which is the + /// transport/media type (e.g. `"image/png"`): `content_type` says how to + /// decode the bytes, `artifact_type` says what the artifact *means*. A + /// first-class column so it can be filtered / grouped-by / projected + /// without materializing the free-form `metadata` JSON. + pub artifact_type: Option, /// Harness metadata — the open-ended escape hatch, for genuinely - /// unstructured fields only (e.g. an artifact's `filename` / `artifact_type`). + /// unstructured fields only (e.g. an artifact's `filename`). Semantic + /// categories that you filter/group-by belong in `artifact_type` instead. pub metadata: Option, } diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 9a41a54..c3788a4 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -423,6 +423,7 @@ impl RolloutStore { let mut binary_payload_builder = LargeBinaryBuilder::new(); let mut payload_size_builder = Int64Builder::new(); let mut payload_checksum_builder = StringBuilder::new(); + let mut artifact_type_builder = StringBuilder::new(); let mut metadata_builder = LargeStringBuilder::new(); for record in records { @@ -482,6 +483,7 @@ impl RolloutStore { } payload_size_builder.append_option(record.payload_size); payload_checksum_builder.append_option(record.payload_checksum.as_deref()); + artifact_type_builder.append_option(record.artifact_type.as_deref()); match &record.metadata { Some(metadata) => metadata_builder.append_value(metadata.to_string()), None => metadata_builder.append_null(), @@ -589,6 +591,10 @@ impl RolloutStore { "payload_checksum".to_string(), Arc::new(payload_checksum_builder.finish()), ); + arrays_by_name.insert( + "artifact_type".to_string(), + Arc::new(artifact_type_builder.finish()), + ); if include_metadata { arrays_by_name.insert("metadata".to_string(), Arc::new(metadata_builder.finish())); } @@ -682,6 +688,7 @@ pub fn rollout_schema() -> Schema { binary_field, Field::new("payload_size", DataType::Int64, true), Field::new("payload_checksum", DataType::Utf8, true), + Field::new("artifact_type", DataType::Utf8, true), Field::new("metadata", DataType::LargeUtf8, true), ]; @@ -767,6 +774,7 @@ fn batch_to_rollout_records(batch: &RecordBatch) -> LanceResult(batch, "binary_payload"); let payload_size_array = column_as_optional::(batch, "payload_size"); let payload_checksum_array = column_as_optional::(batch, "payload_checksum"); + let artifact_type_array = column_as_optional::(batch, "artifact_type"); let metadata_array = column_as_optional::(batch, "metadata"); let mut results = Vec::with_capacity(batch.num_rows()); @@ -855,6 +863,7 @@ fn batch_to_rollout_records(batch: &RecordBatch) -> LanceResult RolloutRecord { binary_payload: r.binary_payload.clone(), payload_size: r.payload_size, payload_checksum: r.payload_checksum.clone(), + artifact_type: r.artifact_type.clone(), metadata: r.metadata.clone(), } } @@ -482,6 +483,7 @@ fn rollout_record_to_dto(r: RolloutRecord) -> RolloutRecordDto { binary_payload: r.binary_payload, payload_size: r.payload_size, payload_checksum: r.payload_checksum, + artifact_type: r.artifact_type, metadata: r.metadata, } } diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index ea8d395..d9fedeb 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -2458,9 +2458,16 @@ def _rollout_record_to_native(record: Mapping[str, Any]) -> dict[str, Any]: """Normalize one rollout record dict for the JSON FFI boundary. Accepts ``binary_payload`` as raw ``bytes``/``bytearray`` and base64-encodes - it (the DTO's JSON representation), leaving everything else untouched. + it (the DTO's JSON representation). If ``id`` is missing/empty, a UUID4 is + generated. Everything else is left untouched. """ out = dict(record) + # `id` is required by the store; if the caller omits it (or passes an empty + # value) generate a UUID4 so users never have to mint ids by hand. + if not out.get("id"): + import uuid + + out["id"] = str(uuid.uuid4()) blob = out.get(_ROLLOUT_BLOB_FIELD) if isinstance(blob, (bytes, bytearray)): import base64 diff --git a/python/tests/test_rollout.py b/python/tests/test_rollout.py index f196f46..028b3e8 100644 --- a/python/tests/test_rollout.py +++ b/python/tests/test_rollout.py @@ -65,6 +65,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_id_autogenerated_when_omitted(store_uri): + import uuid + + store = RolloutStore.open(store_uri) + # No `id` field -> the binding mints a UUID4 for us. + resp = store.add({"rollout_id": "traj-1", "reward": 1.0}) + assert resp["count"] == 1 + (generated,) = resp["ids"] + # Round-trips as a valid UUID and is fetchable by the generated id. + assert uuid.UUID(generated).version == 4 + row = store.get(generated) + assert row is not None and row["id"] == generated + + # add_one without id also works and yields distinct ids. + r2 = store.add_one(rollout_id="traj-1") + assert r2["ids"][0] != generated + + def test_binary_payload_roundtrip(store_uri): store = RolloutStore.open(store_uri) blob = b"\x00\x01\x02trace-bytes" @@ -73,7 +91,8 @@ def test_binary_payload_roundtrip(store_uri): "id": "art-0", "rollout_id": "traj-1", "role": "artifact", - "content_type": "application/octet-stream", + "content_type": "image/png", + "artifact_type": "excel_grade_screenshot", "binary_payload": blob, "payload_size": len(blob), } @@ -84,6 +103,9 @@ def test_binary_payload_roundtrip(store_uri): assert row is not None assert row.get("binary_payload") is None assert row["payload_size"] == len(blob) + # artifact_type is a first-class column, orthogonal to content_type. + assert row["content_type"] == "image/png" + assert row["artifact_type"] == "excel_grade_screenshot" # get_blob materializes the bytes on demand. assert store.get_blob("art-0") == blob