feat(server): scale rollout stores via registry + LRU + global sweeper

crates/lance-context-api/src/lib.rs

@@ -580,7 +580,12 @@ pub struct CreateRolloutStoreRequest {
 pub struct RolloutStoreInfo {
     pub name: String,
     pub uri: String,
-    pub version: u64,
+    /// Dataset version. `None` in list responses, which are served from the
+    /// durable registry without opening each dataset (a store may not be
+    /// resident in memory). Single-store lookups (`get`/`create`) always
+    /// populate it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub version: Option<u64>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/lance-context-client/src/lib.rs

@@ -264,7 +264,7 @@ impl RemoteRolloutStore {
         Ok(Self {
             client,
             store_name: store_name.to_string(),
-            cached_version: info.version,
+            cached_version: info.version.unwrap_or(0),
         })
     }
 
@@ -281,7 +281,7 @@ impl RemoteRolloutStore {
         Ok(Self {
             client,
             store_name: req.name.clone(),
-            cached_version: info.version,
+            cached_version: info.version.unwrap_or(0),
         })
     }
 }

crates/lance-context-core/src/lib.rs

@@ -8,6 +8,7 @@ mod export;
 mod id;
 mod namespace;
 mod record;
+mod registry;
 mod rollout;
 mod rollout_store;
 pub mod serde;
@@ -31,6 +32,7 @@ pub use record::{
     RetrieveResult, SearchResult, StateMetadata, UpdateResult, UpsertResult, LIFECYCLE_ACTIVE,
     LIFECYCLE_CONTRADICTED,
 };
+pub use registry::{RegistryEntry, RolloutRegistry};
 pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL};
 pub use rollout_store::{rollout_schema, RolloutStore, RolloutStoreOptions};
 pub use store::{

crates/lance-context-core/src/registry.rs

@@ -0,0 +1,310 @@
+//! Persistent directory of rollout stores.
+//!
+//! # Why this exists
+//!
+//! With `experiment_name` used as the physical partition key, a deployment may
+//! hold **tens or hundreds of thousands** of independent rollout datasets — one
+//! per experiment. The server can no longer keep every opened store resident in
+//! memory (it uses a bounded LRU), so "is this store in the in-memory map?" is
+//! no longer a valid existence check, and `list` can no longer enumerate stores
+//! by walking the cache.
+//!
+//! [`RolloutRegistry`] is a single, small Lance dataset — one row per rollout
+//! store (`name`, `uri`, `created_at`) — that serves as the durable source of
+//! truth for *which* stores exist. It is consulted on a cache miss to decide
+//! whether to lazily open a store (vs. return 404), and it backs the `list`
+//! endpoint without touching object storage for every experiment.
+//!
+//! It deliberately holds only cheap directory metadata, never rollout rows.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator, StringArray};
+use arrow_schema::{ArrowError, DataType, Field, Schema};
+use chrono::Utc;
+use futures::TryStreamExt;
+use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams};
+use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
+use lance::{Error as LanceError, Result as LanceResult};
+
+/// One entry in the rollout-store directory.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RegistryEntry {
+    /// Logical store name (the experiment name); unique key.
+    pub name: String,
+    /// Physical dataset URI, as produced by the server's `rollout_uri`.
+    pub uri: String,
+    /// Creation time, Unix milliseconds.
+    pub created_at: i64,
+}
+
+/// Durable directory of rollout stores, backed by a single Lance dataset.
+///
+/// All mutations (`upsert`, `remove`) take `&mut self` and are expected to be
+/// serialized by the caller (the server wraps this in a lock), so the registry
+/// itself does no internal locking. Reads (`list`, `contains`) take `&self`.
+pub struct RolloutRegistry {
+    dataset: Dataset,
+    uri: String,
+    storage_options: Option<HashMap<String, String>>,
+}
+
+fn registry_schema() -> Schema {
+    Schema::new(vec![
+        Field::new("name", DataType::Utf8, false),
+        Field::new("uri", DataType::Utf8, false),
+        Field::new("created_at", DataType::Int64, false),
+    ])
+}
+
+impl RolloutRegistry {
+    /// Open the registry dataset at `uri`, creating an empty one if it does not
+    /// yet exist. Idempotent across process restarts.
+    pub async fn open_or_create(
+        uri: &str,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> LanceResult<Self> {
+        let dataset = match Self::load(uri, storage_options.clone()).await {
+            Ok(dataset) => dataset,
+            Err(LanceError::DatasetNotFound { .. }) => {
+                Self::create(uri, storage_options.clone()).await?
+            }
+            Err(err) => return Err(err),
+        };
+        Ok(Self {
+            dataset,
+            uri: uri.to_string(),
+            storage_options,
+        })
+    }
+
+    async fn load(
+        uri: &str,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> LanceResult<Dataset> {
+        if let Some(options) = storage_options {
+            DatasetBuilder::from_uri(uri)
+                .with_storage_options(options)
+                .load()
+                .await
+        } else {
+            Dataset::open(uri).await
+        }
+    }
+
+    async fn create(
+        uri: &str,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> LanceResult<Dataset> {
+        let schema = Arc::new(registry_schema());
+        let empty = RecordBatch::new_empty(schema.clone());
+        let batches = RecordBatchIterator::new(
+            vec![Ok::<RecordBatch, ArrowError>(empty)].into_iter(),
+            schema.clone(),
+        );
+        let params = Self::write_params(WriteMode::Create, storage_options);
+        Dataset::write(batches, uri, Some(params)).await
+    }
+
+    fn write_params(
+        mode: WriteMode,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> WriteParams {
+        let mut params = WriteParams {
+            mode,
+            ..Default::default()
+        };
+        if let Some(options) = storage_options {
+            params.store_params = Some(ObjectStoreParams {
+                storage_options_accessor: Some(Arc::new(
+                    StorageOptionsAccessor::with_static_options(options),
+                )),
+                ..Default::default()
+            });
+        }
+        params
+    }
+
+    /// Insert or replace the directory entry for `name`.
+    ///
+    /// Implemented as delete-same-name-then-append so it is **idempotent**: a
+    /// `create` retried after a crash (dataset on disk, registry row possibly
+    /// present) converges to exactly one row. Callers must serialize mutations.
+    pub async fn upsert(&mut self, name: &str, uri: &str) -> LanceResult<()> {
+        self.delete_row(name).await?;
+        let schema = Arc::new(registry_schema());
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(StringArray::from(vec![name])),
+                Arc::new(StringArray::from(vec![uri])),
+                Arc::new(Int64Array::from(vec![Utc::now().timestamp_millis()])),
+            ],
+        )?;
+        let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema);
+        let params = Self::write_params(WriteMode::Append, self.storage_options.clone());
+        self.dataset.append(reader, Some(params)).await?;
+        Ok(())
+    }
+
+    /// Remove the directory entry for `name`, if present. No-op when absent.
+    pub async fn remove(&mut self, name: &str) -> LanceResult<()> {
+        self.delete_row(name).await
+    }
+
+    async fn delete_row(&mut self, name: &str) -> LanceResult<()> {
+        let escaped = name.replace('\'', "''");
+        self.dataset
+            .delete(&format!("name = '{}'", escaped))
+            .await?;
+        Ok(())
+    }
+
+    /// Whether a store named `name` exists in the directory.
+    pub async fn contains(&self, name: &str) -> LanceResult<bool> {
+        let escaped = name.replace('\'', "''");
+        let mut scanner = self.dataset.scan();
+        scanner.project(&["name"])?;
+        scanner.filter(&format!("name = '{}'", escaped))?;
+        scanner.limit(Some(1), None)?;
+        let mut stream = scanner.try_into_stream().await?;
+        while let Some(batch) = stream.try_next().await? {
+            if batch.num_rows() > 0 {
+                return Ok(true);
+            }
+        }
+        Ok(false)
+    }
+
+    /// All directory entries, ordered as stored (unspecified). The registry is a
+    /// narrow three-column table, so even hundreds of thousands of rows scan
+    /// quickly; pagination can be layered on later if needed.
+    pub async fn list(&self) -> LanceResult<Vec<RegistryEntry>> {
+        let mut scanner = self.dataset.scan();
+        scanner.project(&["name", "uri", "created_at"])?;
+        let mut stream = scanner.try_into_stream().await?;
+        let mut out = Vec::new();
+        while let Some(batch) = stream.try_next().await? {
+            let names = batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .ok_or_else(|| {
+                    LanceError::from(ArrowError::InvalidArgumentError(
+                        "registry 'name' column is not Utf8".into(),
+                    ))
+                })?;
+            let uris = batch
+                .column(1)
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .ok_or_else(|| {
+                    LanceError::from(ArrowError::InvalidArgumentError(
+                        "registry 'uri' column is not Utf8".into(),
+                    ))
+                })?;
+            let created = batch
+                .column(2)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .ok_or_else(|| {
+                    LanceError::from(ArrowError::InvalidArgumentError(
+                        "registry 'created_at' column is not Int64".into(),
+                    ))
+                })?;
+            for i in 0..batch.num_rows() {
+                out.push(RegistryEntry {
+                    name: names.value(i).to_string(),
+                    uri: uris.value(i).to_string(),
+                    created_at: created.value(i),
+                });
+            }
+        }
+        Ok(out)
+    }
+
+    /// The registry dataset URI.
+    #[must_use]
+    pub fn uri(&self) -> &str {
+        &self.uri
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::TempDir;
+
+    async fn new_registry(dir: &TempDir) -> RolloutRegistry {
+        let uri = dir.path().join("_registry.rollout.lance");
+        RolloutRegistry::open_or_create(uri.to_str().unwrap(), None)
+            .await
+            .unwrap()
+    }
+
+    #[tokio::test]
+    async fn upsert_is_idempotent() {
+        let dir = TempDir::new().unwrap();
+        let mut reg = new_registry(&dir).await;
+        reg.upsert("exp-a", "/data/exp-a.rollout.lance")
+            .await
+            .unwrap();
+        // Re-upserting the same name must not create a duplicate row.
+        reg.upsert("exp-a", "/data/exp-a.rollout.lance")
+            .await
+            .unwrap();
+        let entries = reg.list().await.unwrap();
+        assert_eq!(entries.len(), 1);
+        assert_eq!(entries[0].name, "exp-a");
+        assert!(reg.contains("exp-a").await.unwrap());
+    }
+
+    #[tokio::test]
+    async fn contains_and_remove() {
+        let dir = TempDir::new().unwrap();
+        let mut reg = new_registry(&dir).await;
+        assert!(!reg.contains("missing").await.unwrap());
+        reg.upsert("exp-b", "/data/exp-b.rollout.lance")
+            .await
+            .unwrap();
+        assert!(reg.contains("exp-b").await.unwrap());
+        reg.remove("exp-b").await.unwrap();
+        assert!(!reg.contains("exp-b").await.unwrap());
+        assert!(reg.list().await.unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn list_returns_all_entries() {
+        let dir = TempDir::new().unwrap();
+        let mut reg = new_registry(&dir).await;
+        for i in 0..5 {
+            reg.upsert(&format!("exp-{i}"), &format!("/data/exp-{i}.rollout.lance"))
+                .await
+                .unwrap();
+        }
+        let mut names: Vec<String> = reg
+            .list()
+            .await
+            .unwrap()
+            .into_iter()
+            .map(|e| e.name)
+            .collect();
+        names.sort();
+        assert_eq!(names, vec!["exp-0", "exp-1", "exp-2", "exp-3", "exp-4"]);
+    }
+
+    #[tokio::test]
+    async fn survives_reopen() {
+        let dir = TempDir::new().unwrap();
+        {
+            let mut reg = new_registry(&dir).await;
+            reg.upsert("persist", "/data/persist.rollout.lance")
+                .await
+                .unwrap();
+        }
+        // Reopen the same path: the entry must still be present.
+        let reg = new_registry(&dir).await;
+        assert!(reg.contains("persist").await.unwrap());
+    }
+}

crates/lance-context-server/Cargo.toml

@@ -18,6 +18,7 @@ lance-context-api = { version = "0.6.3", path = "../lance-context-api" }
 axum = { version = "0.8", features = ["json", "multipart"] }
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
 clap = { version = "4", features = ["derive", "env"] }
+lru = "0.12"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] }

crates/lance-context-server/src/config.rs

@@ -50,6 +50,15 @@ pub struct ServerConfig {
     /// set. Defaults to `1` (clean up whenever any generation is present).
     #[arg(long, env = "ROLLOUT_CLEANUP_MIN_GENERATIONS", default_value = "1")]
     pub rollout_cleanup_min_generations: usize,
+
+    /// Upper bound on the number of resident rollout-store handles kept in
+    /// memory (an LRU). With one physical dataset per experiment a deployment
+    /// may hold hundreds of thousands of stores; this bounds how many stay open
+    /// at once. Size it for peak *concurrent* experiments, not the total count.
+    /// Evicted stores are transparently reopened on the next request. `0` falls
+    /// back to the built-in default.
+    #[arg(long, env = "ROLLOUT_CACHE_CAPACITY", default_value = "2000")]
+    pub rollout_cache_capacity: usize,
 }
 
 impl ServerConfig {