From 4927d3b69077c311a728dfd73309e91ae9818433 Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 15 Jul 2026 04:56:06 +0000 Subject: [PATCH 1/3] fix(rollout): make WAL time trigger a true OR with count trigger The periodic time-based WAL cleanup was gated by a second count threshold (cleanup_min_generations), so "interval elapsed" would not merge unless enough generations had also accumulated. That turned the intended "time OR count, whichever fires first" trigger into an AND: a shard whose interval elapsed but stayed below the count threshold never merged. Make the time trigger unconditional on count: cleanup_own_shard now merges whatever is pending (>=1 generation) the moment a pass runs. Remove the cleanup_min_generations option, the rollout_cleanup_min_generations server config, and the ROLLOUT_CLEANUP_MIN_GENERATIONS env/CLI flag. The count trigger (merge_after_generations) is unchanged; the two are now a strict OR. Co-Authored-By: Claude Opus 4 --- .../lance-context-core/src/rollout_store.rs | 75 +++++++++---------- crates/lance-context-server/src/config.rs | 16 ++-- .../src/routes/rollouts.rs | 1 - crates/lance-context-server/src/state.rs | 6 -- 4 files changed, 42 insertions(+), 56 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 2abb3a1..d3cfbda 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -148,13 +148,13 @@ pub struct RolloutStoreOptions { /// /// Spawn the task with [`RolloutStore::spawn_periodic_cleanup`] once the /// store is behind an `Arc>` (the server's ownership model). + /// + /// The two triggers are a strict **OR**: whichever fires first merges. The + /// time trigger is *not* gated by any generation count — once the interval + /// elapses, every pending generation is folded into the base table even if + /// the count trigger's threshold was never crossed. (A pass with nothing + /// pending is a no-op.) pub cleanup_interval_secs: Option, - /// The *count* half of the periodic cleanup trigger: the timer only merges - /// when this instance's shard has at least this many flushed generations, - /// skipping the pass otherwise (avoids rewriting the base table to reclaim a - /// single small generation). `None` defaults to `1` (clean up whenever any - /// generation is present on a tick). - pub cleanup_min_generations: Option, } /// A Lance-backed store for RL rollout trajectories. @@ -171,10 +171,6 @@ pub struct RolloutStore { /// Periodic-cleanup interval in seconds; `0` disables the timer. See /// [`RolloutStoreOptions::cleanup_interval_secs`]. cleanup_interval_secs: u64, - /// Minimum flushed generations before a periodic-cleanup tick merges. - /// Normalized to at least `1`. See - /// [`RolloutStoreOptions::cleanup_min_generations`]. - cleanup_min_generations: usize, /// Timestamp of the last successful [`Self::compact`] on this handle. last_compaction: Option>, /// Number of successful compactions performed by this handle. @@ -236,7 +232,6 @@ impl RolloutStore { storage_options, merge_after_generations: options.merge_after_generations.unwrap_or(0), cleanup_interval_secs: options.cleanup_interval_secs.unwrap_or(0), - cleanup_min_generations: options.cleanup_min_generations.unwrap_or(1).max(1), last_compaction: None, total_compactions: 0, last_compaction_error: None, @@ -320,9 +315,10 @@ impl RolloutStore { /// /// This is the shared core of both triggers: the synchronous count trigger /// in [`Self::add`] (with `threshold = merge_after_generations`) and the - /// periodic timer in [`Self::spawn_periodic_cleanup`] (with - /// `threshold = cleanup_min_generations`). Both merge only the shard this - /// instance owns and writes, so the epoch claim never fences another writer. + /// periodic timer in [`Self::spawn_periodic_cleanup`] (with `threshold = 1`, + /// i.e. merge whatever is pending once the interval elapses). Both merge only + /// the shard this instance owns and writes, so the epoch claim never fences + /// another writer. async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { let object_store = self.dataset.object_store(None).await?; let branch_location = self.dataset.branch_location(); @@ -343,10 +339,13 @@ impl RolloutStore { Ok(pending) } - /// Run one periodic WAL-cleanup pass over this instance's own shard: fold any - /// flushed generations into the base table once at least - /// `cleanup_min_generations` have accumulated. Returns the number of - /// generations reclaimed (`0` if below the threshold or nothing pending). + /// Run one periodic WAL-cleanup pass over this instance's own shard: fold + /// **every** flushed generation into the base table. This is the *time* half + /// of the "time OR count" trigger, so it is deliberately *not* gated by any + /// generation-count threshold: once the interval elapses, whatever is pending + /// gets merged even if the count trigger + /// ([`RolloutStoreOptions::merge_after_generations`]) never fired. Returns the + /// number of generations reclaimed (`0` only when nothing was pending). /// /// Exposed for callers that drive cleanup on their own schedule instead of /// (or in addition to) the built-in timer from @@ -355,18 +354,21 @@ impl RolloutStore { /// with this instance's own appends but must not target another instance's /// shard. pub async fn cleanup_own_shard(&mut self) -> LanceResult { - self.merge_own_shard_if_ready(self.cleanup_min_generations) - .await + // Threshold `1`: merge whenever at least one generation is pending. The + // time trigger must not depend on the count threshold — that is what + // makes the two triggers a true OR. + self.merge_own_shard_if_ready(1).await } /// Spawn a background timer that periodically reclaims this instance's /// flushed MemWAL generations into the base table (the *time* trigger). /// /// Every `cleanup_interval_secs` seconds the task acquires the write lock and - /// calls [`Self::cleanup_own_shard`], which merges only when at least - /// `cleanup_min_generations` are pending. This bounds read amplification and - /// reclaims stale generation datasets even on shards that never cross the - /// synchronous count threshold ([`RolloutStoreOptions::merge_after_generations`]). + /// calls [`Self::cleanup_own_shard`], which merges whatever generations are + /// pending (no count threshold — time and count are a strict OR). This bounds + /// read amplification and reclaims stale generation datasets even on shards + /// that never cross the synchronous count threshold + /// ([`RolloutStoreOptions::merge_after_generations`]). /// /// Returns `None` (spawning nothing) when `cleanup_interval_secs` is `0` /// (disabled). Otherwise returns the task handle. @@ -1952,7 +1954,6 @@ mod tests { RolloutStoreOptions { storage_options: None, shard_id: Some("rollout-1".to_string()), - cleanup_min_generations: Some(1), ..Default::default() }, ) @@ -2106,11 +2107,12 @@ mod tests { } #[test] - fn cleanup_own_shard_merges_only_at_min_generations() { + fn cleanup_own_shard_merges_whatever_is_pending() { // The periodic-cleanup entry point (`cleanup_own_shard`) is the time - // trigger's per-pass body. With count self-merge disabled, generations - // accumulate; a cleanup pass below `cleanup_min_generations` is a no-op, - // and one at/above the threshold drains the shard into the base table. + // trigger's per-pass body. Time and count are a strict OR, so the time + // trigger is NOT gated by any generation count: with count self-merge + // disabled, a single pending generation is merged the moment the pass + // runs. A pass with nothing pending is a no-op. let dir = TempDir::new().unwrap(); let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); @@ -2121,7 +2123,6 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // count trigger off - cleanup_min_generations: Some(2), ..Default::default() }, ) @@ -2129,13 +2130,14 @@ mod tests { .unwrap(); store.add(&[assistant_record("a-0")]).await.unwrap(); - // One generation pending, below min (2): cleanup is a no-op. - assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); - assert_eq!(flushed_generation_count(&store).await, 1); + // One generation pending: the time trigger merges it immediately — + // it does not wait for a count threshold. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 1); + assert_eq!(flushed_generation_count(&store).await, 0); store.add(&[assistant_record("a-1")]).await.unwrap(); - // Two generations pending, at min: cleanup merges both. - assert_eq!(store.cleanup_own_shard().await.unwrap(), 2); + // Next generation is likewise merged on the following pass. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 1); assert_eq!(flushed_generation_count(&store).await, 0); // Rows survive the merge, readable exactly once. @@ -2169,7 +2171,6 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // only the timer merges cleanup_interval_secs: Some(1), - cleanup_min_generations: Some(1), }, ) .await @@ -2238,7 +2239,6 @@ mod tests { RolloutStoreOptions { storage_options: None, shard_id: Some(shard), - cleanup_min_generations: Some(1), ..Default::default() }, ) @@ -2309,7 +2309,6 @@ mod tests { RolloutStoreOptions { storage_options: None, shard_id: Some("rollout-0".to_string()), - cleanup_min_generations: Some(1), ..Default::default() }, ) diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 8337fb8..ed262da 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -37,20 +37,14 @@ pub struct ServerConfig { /// Interval, in seconds, for the periodic per-shard WAL cleanup task. When /// non-zero, each rollout store spawns a background timer that folds this /// instance's flushed MemWAL generations into the base table on a schedule — - /// the *time* half of the "time + count" trigger, complementing - /// `--rollout-merge-after-generations`. It reclaims stale generations even on - /// low-traffic shards that never cross the count threshold. `0` (the default) - /// disables the timer. + /// the *time* half of the "time OR count" trigger, complementing + /// `--rollout-merge-after-generations`. Whichever fires first merges: the + /// timer reclaims whatever is pending regardless of count, so stale + /// generations are folded in even on low-traffic shards that never cross the + /// count threshold. `0` (the default) disables the timer. #[arg(long, env = "ROLLOUT_CLEANUP_INTERVAL_SECS", default_value = "0")] pub rollout_cleanup_interval_secs: u64, - /// Minimum flushed generations a periodic cleanup tick requires before it - /// merges (avoids rewriting the base table to reclaim a single small - /// generation). Only meaningful when `--rollout-cleanup-interval-secs` is - /// 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 diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index e0288da..deb1812 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -54,7 +54,6 @@ pub async fn create_rollout_store( .then_some(state.rollout_merge_after_generations), cleanup_interval_secs: (state.rollout_cleanup_interval_secs > 0) .then_some(state.rollout_cleanup_interval_secs), - cleanup_min_generations: Some(state.rollout_cleanup_min_generations), }; let store = RolloutStore::open_with_options(&uri, options) diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index ba0508d..6c4eb27 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -47,9 +47,6 @@ pub struct AppState { /// Periodic per-shard WAL-cleanup interval in seconds; `0` disables the /// global sweeper. See [`Self::spawn_global_sweeper`]. pub rollout_cleanup_interval_secs: u64, - /// Minimum flushed generations before a periodic cleanup tick merges. See - /// `RolloutStoreOptions::cleanup_min_generations`. - pub rollout_cleanup_min_generations: usize, } impl AppState { @@ -72,7 +69,6 @@ impl AppState { instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, - rollout_cleanup_min_generations: config.rollout_cleanup_min_generations, }) } @@ -109,7 +105,6 @@ impl AppState { instance_id, rollout_merge_after_generations: 0, rollout_cleanup_interval_secs: 0, - rollout_cleanup_min_generations: 1, } } @@ -130,7 +125,6 @@ impl AppState { .then_some(self.rollout_merge_after_generations), cleanup_interval_secs: (self.rollout_cleanup_interval_secs > 0) .then_some(self.rollout_cleanup_interval_secs), - cleanup_min_generations: Some(self.rollout_cleanup_min_generations), } } From 1e6f4c0756c2cd58252139cb6ca215276d17c0e0 Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 15 Jul 2026 05:41:50 +0000 Subject: [PATCH 2/3] feat(master): unified task scheduler for compaction and WAL merge Generalize the master's single-serial compaction driver into a unified task scheduler that runs Compact and MergeWal tasks with bounded concurrency. Same-dataset compaction stays serialized via a per-name in-flight gate; distinct experiments run concurrently under a global semaphore. MergeWal fans out to configured worker endpoints so each worker merges its own MemWAL shard (the master cannot fence a live writer). Adds a task queue API + UI view and a "Merge WAL" action. Co-Authored-By: Claude Opus 4 --- Cargo.lock | 1 + crates/lance-context-api/src/lib.rs | 122 ++++++ crates/lance-context-master/Cargo.toml | 1 + crates/lance-context-master/src/config.rs | 16 + crates/lance-context-master/src/routes.rs | 114 ++++- crates/lance-context-master/src/scheduler.rs | 411 +++++++++++++----- crates/lance-context-master/src/state.rs | 37 +- crates/lance-context-master/ui/src/App.tsx | 135 +++++- crates/lance-context-master/ui/src/api.ts | 43 ++ crates/lance-context-master/ui/src/store.ts | 4 + crates/lance-context-master/ui/src/styles.css | 24 + crates/lance-context-server/src/routes/mod.rs | 4 + .../src/routes/rollouts.rs | 68 +++ 13 files changed, 846 insertions(+), 134 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0e03f9f..69a324a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5482,6 +5482,7 @@ dependencies = [ "lance-context-core", "lance-context-metrics", "metrics", + "reqwest 0.12.28", "serde", "serde_json", "tempfile", diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index d5ad06a..9442579 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -902,6 +902,78 @@ pub enum CompactJobStatus { None, } +// --------------------------------------------------------------------------- +// Unified task scheduler (master control-plane) +// --------------------------------------------------------------------------- + +/// The kind of work a scheduled task performs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TaskKind { + /// Compact an experiment's base table (rewrites fragments; runs on the + /// master, serialized per experiment because two `Rewrite`s conflict). + Compact, + /// Fold flushed MemWAL generations back into the base table. The master + /// cannot do this directly without fencing the live shard writer, so this + /// task fans out to the configured worker endpoints and each worker merges + /// its own shard. + MergeWal, +} + +/// Lifecycle state of a scheduled task, generalized from [`CompactJobStatus`] +/// so it applies uniformly to every [`TaskKind`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TaskState { + /// Accepted and waiting for a scheduler slot. + Queued, + /// Currently executing. + Running, + /// Finished successfully. + Done, + /// Finished with an error (see [`TaskRecord::error`]). + Failed, +} + +/// One unit of scheduled work plus its lifecycle, as surfaced to the queue UI. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TaskRecord { + /// Time-ordered unique id (UUIDv7). + pub id: String, + pub kind: TaskKind, + /// Experiment / rollout-store name this task acts on. + pub target: String, + pub state: TaskState, + /// Error message when `state == Failed`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Human-readable outcome summary when `state == Done` + /// (e.g. `"removed 3 / added 1 fragments"` or `"merged 4 gens on 2/3 workers"`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub detail: Option, + /// When the task was enqueued, Unix milliseconds. + pub enqueued_at: i64, + /// When execution began, Unix milliseconds, if started. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub started_at: Option, + /// When the task reached a terminal state, Unix milliseconds, if finished. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub finished_at: Option, +} + +/// Request body for `POST /api/v1/tasks`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EnqueueTaskRequest { + pub kind: TaskKind, + pub target: String, +} + +/// Response for `GET /api/v1/tasks`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskListResponse { + pub tasks: Vec, +} + #[cfg(test)] mod tests { use super::*; @@ -990,4 +1062,54 @@ mod tests { .unwrap(); assert_eq!(legacy.payload_uri, None); } + + #[test] + fn task_record_roundtrips_and_omits_empty_optionals() { + let queued = TaskRecord { + id: "0190-abc".to_string(), + kind: TaskKind::MergeWal, + target: "exp-1".to_string(), + state: TaskState::Queued, + error: None, + detail: None, + enqueued_at: 1_700_000_000_000, + started_at: None, + finished_at: None, + }; + let json = serde_json::to_string(&queued).unwrap(); + // snake_case tags for the enums. + assert!(json.contains(r#""kind":"merge_wal""#)); + assert!(json.contains(r#""state":"queued""#)); + // Absent optionals are not serialized. + assert!(!json.contains("error")); + assert!(!json.contains("started_at")); + assert!(!json.contains("finished_at")); + let back: TaskRecord = serde_json::from_str(&json).unwrap(); + assert_eq!(back, queued); + } + + #[test] + fn task_record_done_and_failed_carry_details() { + let done = TaskRecord { + id: "id".to_string(), + kind: TaskKind::Compact, + target: "exp".to_string(), + state: TaskState::Done, + error: None, + detail: Some("removed 3 / added 1 fragments".to_string()), + enqueued_at: 1, + started_at: Some(2), + finished_at: Some(3), + }; + let back: TaskRecord = serde_json::from_str(&serde_json::to_string(&done).unwrap()).unwrap(); + assert_eq!(back, done); + } + + #[test] + fn enqueue_task_request_parses_snake_case_kind() { + let req: EnqueueTaskRequest = + serde_json::from_str(r#"{"kind":"compact","target":"exp-7"}"#).unwrap(); + assert_eq!(req.kind, TaskKind::Compact); + assert_eq!(req.target, "exp-7"); + } } diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index 16c36c5..aaa4d26 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -24,6 +24,7 @@ clap = { version = "4", features = ["derive", "env"] } futures = "0.3" lance = "7.0.0" metrics = "0.24" +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "sync"] } diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index 546966d..b81a9a0 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -40,6 +40,22 @@ pub struct MasterConfig { #[arg(long, env = "TARGET_ROWS_PER_FRAGMENT", default_value_t = 1_048_576)] pub target_rows_per_fragment: usize, + /// Data-plane worker base URLs (comma-separated), e.g. + /// `http://rollout-0:3000,http://rollout-1:3000`. A `MergeWal` task fans out + /// to every endpoint so each worker merges its own MemWAL shard (the master + /// cannot merge a shard it does not own without fencing the live writer). + /// Empty (the default) means WAL-merge tasks have nowhere to go and fail + /// fast with a clear message. + #[arg(long, env = "WORKER_ENDPOINTS", value_delimiter = ',')] + pub worker_endpoints: Vec, + + /// Maximum number of scheduler tasks executing concurrently. Compaction of + /// the *same* experiment is always serialized regardless of this value + /// (two `Rewrite`s on one dataset conflict); this only bounds how many + /// *distinct* experiments/tasks run at once. + #[arg(long, env = "TASK_CONCURRENCY", default_value_t = 4)] + pub task_concurrency: usize, + /// Directory of built UI assets to serve. When unset, only the JSON API is /// exposed. #[arg(long, env = "UI_DIR")] diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 0371b4a..e4133e9 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -9,7 +9,8 @@ use axum::{Json, Router}; use serde::Deserialize; use lance_context_api::{ - CompactJobStatus, ExperimentDetail, ExperimentListResponse, ExperimentSummary, + CompactJobStatus, EnqueueTaskRequest, ExperimentDetail, ExperimentListResponse, + ExperimentSummary, TaskKind, TaskListResponse, TaskRecord, TaskState, }; use crate::error::MasterError; @@ -105,8 +106,58 @@ pub async fn rescan( Ok(Json(serde_json::json!({ "scanned": n }))) } +/// Map a task's terminal `detail` string back into the fragment counts the +/// legacy `CompactJobStatus::Done` shape carries. Best-effort: unparsable +/// details yield zeros (the UI only needs the state to stop polling). +fn parse_fragment_detail(detail: Option<&str>) -> (usize, usize) { + // Detail format: "removed {r} / added {a} fragments". + let parse = |s: &str, key: &str| -> usize { + s.split_whitespace() + .skip_while(|w| *w != key) + .nth(1) + .and_then(|w| w.parse().ok()) + .unwrap_or(0) + }; + match detail { + Some(d) => (parse(d, "removed"), parse(d, "added")), + None => (0, 0), + } +} + +/// Project a [`TaskRecord`] into the legacy [`CompactJobStatus`] shape so the +/// existing compaction UI keeps working while the queue view rolls out. +fn task_to_compact_status(task: &TaskRecord) -> CompactJobStatus { + match task.state { + TaskState::Queued => CompactJobStatus::Queued, + TaskState::Running => CompactJobStatus::Running, + TaskState::Done => { + let (fragments_removed, fragments_added) = parse_fragment_detail(task.detail.as_deref()); + CompactJobStatus::Done { + fragments_removed, + fragments_added, + } + } + TaskState::Failed => CompactJobStatus::Failed { + error: task.error.clone().unwrap_or_default(), + }, + } +} + +/// The most recent `Compact` task for `name`, if any (by `enqueued_at`). +async fn latest_compact_task(state: &Arc, name: &str) -> Option { + state + .tasks + .lock() + .await + .values() + .filter(|t| t.kind == TaskKind::Compact && t.target == name) + .max_by_key(|t| t.enqueued_at) + .cloned() +} + /// `POST /api/v1/experiments/{name}/compact` — enqueue a manual compaction. -/// Returns 202 Accepted with the (possibly de-duped) job status. +/// Returns 202 Accepted with the (possibly de-duped) job status. Retained for +/// backward compatibility; internally this is just a `Compact` task. pub async fn compact_experiment( State(state): State>, Path(name): Path, @@ -125,8 +176,8 @@ pub async fn compact_experiment( name ))); } - let status = scheduler::enqueue(&state, &name).await; - Ok((StatusCode::ACCEPTED, Json(status))) + let task = scheduler::enqueue(&state, TaskKind::Compact, &name).await; + Ok((StatusCode::ACCEPTED, Json(task_to_compact_status(&task)))) } /// `GET /api/v1/experiments/{name}/compact/status` — latest compaction job @@ -135,14 +186,51 @@ pub async fn compact_status( State(state): State>, Path(name): Path, ) -> Json { - let status = state - .jobs - .lock() + match latest_compact_task(&state, &name).await { + Some(task) => Json(task_to_compact_status(&task)), + None => Json(CompactJobStatus::None), + } +} + +/// `POST /api/v1/tasks` — enqueue a task of any kind. Returns 202 Accepted with +/// the (possibly de-duped) task record. +pub async fn enqueue_task( + State(state): State>, + Json(req): Json, +) -> Result<(StatusCode, Json), MasterError> { + let exists = state + .registry + .write() .await - .get(&name) - .cloned() - .unwrap_or(CompactJobStatus::None); - Json(status) + .contains(&req.target) + .await + .map_err(MasterError::from_lance)?; + if !exists { + return Err(MasterError::NotFound(format!( + "experiment '{}' does not exist", + req.target + ))); + } + let task = scheduler::enqueue(&state, req.kind, &req.target).await; + Ok((StatusCode::ACCEPTED, Json(task))) +} + +/// `GET /api/v1/tasks` — all tasks (queue + recent history), newest first. +pub async fn list_tasks(State(state): State>) -> Json { + let mut tasks: Vec = state.tasks.lock().await.values().cloned().collect(); + tasks.sort_by_key(|t| std::cmp::Reverse(t.enqueued_at)); + Json(TaskListResponse { tasks }) +} + +/// `GET /api/v1/tasks/{id}` — a single task by id. +pub async fn get_task( + State(state): State>, + Path(id): Path, +) -> Result, MasterError> { + match state.tasks.lock().await.get(&id).cloned() { + Some(task) => Ok(Json(task)), + None => Err(MasterError::NotFound(format!("task '{}' not found", id))), + } } /// Build the admin API router (mounted under `/api/v1`). @@ -152,6 +240,8 @@ pub fn api_router() -> Router> { .route("/experiments/{name}", get(get_experiment)) .route("/experiments/{name}/compact", post(compact_experiment)) .route("/experiments/{name}/compact/status", get(compact_status)) + .route("/tasks", post(enqueue_task).get(list_tasks)) + .route("/tasks/{id}", get(get_task)) .route("/rescan", post(rescan)) } @@ -172,6 +262,8 @@ mod tests { compaction_interval_secs: 0, min_fragments: 16, target_rows_per_fragment: 1_048_576, + worker_endpoints: vec![], + task_concurrency: 4, ui_dir: None, } } diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 6a75ab6..dd0a82f 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -1,19 +1,29 @@ -//! Centralized compaction scheduler. +//! Unified task scheduler. //! -//! Compaction rewrites fragments (`Rewrite` in Lance's conflict matrix). Two -//! concurrent `Rewrite`s on the same dataset conflict, so compaction must have a -//! **single serial driver**. This module is that driver: one background task -//! drains a single queue, so automatic sweeps and manual API triggers share the -//! same serial execution path and can never overlap. `Rewrite` vs the -//! data-plane's `Append` is non-conflicting, so this runs safely alongside live -//! ingest. +//! The master runs a single scheduler that drains one queue with **bounded +//! concurrency**. It executes two kinds of task ([`TaskKind`]): +//! +//! - **Compact** — rewrites an experiment's base-table fragments. Two `Rewrite`s +//! on the *same* dataset conflict in Lance's conflict matrix, so compaction of +//! one experiment is serialized against itself via a per-name in-flight gate +//! ([`MasterState::inflight_compacts`]). Distinct experiments compact +//! concurrently. `Rewrite` vs the data-plane's `Append` is non-conflicting, so +//! this runs safely alongside live ingest. +//! - **MergeWal** — folds flushed MemWAL generations back into the base table. +//! The master cannot do this itself without fencing the live shard writer, so +//! it fans out to every configured worker endpoint and each worker merges its +//! own shard (`POST /api/v1/internal/merge-wal/{name}`). +//! +//! Each task runs in its own `tokio::spawn`, so one task's failure never affects +//! another. A global [`Semaphore`] bounds how many run at once. use std::sync::Arc; use std::time::Duration; use chrono::Utc; -use lance_context_api::CompactJobStatus; -use lance_context_core::{CompactionConfig, RolloutStore, RolloutStoreOptions}; +use lance_context_api::{TaskKind, TaskRecord, TaskState}; +use lance_context_core::{generate_id, CompactionConfig, RolloutStore, RolloutStoreOptions}; +use tokio::sync::Semaphore; use tokio::task::JoinHandle; use crate::state::MasterState; @@ -29,78 +39,204 @@ pub fn compaction_config(state: &MasterState) -> CompactionConfig { } } -/// Enqueue an experiment for compaction and mark it `Queued`. De-dupes: if a -/// job for `name` is already queued or running, this is a no-op returning the -/// existing status. -pub async fn enqueue(state: &Arc, name: &str) -> CompactJobStatus { - { - let jobs = state.jobs.lock().await; - if matches!( - jobs.get(name), - Some(CompactJobStatus::Queued) | Some(CompactJobStatus::Running) - ) { - return jobs.get(name).cloned().unwrap(); +/// Current Unix-millisecond timestamp. +fn now_ms() -> i64 { + Utc::now().timestamp_millis() +} + +fn kind_label(kind: TaskKind) -> &'static str { + match kind { + TaskKind::Compact => "compact", + TaskKind::MergeWal => "merge_wal", + } +} + +/// Enqueue a task and return its record. For [`TaskKind::Compact`] this de-dupes +/// against an existing non-terminal task for the same target: if one is already +/// `Queued` or `Running`, its record is returned unchanged and nothing new is +/// enqueued. `MergeWal` tasks are not de-duped (each fan-out is independent). +pub async fn enqueue(state: &Arc, kind: TaskKind, target: &str) -> TaskRecord { + if matches!(kind, TaskKind::Compact) { + let tasks = state.tasks.lock().await; + if let Some(existing) = tasks.values().find(|t| { + t.kind == TaskKind::Compact + && t.target == target + && matches!(t.state, TaskState::Queued | TaskState::Running) + }) { + return existing.clone(); } } + + let record = TaskRecord { + id: generate_id(), + kind, + target: target.to_string(), + state: TaskState::Queued, + error: None, + detail: None, + enqueued_at: now_ms(), + started_at: None, + finished_at: None, + }; state - .jobs + .tasks .lock() .await - .insert(name.to_string(), CompactJobStatus::Queued); + .insert(record.id.clone(), record.clone()); // Unbounded send only fails if the receiver was dropped (scheduler gone). - let _ = state.compact_tx.send(name.to_string()); - metrics::counter!("master_compaction_enqueued_total").increment(1); - metrics::gauge!("master_compaction_queue_depth").increment(1.0); - CompactJobStatus::Queued + let _ = state.task_tx.send(record.id.clone()); + metrics::counter!("master_task_enqueued_total", "kind" => kind_label(kind)).increment(1); + metrics::gauge!("master_task_queue_depth").increment(1.0); + record } -/// Compact one experiment now (serial, on the scheduler task). Updates the job -/// status map and the stats table's compaction counters on success. -async fn run_compaction(state: &Arc, name: &str) { - // This job is leaving the queue and entering execution. - metrics::gauge!("master_compaction_queue_depth").decrement(1.0); - state - .jobs - .lock() - .await - .insert(name.to_string(), CompactJobStatus::Running); +/// Mutate a task record in place under the lock, if it still exists. +async fn update_task(state: &Arc, id: &str, f: impl FnOnce(&mut TaskRecord)) { + if let Some(rec) = state.tasks.lock().await.get_mut(id) { + f(rec); + } +} + +/// Execute one task by id: dispatch on its kind, recording lifecycle timestamps +/// and the terminal state. +async fn run_task(state: &Arc, id: String) { + metrics::gauge!("master_task_queue_depth").decrement(1.0); + + let Some(task) = state.tasks.lock().await.get(&id).cloned() else { + return; // task vanished (should not happen) + }; + update_task(state, &id, |t| { + t.state = TaskState::Running; + t.started_at = Some(now_ms()); + }) + .await; + + let started = std::time::Instant::now(); + let outcome = match task.kind { + TaskKind::Compact => run_compaction(state, &task.target).await, + TaskKind::MergeWal => run_merge_wal(state, &task.target).await, + }; + + metrics::histogram!("master_task_duration_seconds", "kind" => kind_label(task.kind)) + .record(started.elapsed().as_secs_f64()); + let result = if outcome.is_ok() { "success" } else { "failed" }; + metrics::counter!("master_tasks_total", "kind" => kind_label(task.kind), "result" => result) + .increment(1); + + update_task(state, &id, |t| { + t.finished_at = Some(now_ms()); + match outcome { + Ok(detail) => { + t.state = TaskState::Done; + t.detail = Some(detail); + } + Err(error) => { + tracing::warn!(task = %id, target = %t.target, error = %error, "task failed"); + t.state = TaskState::Failed; + t.error = Some(error); + } + } + }) + .await; +} + +/// Compact one experiment. Serialized per experiment via the in-flight gate so +/// two `Rewrite`s on the same dataset never race. Returns a human-readable +/// outcome summary on success. +async fn run_compaction(state: &Arc, name: &str) -> Result { + // Per-name serial gate: refuse if a compaction for this experiment is + // already running. (De-dup at enqueue makes this rare, but the gate is the + // real guarantee against concurrent Rewrites.) + { + let mut inflight = state.inflight_compacts.lock().await; + if !inflight.insert(name.to_string()) { + return Err(format!("a compaction for '{name}' is already in progress")); + } + } + // Ensure the gate is released no matter how the inner call exits. + let result = compact_inner(state, name).await; + state.inflight_compacts.lock().await.remove(name); + result +} +async fn compact_inner(state: &Arc, name: &str) -> Result { let uri = state.rollout_uri(name); let config = compaction_config(state); let opts = RolloutStoreOptions::default(); - let compact_start = std::time::Instant::now(); - let status = match RolloutStore::open_existing_with_options(&uri, opts).await { - Ok(mut store) => match store.compact(Some(config)).await { - Ok(metrics) => { - update_stats_after_compaction(state, name, &store).await; - CompactJobStatus::Done { - fragments_removed: metrics.fragments_removed, - fragments_added: metrics.fragments_added, - } + let mut store = RolloutStore::open_existing_with_options(&uri, opts) + .await + .map_err(|e| e.to_string())?; + let metrics = store + .compact(Some(config)) + .await + .map_err(|e| e.to_string())?; + update_stats_after_compaction(state, name, &store).await; + Ok(format!( + "removed {} / added {} fragments", + metrics.fragments_removed, metrics.fragments_added + )) +} + +/// Shape of the worker's merge-wal response (`{ "reclaimed": n }`). +#[derive(serde::Deserialize)] +struct MergeWalReply { + reclaimed: usize, +} + +/// Fan a WAL-merge out to every configured worker endpoint. Each worker merges +/// its own shard; a worker that owns no data for `name` reports 0 (or 404, which +/// we tolerate). Succeeds if at least one endpoint responded; fails only when +/// there are no endpoints or every one errored. +async fn run_merge_wal(state: &Arc, name: &str) -> Result { + let endpoints = &state.config.worker_endpoints; + if endpoints.is_empty() { + return Err("no worker endpoints configured (--worker-endpoints)".to_string()); + } + + let calls = endpoints.iter().map(|ep| { + let http = state.http.clone(); + let url = format!( + "{}/api/v1/internal/merge-wal/{}", + ep.trim_end_matches('/'), + name + ); + async move { + let resp = http.post(&url).send().await.map_err(|e| e.to_string())?; + let status = resp.status(); + if status == reqwest::StatusCode::NOT_FOUND { + // This worker owns no shard for `name`; treat as 0 reclaimed. + return Ok::(0); } - Err(e) => CompactJobStatus::Failed { - error: e.to_string(), - }, - }, - Err(e) => CompactJobStatus::Failed { - error: e.to_string(), - }, - }; + if !status.is_success() { + return Err(format!("{url}: HTTP {status}")); + } + let body: MergeWalReply = resp.json().await.map_err(|e| e.to_string())?; + Ok(body.reclaimed) + } + }); - metrics::histogram!("master_compaction_duration_seconds") - .record(compact_start.elapsed().as_secs_f64()); - let result = if matches!(status, CompactJobStatus::Done { .. }) { - "success" - } else { - "failed" - }; - metrics::counter!("master_compactions_total", "result" => result).increment(1); + let results = futures::future::join_all(calls).await; + let total_workers = results.len(); + let mut reclaimed = 0usize; + let mut ok_workers = 0usize; + let mut last_err = None; + for r in results { + match r { + Ok(n) => { + reclaimed += n; + ok_workers += 1; + } + Err(e) => last_err = Some(e), + } + } - if let CompactJobStatus::Failed { error } = &status { - tracing::warn!(store = %name, error = %error, "compaction failed"); + if ok_workers == 0 { + return Err(last_err.unwrap_or_else(|| "all workers failed".to_string())); } - state.jobs.lock().await.insert(name.to_string(), status); + Ok(format!( + "merged {reclaimed} generations across {ok_workers}/{total_workers} workers" + )) } /// Refresh the stats row for `name` after a successful compaction: re-observe @@ -134,8 +270,9 @@ async fn update_stats_after_compaction(state: &Arc, name: &str, sto } } -/// Enqueue every experiment whose fragment count is at or above the configured -/// threshold, honoring quiet hours. Reads candidates from the stats table. +/// Enqueue a `Compact` task for every experiment whose fragment count is at or +/// above the configured threshold, honoring quiet hours. Reads candidates from +/// the stats table. pub async fn sweep_candidates(state: &Arc) -> lance::Result { let config = compaction_config(state); // Quiet-hours gate applies to the whole sweep. @@ -146,7 +283,7 @@ pub async fn sweep_candidates(state: &Arc) -> lance::Result let mut queued = 0; for row in rows { if row.fragment_count as usize >= config.min_fragments { - enqueue(state, &row.name).await; + enqueue(state, TaskKind::Compact, &row.name).await; queued += 1; } } @@ -165,13 +302,13 @@ fn in_quiet_hours(config: &CompactionConfig) -> bool { .any(|(start, end)| hour >= *start && hour < *end) } -/// Spawn the single serial compaction worker plus (optionally) the periodic -/// auto-sweep. The worker owns the queue receiver and processes one experiment -/// at a time. Returns the worker handle. Panics if called twice (receiver -/// already taken). +/// Spawn the scheduler dispatcher plus (optionally) the periodic auto-sweep. The +/// dispatcher owns the queue receiver and runs up to `task_concurrency` tasks at +/// once, each in its own detached task. Returns the dispatcher handle. Panics if +/// called twice (receiver already taken). pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { let mut rx = state - .compact_rx + .task_rx .try_lock() .expect("scheduler: state lock") .take() @@ -195,10 +332,22 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { }); } - let worker_state = state.clone(); + let concurrency = state.config.task_concurrency.max(1); + let sem = Arc::new(Semaphore::new(concurrency)); + let dispatch_state = state.clone(); tokio::spawn(async move { - while let Some(name) = rx.recv().await { - run_compaction(&worker_state, &name).await; + while let Some(id) = rx.recv().await { + // Acquire a slot before spawning so at most `concurrency` run at once. + let permit = sem + .clone() + .acquire_owned() + .await + .expect("semaphore never closed"); + let st = dispatch_state.clone(); + tokio::spawn(async move { + run_task(&st, id).await; + drop(permit); + }); } }) } @@ -220,12 +369,27 @@ mod tests { // Low threshold so a handful of appends crosses it. min_fragments: 2, target_rows_per_fragment: 1_048_576, + worker_endpoints: vec![], + task_concurrency: 4, ui_dir: None, } } - /// Manual enqueue -> serial worker compacts -> job reaches Done and the - /// stats table records a compaction. + /// Wait until the task reaches a terminal state, returning its final record. + async fn await_terminal(state: &Arc, id: &str) -> TaskRecord { + for _ in 0..100 { + tokio::time::sleep(Duration::from_millis(50)).await; + if let Some(t) = state.tasks.lock().await.get(id) { + if matches!(t.state, TaskState::Done | TaskState::Failed) { + return t.clone(); + } + } + } + panic!("task {id} did not reach a terminal state"); + } + + /// Manual enqueue -> dispatcher compacts -> task reaches Done and the stats + /// table records a compaction. #[tokio::test] async fn manual_compaction_runs_and_updates_stats() { let dir = TempDir::new().unwrap(); @@ -238,7 +402,7 @@ mod tests { { let mut store = RolloutStore::open(&uri).await.unwrap(); for i in 0..4 { - let rec = crate::scheduler::tests::rollout_record(&format!("r{i}")); + let rec = rollout_record(&format!("r{i}")); store.add(&[rec]).await.unwrap(); store.cleanup_own_shard().await.unwrap(); } @@ -253,27 +417,9 @@ mod tests { // Seed a stats row so post-compaction upsert has a prior counter. crate::scanner::scan_once(&state).await.unwrap(); - enqueue(&state, name).await; - - // Wait for the worker to reach a terminal state. - let mut done = None; - for _ in 0..50 { - tokio::time::sleep(Duration::from_millis(50)).await; - if let Some(s) = state.jobs.lock().await.get(name) { - if matches!( - s, - CompactJobStatus::Done { .. } | CompactJobStatus::Failed { .. } - ) { - done = Some(s.clone()); - break; - } - } - } - let status = done.expect("job reached terminal state"); - assert!( - matches!(status, CompactJobStatus::Done { .. }), - "expected Done, got {status:?}" - ); + let rec = enqueue(&state, TaskKind::Compact, name).await; + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Done, "got {status:?}"); let row = state.stats.lock().await.get(name).await.unwrap().unwrap(); assert_eq!(row.total_compactions, 1); @@ -282,17 +428,18 @@ mod tests { worker.abort(); } + /// Enqueuing the same experiment twice while queued de-dupes to one task. #[tokio::test] - async fn enqueue_dedupes_queued_jobs() { + async fn enqueue_dedupes_queued_compactions() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); - // Do NOT spawn the worker, so jobs stay Queued. - let s1 = enqueue(&state, "x").await; - let s2 = enqueue(&state, "x").await; - assert!(matches!(s1, CompactJobStatus::Queued)); - assert!(matches!(s2, CompactJobStatus::Queued)); - // Only one message should be in the channel; drain and count. - let mut rx = state.compact_rx.lock().await.take().unwrap(); + // Do NOT spawn the dispatcher, so the first task stays Queued. + let a = enqueue(&state, TaskKind::Compact, "x").await; + let b = enqueue(&state, TaskKind::Compact, "x").await; + assert_eq!(a.id, b.id, "second enqueue returns the same task"); + // Only one task in the map, one message on the queue. + assert_eq!(state.tasks.lock().await.len(), 1); + let mut rx = state.task_rx.lock().await.take().unwrap(); let mut n = 0; while rx.try_recv().is_ok() { n += 1; @@ -300,6 +447,54 @@ mod tests { assert_eq!(n, 1); } + /// A MergeWal task with no configured endpoints fails fast with a clear msg. + #[tokio::test] + async fn merge_wal_without_endpoints_fails() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(config(&dir)).await.unwrap(); + let worker = spawn_scheduler(&state); + let rec = enqueue(&state, TaskKind::MergeWal, "exp").await; + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Failed); + assert!(status.error.unwrap().contains("worker endpoints")); + worker.abort(); + } + + /// MergeWal fans out to every configured worker endpoint and sums the + /// reclaimed counts. Uses a tiny in-process stub server per "worker". + #[tokio::test] + async fn merge_wal_broadcasts_and_sums_reclaimed() { + use axum::{routing::post, Json, Router}; + + // A stub worker that always reports `reclaimed` for any merge call. + async fn spawn_stub(reclaimed: usize) -> String { + let app = Router::new().route( + "/api/v1/internal/merge-wal/{name}", + post(move || async move { Json(serde_json::json!({ "reclaimed": reclaimed })) }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + format!("http://{addr}") + } + + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.worker_endpoints = vec![spawn_stub(3).await, spawn_stub(2).await]; + let state = MasterState::new(cfg).await.unwrap(); + let worker = spawn_scheduler(&state); + + let rec = enqueue(&state, TaskKind::MergeWal, "exp").await; + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Done, "got {status:?}"); + let detail = status.detail.unwrap(); + assert!(detail.contains("merged 5 generations"), "detail: {detail}"); + assert!(detail.contains("2/2 workers"), "detail: {detail}"); + worker.abort(); + } + /// Minimal rollout record builder for tests (the core struct has no /// `Default`). pub fn rollout_record(id: &str) -> lance_context_core::RolloutRecord { diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index c3bc0f9..435fb1b 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -1,9 +1,9 @@ //! Shared master state: durable registry + stats table. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use lance_context_api::CompactJobStatus; +use lance_context_api::TaskRecord; use lance_context_core::{join_uri, RolloutRegistry}; use tokio::sync::{mpsc, Mutex, RwLock}; @@ -27,14 +27,21 @@ pub struct MasterState { pub base_uri: String, /// Effective configuration. pub config: MasterConfig, - /// Enqueues an experiment name for (serial) compaction. Both automatic - /// sweeps and manual API triggers push here, so compaction has a single - /// serial driver — never two concurrent `Rewrite`s. - pub compact_tx: mpsc::UnboundedSender, + /// Enqueues a task id for the scheduler. Both automatic sweeps and manual + /// API triggers push here; the scheduler drains it with bounded concurrency + /// (see [`crate::scheduler`]). + pub task_tx: mpsc::UnboundedSender, /// Receiver half, taken once by [`crate::scheduler::spawn_scheduler`]. - pub compact_rx: Mutex>>, - /// Last-known compaction job state per experiment, for the status endpoint. - pub jobs: Mutex>, + pub task_rx: Mutex>>, + /// Every task (queued / running / terminal) keyed by task id, for the queue + /// API and UI. Retained after completion so the UI can show recent history. + pub tasks: Mutex>, + /// Experiment names with a compaction currently executing — the per-name + /// serial gate. A `Compact` task must not run while another `Compact` on the + /// same experiment is in flight (two `Rewrite`s on one dataset conflict). + pub inflight_compacts: Mutex>, + /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. + pub http: reqwest::Client, } impl MasterState { @@ -52,15 +59,17 @@ impl MasterState { ); } let stats = StatsStore::open_or_create(&stats_uri, None).await?; - let (compact_tx, compact_rx) = mpsc::unbounded_channel(); + let (task_tx, task_rx) = mpsc::unbounded_channel(); Ok(Arc::new(Self { registry: RwLock::new(registry), stats: Mutex::new(stats), base_uri, config, - compact_tx, - compact_rx: Mutex::new(Some(compact_rx)), - jobs: Mutex::new(HashMap::new()), + task_tx, + task_rx: Mutex::new(Some(task_rx)), + tasks: Mutex::new(HashMap::new()), + inflight_compacts: Mutex::new(HashSet::new()), + http: reqwest::Client::new(), })) } @@ -87,6 +96,8 @@ mod tests { compaction_interval_secs: 0, min_fragments: 16, target_rows_per_fragment: 1_048_576, + worker_endpoints: vec![], + task_concurrency: 4, ui_dir: None, } } diff --git a/crates/lance-context-master/ui/src/App.tsx b/crates/lance-context-master/ui/src/App.tsx index 3ae3682..fcefc1d 100644 --- a/crates/lance-context-master/ui/src/App.tsx +++ b/crates/lance-context-master/ui/src/App.tsx @@ -2,11 +2,15 @@ import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { useEffect, useState } from "react"; import { compactionStatus, + enqueueTask, getExperiment, listExperiments, + listTasks, triggerCompaction, type CompactJobStatus, type ExperimentSummary, + type TaskRecord, + type TaskState, } from "./api"; import { useUiStore } from "./store"; @@ -236,6 +240,7 @@ function Drawer({ name }: { name: string }) {
+
@@ -256,14 +261,120 @@ function StatCard({ label, value, unit }: { label: string; value: string; unit?: ); } +/* ---- task queue view ----------------------------------------------------- */ + +/** Generic state pill for a scheduler task (queued/running/done/failed). */ +function TaskStatePill({ state, title }: { state: TaskState; title?: string }) { + const cls = + state === "queued" + ? "pill--queued" + : state === "running" + ? "pill--running" + : state === "done" + ? "pill--done" + : "pill--failed"; + return ( + + + {state} + + ); +} + +function taskDuration(t: TaskRecord): string { + const end = t.finished_at ?? Date.now(); + const start = t.started_at ?? t.enqueued_at; + if (!start) return "—"; + const s = Math.max(0, Math.round((end - start) / 1000)); + if (s < 60) return `${s}s`; + const m = Math.floor(s / 60); + return `${m}m ${s % 60}s`; +} + +/** "Merge WAL" trigger — enqueues a merge_wal task for one experiment. */ +function MergeWalButton({ name }: { name: string }) { + const qc = useQueryClient(); + const setView = useUiStore((s) => s.setView); + const trigger = useMutation({ + mutationFn: () => enqueueTask({ kind: "merge_wal", target: name }), + onSuccess: () => { + qc.invalidateQueries({ queryKey: ["tasks"] }); + setView("tasks"); + }, + }); + return ( + + ); +} + +function TaskQueue() { + const tasks = useQuery({ + queryKey: ["tasks"], + queryFn: () => listTasks(), + refetchInterval: 1000, + }); + const rows = tasks.data?.tasks ?? []; + const active = rows.filter((t) => t.state === "queued" || t.state === "running").length; + + return ( + <> +
+ + +
+
+ {tasks.isError &&
{String(tasks.error)}
} + {!tasks.isError && ( + + + + + + + + + + + + + {rows.map((t) => ( + + + + + + + + + ))} + +
KindTargetStateDetailDurationEnqueued
{t.kind}{t.target} + + {t.detail ?? t.error ?? "—"}{taskDuration(t)}{relTime(t.enqueued_at)}
+ )} + {!tasks.isLoading && !tasks.isError && rows.length === 0 && ( +
No tasks in the queue.
+ )} +
+ + ); +} + /* ---- app ----------------------------------------------------------------- */ export function App() { - const { search, page, pageSize, selected, setSearch, setPage } = useUiStore(); + const { view, search, page, pageSize, selected, setView, setSearch, setPage } = useUiStore(); const list = useQuery({ queryKey: ["experiments", search, page, pageSize], queryFn: () => listExperiments(search, pageSize, page * pageSize), placeholderData: (prev) => prev, + enabled: view === "experiments", }); const total = list.data?.total ?? 0; @@ -284,13 +395,31 @@ export function App() { control plane
+
{list.isFetching ? "syncing" : "live"}
-
+ {view === "tasks" ? ( + + ) : ( + <> +
@@ -364,6 +493,8 @@ export function App() {
)} + + )} {selected && }
diff --git a/crates/lance-context-master/ui/src/api.ts b/crates/lance-context-master/ui/src/api.ts index b9ff1a9..c733c8b 100644 --- a/crates/lance-context-master/ui/src/api.ts +++ b/crates/lance-context-master/ui/src/api.ts @@ -25,6 +25,31 @@ export type CompactJobStatus = | { state: "failed"; error: string } | { state: "none" }; +// Unified scheduler task DTOs (mirror `TaskRecord` etc. in `lance-context-api`). +export type TaskKind = "compact" | "merge_wal"; +export type TaskState = "queued" | "running" | "done" | "failed"; + +export interface TaskRecord { + id: string; + kind: TaskKind; + target: string; + state: TaskState; + error?: string | null; + detail?: string | null; + enqueued_at: number; + started_at?: number | null; + finished_at?: number | null; +} + +export interface TaskListResponse { + tasks: TaskRecord[]; +} + +export interface EnqueueTaskRequest { + kind: TaskKind; + target: string; +} + const API = "/api/v1"; async function json(res: Response): Promise { @@ -68,3 +93,21 @@ export async function compactionStatus(name: string): Promise await fetch(`${API}/experiments/${encodeURIComponent(name)}/compact/status`), ); } + +export async function listTasks(): Promise { + return json(await fetch(`${API}/tasks`)); +} + +export async function getTask(id: string): Promise { + return json(await fetch(`${API}/tasks/${encodeURIComponent(id)}`)); +} + +export async function enqueueTask(req: EnqueueTaskRequest): Promise { + return json( + await fetch(`${API}/tasks`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(req), + }), + ); +} diff --git a/crates/lance-context-master/ui/src/store.ts b/crates/lance-context-master/ui/src/store.ts index 478cf50..4737b69 100644 --- a/crates/lance-context-master/ui/src/store.ts +++ b/crates/lance-context-master/ui/src/store.ts @@ -2,20 +2,24 @@ import { create } from "zustand"; /// UI-local view state: the current search query and pagination cursor. interface UiState { + view: "experiments" | "tasks"; search: string; page: number; pageSize: number; selected: string | null; + setView: (view: "experiments" | "tasks") => void; setSearch: (search: string) => void; setPage: (page: number) => void; select: (name: string | null) => void; } export const useUiStore = create((set) => ({ + view: "experiments", search: "", page: 0, pageSize: 25, selected: null, + setView: (view) => set({ view }), // Changing the search resets pagination to the first page. setSearch: (search) => set({ search, page: 0 }), setPage: (page) => set({ page }), diff --git a/crates/lance-context-master/ui/src/styles.css b/crates/lance-context-master/ui/src/styles.css index 09ba623..ef9de92 100644 --- a/crates/lance-context-master/ui/src/styles.css +++ b/crates/lance-context-master/ui/src/styles.css @@ -155,6 +155,30 @@ body { flex: 1; } +.tabs { + display: flex; + gap: 4px; + margin-right: 16px; +} +.tab { + background: transparent; + border: 1px solid transparent; + color: var(--text-muted); + font-family: var(--mono); + font-size: 12px; + padding: 5px 12px; + border-radius: 6px; + cursor: pointer; +} +.tab:hover { + color: var(--text); +} +.tab--active { + color: var(--text); + border-color: var(--border); + background: var(--surface, rgba(255, 255, 255, 0.04)); +} + .live { display: flex; align-items: center; diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index a8d5b1e..c42eb15 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -113,6 +113,10 @@ pub fn router() -> Router> { "/api/v1/rollouts/{name}/compact/stats", get(rollouts::compact_rollout_stats), ) + .route( + "/api/v1/internal/merge-wal/{name}", + post(rollouts::merge_wal), + ) } #[cfg(test)] diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index deb1812..b2e5fdc 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -26,6 +26,13 @@ use crate::state::AppState; /// raises the ceiling well above it while still bounding memory. pub const MAX_ROLLOUT_UPLOAD_BYTES: usize = 1024 * 1024 * 1024; +/// Response for the internal WAL-merge trigger: how many flushed generations +/// this worker's shard folded into the base table. +#[derive(Debug, serde::Serialize, serde::Deserialize)] +pub struct MergeWalResponse { + pub reclaimed: usize, +} + pub async fn create_rollout_store( State(state): State>, Json(req): Json, @@ -452,6 +459,28 @@ pub async fn compact_rollout_stats( })) } +/// Trigger a WAL merge on this worker's own shard for `name`. +/// +/// This is the worker half of the master-driven "MergeWal" task: the master +/// cannot merge a shard it does not own without fencing the live writer, so it +/// fans this call out to every worker and each worker folds *its own* flushed +/// MemWAL generations into the base table via [`RolloutStore::cleanup_own_shard`] +/// (which merges whatever is pending, no count threshold). A worker whose shard +/// has nothing pending simply reports `reclaimed: 0`. +pub async fn merge_wal( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { + let store_lock = state.get_or_open_rollout_store(&name).await?; + let mut store = store_lock.write().await; + let reclaimed = store.cleanup_own_shard().await.map_err(AppError::from_lance)?; + if reclaimed > 0 { + ::metrics::counter!("rollout_wal_cleanup_total", "result" => "merged").increment(1); + ::metrics::counter!("rollout_wal_generations_reclaimed_total").increment(reclaimed as u64); + } + Ok(Json(MergeWalResponse { reclaimed })) +} + fn content_type_is(content_type: &str, expected: &str) -> bool { content_type .split(';') @@ -845,6 +874,45 @@ mod tests { assert!(matches!(err, AppError::InvalidRequest(_))); } + /// The internal merge-wal endpoint folds this worker's pending flushed + /// generations into the base table and reports how many it reclaimed. A + /// second call with nothing pending is a no-op reporting `0`. + #[tokio::test] + async fn merge_wal_reclaims_pending_then_is_noop() { + let (state, _dir) = rollout_state().await; + // One append flushes one generation into this instance's shard. + let body = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", None)], + }) + .unwrap(); + let _ = add_rollouts( + State(state.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .expect("append succeeds"); + + let Json(first) = merge_wal(State(state.clone()), Path("rl".to_string())) + .await + .expect("merge succeeds"); + assert_eq!(first.reclaimed, 1, "one pending generation merged"); + + let Json(second) = merge_wal(State(state.clone()), Path("rl".to_string())) + .await + .expect("second merge succeeds"); + assert_eq!(second.reclaimed, 0, "nothing left to merge"); + + // The row survives the merge, readable exactly once. + let Json(got) = get_rollout( + State(state.clone()), + Path(("rl".to_string(), "r0".to_string())), + ) + .await + .unwrap(); + assert!(got.record.is_some()); + } + /// A second server instance (distinct in-memory cache, shared data dir) /// that never `create`d the store must still serve reads/writes for it by /// lazily loading from storage — the multi-replica 404 regression. From 781ca789aad3b3b15c025d61719594c9ec2c5c71 Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 15 Jul 2026 06:07:12 +0000 Subject: [PATCH 3/3] style: apply rustfmt to scheduler + merge-wal code --- crates/lance-context-api/src/lib.rs | 3 ++- crates/lance-context-master/src/routes.rs | 3 ++- crates/lance-context-server/src/routes/rollouts.rs | 5 ++++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 1785916..f29cd1e 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -1111,7 +1111,8 @@ mod tests { started_at: Some(2), finished_at: Some(3), }; - let back: TaskRecord = serde_json::from_str(&serde_json::to_string(&done).unwrap()).unwrap(); + let back: TaskRecord = + serde_json::from_str(&serde_json::to_string(&done).unwrap()).unwrap(); assert_eq!(back, done); } diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 00a8f50..1b4468f 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -231,7 +231,8 @@ fn task_to_compact_status(task: &TaskRecord) -> CompactJobStatus { TaskState::Queued => CompactJobStatus::Queued, TaskState::Running => CompactJobStatus::Running, TaskState::Done => { - let (fragments_removed, fragments_added) = parse_fragment_detail(task.detail.as_deref()); + let (fragments_removed, fragments_added) = + parse_fragment_detail(task.detail.as_deref()); CompactJobStatus::Done { fragments_removed, fragments_added, diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index b2e5fdc..6884df9 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -473,7 +473,10 @@ pub async fn merge_wal( ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; let mut store = store_lock.write().await; - let reclaimed = store.cleanup_own_shard().await.map_err(AppError::from_lance)?; + let reclaimed = store + .cleanup_own_shard() + .await + .map_err(AppError::from_lance)?; if reclaimed > 0 { ::metrics::counter!("rollout_wal_cleanup_total", "result" => "merged").increment(1); ::metrics::counter!("rollout_wal_generations_reclaimed_total").increment(reclaimed as u64);