diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b8acd7c..cd1f66b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -56,7 +56,7 @@ jobs: - name: Install system dependencies run: | sudo apt update - sudo apt install -y protobuf-compiler libssl-dev + sudo apt install -y clang libclang-dev protobuf-compiler libssl-dev - uses: actions-rust-lang/setup-rust-toolchain@v1 with: diff --git a/.github/workflows/rust-publish.yml b/.github/workflows/rust-publish.yml index 7906824..90fd193 100644 --- a/.github/workflows/rust-publish.yml +++ b/.github/workflows/rust-publish.yml @@ -36,7 +36,7 @@ jobs: - name: Install dependencies run: | sudo apt update - sudo apt install -y protobuf-compiler libssl-dev + sudo apt install -y clang libclang-dev protobuf-compiler libssl-dev - uses: actions-rust-lang/setup-rust-toolchain@v1 with: @@ -46,10 +46,15 @@ jobs: - uses: rui314/setup-mold@v1 + - name: Validate Rust crate packages + if: github.event_name == 'pull_request' + run: cargo package --workspace --exclude lance-context-python --allow-dirty --no-verify --all-features --locked + - name: Publish Rust crates + if: github.event_name != 'pull_request' uses: katyo/publish-crates@v2 with: registry-token: ${{ secrets.CARGO_REGISTRY_TOKEN }} args: "--all-features" path: . - dry-run: ${{ github.event_name == 'pull_request' || (github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'dry_run') }} + dry-run: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'dry_run' }} diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index 9f9b0af..7430d2f 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -43,10 +43,10 @@ jobs: - name: Install dependencies run: | sudo apt update - sudo apt install -y protobuf-compiler + sudo apt install -y clang libclang-dev protobuf-compiler - name: Build tests - run: cargo test -p lance-context-core --no-run + run: cargo test -p lance-context-core -p lance-context-master --no-run - name: Run unit tests - run: cargo test -p lance-context-core --lib + run: cargo test -p lance-context-core -p lance-context-master --lib - name: Run doc tests run: cargo test -p lance-context-core --doc diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index 6266fa1..cf83a77 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -33,7 +33,7 @@ jobs: clippy: runs-on: ubuntu-24.04 - timeout-minutes: 15 + timeout-minutes: 30 steps: - uses: actions/checkout@v4 - uses: Swatinem/rust-cache@v2 @@ -45,7 +45,7 @@ jobs: - name: Install dependencies run: | sudo apt update - sudo apt install -y protobuf-compiler + sudo apt install -y clang libclang-dev protobuf-compiler - name: Clippy run: cargo clippy --workspace --all-targets -- -D warnings diff --git a/.gitignore b/.gitignore index 90cdcf1..8de0b92 100644 --- a/.gitignore +++ b/.gitignore @@ -74,6 +74,7 @@ docs/_build/ # PyBuilder .pybuilder/ target/ +/master-data/ # Jupyter Notebook .ipynb_checkpoints diff --git a/Cargo.lock b/Cargo.lock index 69a324a..a9145c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5483,6 +5483,7 @@ dependencies = [ "lance-context-metrics", "metrics", "reqwest 0.12.28", + "rocksdb", "serde", "serde_json", "tempfile", @@ -6566,6 +6567,30 @@ dependencies = [ "libc", ] +[[package]] +name = "librocksdb-sys" +version = "0.17.3+10.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cef2a00ee60fe526157c9023edab23943fae1ce2ab6f4abb2a807c1746835de9" +dependencies = [ + "bindgen", + "bzip2-sys", + "cc", + "libc", + "libz-sys", +] + +[[package]] +name = "libz-sys" +version = "1.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "lindera" version = "3.0.7" @@ -9096,6 +9121,16 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" +[[package]] +name = "rocksdb" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddb7af00d2b17dbd07d82c0063e25411959748ff03e8d4f96134c2ff41fce34f" +dependencies = [ + "libc", + "librocksdb-sys", +] + [[package]] name = "rsa" version = "0.9.10" diff --git a/README.md b/README.md index d0d5dc0..c4427e7 100644 --- a/README.md +++ b/README.md @@ -128,8 +128,10 @@ crates/lance-context-core # Rust engine: Context + Rollout stores (no Python crates/lance-context-api # Shared request/response types (DTOs) crates/lance-context-server # HTTP server for remote access crates/lance-context-client # HTTP client for the server +crates/lance-context-master # Control plane, scheduler, and admin UI crates/lance-context # Re-export crate for downstream clients python/ # Python bindings (PyO3) + tests +deploy/kubernetes/ # Kubernetes master + RocksDB PVC example examples/ # Runnable example projects ``` diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index aaa4d26..4d340ed 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -25,6 +25,7 @@ futures = "0.3" lance = "7.0.0" metrics = "0.24" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +rocksdb = { version = "0.24", default-features = false } 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 b81a9a0..4028dfe 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -56,6 +56,22 @@ pub struct MasterConfig { #[arg(long, env = "TASK_CONCURRENCY", default_value_t = 4)] pub task_concurrency: usize, + /// Local RocksDB directory for durable scheduler task state. This must be + /// backed by persistent local storage in production; it must not be an + /// object-store URI or a volume shared by multiple master processes. + #[arg( + long, + env = "TASK_DB_PATH", + default_value = "./master-data/tasks.rocksdb" + )] + pub task_db_path: String, + + /// Maximum number of terminal scheduler tasks retained in RocksDB and the + /// queue UI. Queued and running tasks are never pruned, and at least one + /// terminal task is retained for status polling. + #[arg(long, env = "TASK_HISTORY_LIMIT", default_value_t = 1_000)] + pub task_history_limit: 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/lib.rs b/crates/lance-context-master/src/lib.rs index 505d6ae..1744518 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -8,3 +8,4 @@ pub mod scanner; pub mod scheduler; pub mod state; pub mod stats_store; +pub mod task_store; diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs index 5498304..a13f37e 100644 --- a/crates/lance-context-master/src/main.rs +++ b/crates/lance-context-master/src/main.rs @@ -22,6 +22,10 @@ async fn main() { let config = MasterConfig::parse(); let addr = format!("{}:{}", config.host, config.port); + // Install the recorder before state recovery and scheduler startup so + // recovered queue depth is represented immediately. + let metrics_handle = lance_context_metrics::install_recorder(); + if let Err(e) = create_local_dir_if_needed(&config.data_dir) { tracing::error!( "Failed to create local data directory '{}': {}", @@ -47,9 +51,6 @@ async fn main() { // Single serial compaction worker (+ optional periodic auto-sweep). let _scheduler = scheduler::spawn_scheduler(&state); - // Install the Prometheus recorder once, before any metrics are emitted. - let metrics_handle = lance_context_metrics::install_recorder(); - let mut app = Router::new().nest("/api/v1", routes::api_router()); // Serve the built UI (SPA) as a fallback when a `--ui-dir` is configured. diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 1b4468f..338f8cb 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -277,7 +277,9 @@ pub async fn compact_experiment( name ))); } - let task = scheduler::enqueue(&state, TaskKind::Compact, &name).await; + let task = scheduler::enqueue(&state, TaskKind::Compact, &name) + .await + .map_err(MasterError::from_lance)?; Ok((StatusCode::ACCEPTED, Json(task_to_compact_status(&task)))) } @@ -312,7 +314,9 @@ pub async fn enqueue_task( req.target ))); } - let task = scheduler::enqueue(&state, req.kind, &req.target).await; + let task = scheduler::enqueue(&state, req.kind, &req.target) + .await + .map_err(MasterError::from_lance)?; Ok((StatusCode::ACCEPTED, Json(task))) } @@ -418,6 +422,12 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_db_path: dir + .path() + .join("master-tasks") + .to_string_lossy() + .to_string(), + task_history_limit: 1_000, ui_dir: None, } } diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 91c0c36..f155a40 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -1,7 +1,10 @@ //! Unified task scheduler. //! -//! The master runs a single scheduler that drains one queue with **bounded -//! concurrency**. It executes two kinds of task ([`TaskKind`]): +//! The master runs a single scheduler that drains one durable queue with +//! **bounded concurrency**. Task records are written to RocksDB before they +//! enter the in-memory dispatch channel; queued and interrupted tasks are +//! reloaded after a master restart. It executes three 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 @@ -30,7 +33,7 @@ use lance_context_core::{generate_id, CompactionConfig, RolloutStore, RolloutSto use tokio::sync::Semaphore; use tokio::task::JoinHandle; -use crate::state::MasterState; +use crate::state::{prune_terminal_history, MasterState}; use crate::stats_store::StatRow; /// Build the [`CompactionConfig`] the scheduler applies, from master config. @@ -61,15 +64,19 @@ fn kind_label(kind: TaskKind) -> &'static str { /// 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 { +pub async fn enqueue( + state: &Arc, + kind: TaskKind, + target: &str, +) -> lance::Result { + let mut tasks = state.tasks.lock().await; if matches!(kind, TaskKind::Compact | TaskKind::IndexId) { - let tasks = state.tasks.lock().await; if let Some(existing) = tasks.values().find(|t| { t.kind == kind && t.target == target && matches!(t.state, TaskState::Queued | TaskState::Running) }) { - return existing.clone(); + return Ok(existing.clone()); } } @@ -84,23 +91,39 @@ pub async fn enqueue(state: &Arc, kind: TaskKind, target: &str) -> started_at: None, finished_at: None, }; - state - .tasks - .lock() - .await - .insert(record.id.clone(), record.clone()); + state.task_store.put(&record)?; + tasks.insert(record.id.clone(), record.clone()); + drop(tasks); // Unbounded send only fails if the receiver was dropped (scheduler gone). 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 + Ok(record) } -/// 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); +/// Persist a task mutation before publishing it to the in-memory cache. +async fn update_task( + state: &Arc, + id: &str, + f: impl FnOnce(&mut TaskRecord), +) -> lance::Result<()> { + let mut tasks = state.tasks.lock().await; + let Some(current) = tasks.get(id) else { + return Ok(()); + }; + let mut updated = current.clone(); + f(&mut updated); + state.task_store.put(&updated)?; + let terminal = matches!(updated.state, TaskState::Done | TaskState::Failed); + tasks.insert(id.to_string(), updated); + if terminal { + prune_terminal_history( + &state.task_store, + &mut tasks, + state.config.task_history_limit, + )?; } + Ok(()) } /// Execute one task by id: dispatch on its kind, recording lifecycle timestamps @@ -115,7 +138,27 @@ async fn run_task(state: &Arc, id: String) { t.state = TaskState::Running; t.started_at = Some(now_ms()); }) - .await; + .await + .unwrap_or_else(|error| { + tracing::error!(task = %id, error = %error, "failed to persist task start"); + }); + + let task_started = state + .tasks + .lock() + .await + .get(&id) + .is_some_and(|task| task.state == TaskState::Running); + if !task_started { + metrics::gauge!("master_task_queue_depth").increment(1.0); + let retry_tx = state.task_tx.clone(); + let retry_id = id.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(5)).await; + let _ = retry_tx.send(retry_id); + }); + return; + } let started = std::time::Instant::now(); let outcome = match task.kind { @@ -144,7 +187,10 @@ async fn run_task(state: &Arc, id: String) { } } }) - .await; + .await + .unwrap_or_else(|error| { + tracing::error!(task = %id, error = %error, "failed to persist task completion"); + }); } /// Compact one experiment. Serialized per experiment via the shared base-table @@ -324,7 +370,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, TaskKind::Compact, &row.name).await; + enqueue(state, TaskKind::Compact, &row.name).await?; queued += 1; } } @@ -374,6 +420,17 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { } let concurrency = state.config.task_concurrency.max(1); + let queue_depth = state + .tasks + .try_lock() + .map(|tasks| { + tasks + .values() + .filter(|task| task.state == TaskState::Queued) + .count() + }) + .unwrap_or(0); + metrics::gauge!("master_task_queue_depth").set(queue_depth as f64); let sem = Arc::new(Semaphore::new(concurrency)); let dispatch_state = state.clone(); tokio::spawn(async move { @@ -412,6 +469,12 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_db_path: dir + .path() + .join("master-tasks") + .to_string_lossy() + .to_string(), + task_history_limit: 1_000, ui_dir: None, } } @@ -458,9 +521,20 @@ mod tests { // Seed a stats row so post-compaction upsert has a prior counter. crate::scanner::scan_once(&state).await.unwrap(); - let rec = enqueue(&state, TaskKind::Compact, name).await; + let rec = enqueue(&state, TaskKind::Compact, name).await.unwrap(); let status = await_terminal(&state, &rec.id).await; assert_eq!(status.state, TaskState::Done, "got {status:?}"); + assert_eq!( + state + .task_store + .list() + .unwrap() + .into_iter() + .find(|task| task.id == rec.id) + .unwrap() + .state, + TaskState::Done + ); let row = state.stats.lock().await.get(name).await.unwrap().unwrap(); assert_eq!(row.total_compactions, 1); @@ -495,7 +569,7 @@ mod tests { .await .unwrap(); - let rec = enqueue(&state, TaskKind::IndexId, name).await; + let rec = enqueue(&state, TaskKind::IndexId, name).await.unwrap(); let status = await_terminal(&state, &rec.id).await; assert_eq!(status.state, TaskState::Done, "got {status:?}"); assert_eq!(status.detail.as_deref(), Some("built zonemap index on id")); @@ -509,11 +583,12 @@ mod tests { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.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; + let a = enqueue(&state, TaskKind::Compact, "x").await.unwrap(); + let b = enqueue(&state, TaskKind::Compact, "x").await.unwrap(); 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); + assert_eq!(state.task_store.list().unwrap().len(), 1); let mut rx = state.task_rx.lock().await.take().unwrap(); let mut n = 0; while rx.try_recv().is_ok() { @@ -528,7 +603,7 @@ mod tests { 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 rec = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); let status = await_terminal(&state, &rec.id).await; assert_eq!(status.state, TaskState::Failed); assert!(status.error.unwrap().contains("worker endpoints")); @@ -561,7 +636,7 @@ mod tests { let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); - let rec = enqueue(&state, TaskKind::MergeWal, "exp").await; + let rec = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); let status = await_terminal(&state, &rec.id).await; assert_eq!(status.state, TaskState::Done, "got {status:?}"); let detail = status.detail.unwrap(); diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index 4a0b04b..af96a97 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -3,13 +3,14 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use lance_context_api::TaskRecord; +use lance_context_api::{TaskRecord, TaskState}; use lance_context_core::{join_uri, RolloutRegistry}; use tokio::sync::{mpsc, Mutex, RwLock}; use crate::config::MasterConfig; use crate::discovery; use crate::stats_store::StatsStore; +use crate::task_store::TaskStore; /// Shared state for the master process. /// @@ -33,8 +34,11 @@ pub struct MasterState { pub task_tx: mpsc::UnboundedSender, /// Receiver half, taken once by [`crate::scheduler::spawn_scheduler`]. pub task_rx: Mutex>>, + /// Durable task records on the master-local RocksDB PVC. + pub task_store: TaskStore, /// 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. + /// API, de-duplication, and UI. This is a cache of `task_store`, rebuilt at + /// startup. Terminal history is bounded by `task_history_limit`. pub tasks: Mutex>, /// Experiment names with a base-table write currently executing — the /// per-name serial gate. A `Compact` or `IndexId` task must not run while @@ -47,7 +51,7 @@ pub struct MasterState { } impl MasterState { - /// Open (or create) the registry and stats datasets under `data_dir`. + /// Open the registry, stats dataset, and local durable task store. pub async fn new(config: MasterConfig) -> lance::Result> { let base_uri = config.data_dir.clone(); let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); @@ -61,18 +65,59 @@ impl MasterState { ); } let stats = StatsStore::open_or_create(&stats_uri, None).await?; + let task_store = TaskStore::open(&config.task_db_path)?; + let mut tasks = task_store + .list()? + .into_iter() + .map(|task| (task.id.clone(), task)) + .collect::>(); + + let mut recovered_running = 0usize; + for task in tasks.values_mut() { + if task.state == TaskState::Running { + task.state = TaskState::Queued; + task.started_at = None; + task.finished_at = None; + task.error = None; + task.detail = None; + task_store.put(task)?; + recovered_running += 1; + } + } + let pruned = prune_terminal_history(&task_store, &mut tasks, config.task_history_limit)?; + let mut pending = tasks + .values() + .filter(|task| task.state == TaskState::Queued) + .map(|task| (task.enqueued_at, task.id.clone())) + .collect::>(); + pending.sort(); + let pending_ids = pending.into_iter().map(|(_, id)| id).collect::>(); + let (task_tx, task_rx) = mpsc::unbounded_channel(); - Ok(Arc::new(Self { + let state = Arc::new(Self { registry: RwLock::new(registry), stats: Mutex::new(stats), base_uri, config, task_tx, task_rx: Mutex::new(Some(task_rx)), - tasks: Mutex::new(HashMap::new()), + task_store, + tasks: Mutex::new(tasks), inflight_dataset_writes: Mutex::new(HashSet::new()), http: reqwest::Client::new(), - })) + }); + for id in &pending_ids { + let _ = state.task_tx.send(id.clone()); + } + if !pending_ids.is_empty() || recovered_running > 0 || pruned > 0 { + tracing::info!( + queued = pending_ids.len(), + recovered_running, + pruned, + "loaded durable scheduler tasks" + ); + } + Ok(state) } /// Physical rollout dataset URI for `name`, matching the data-plane's @@ -82,9 +127,41 @@ impl MasterState { } } +pub(crate) fn prune_terminal_history( + task_store: &TaskStore, + tasks: &mut HashMap, + limit: usize, +) -> lance::Result { + // Keep at least the latest terminal task so legacy per-experiment status + // polling can observe completion even if the configured limit is zero. + let limit = limit.max(1); + let mut terminal = tasks + .values() + .filter(|task| matches!(task.state, TaskState::Done | TaskState::Failed)) + .map(|task| { + ( + task.finished_at.unwrap_or(task.enqueued_at), + task.id.clone(), + ) + }) + .collect::>(); + terminal.sort_by_key(|(finished_at, _)| std::cmp::Reverse(*finished_at)); + let ids = terminal + .into_iter() + .skip(limit) + .map(|(_, id)| id) + .collect::>(); + task_store.delete_many(ids.iter().map(String::as_str))?; + for id in &ids { + tasks.remove(id); + } + Ok(ids.len()) +} + #[cfg(test)] mod tests { use super::*; + use lance_context_api::TaskKind; use lance_context_core::RolloutStore; use tempfile::TempDir; @@ -100,6 +177,12 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_db_path: dir + .path() + .join("master-tasks") + .to_string_lossy() + .to_string(), + task_history_limit: 1_000, ui_dir: None, } } @@ -117,4 +200,79 @@ mod tests { assert_eq!(entries[0].name, "legacy"); assert_eq!(entries[0].uri, uri.to_string_lossy().to_string()); } + + #[tokio::test] + async fn startup_requeues_queued_and_interrupted_tasks() { + let dir = TempDir::new().unwrap(); + let config = test_config(&dir); + { + let state = MasterState::new(config.clone()).await.unwrap(); + for (id, task_state) in [ + ("queued-task", TaskState::Queued), + ("running-task", TaskState::Running), + ("done-task", TaskState::Done), + ] { + let task = TaskRecord { + id: id.to_string(), + kind: TaskKind::Compact, + target: "legacy".to_string(), + state: task_state, + error: None, + detail: None, + enqueued_at: 1, + started_at: (task_state == TaskState::Running).then_some(2), + finished_at: (task_state == TaskState::Done).then_some(3), + }; + state.task_store.put(&task).unwrap(); + } + } + + let state = MasterState::new(config).await.unwrap(); + let tasks = state.tasks.lock().await; + assert_eq!(tasks["queued-task"].state, TaskState::Queued); + assert_eq!(tasks["running-task"].state, TaskState::Queued); + assert_eq!(tasks["running-task"].started_at, None); + assert_eq!(tasks["done-task"].state, TaskState::Done); + drop(tasks); + + let mut rx = state.task_rx.lock().await.take().unwrap(); + let mut recovered = Vec::new(); + while let Ok(id) = rx.try_recv() { + recovered.push(id); + } + recovered.sort(); + assert_eq!(recovered, ["queued-task", "running-task"]); + } + + #[test] + fn terminal_history_pruning_keeps_active_tasks() { + let dir = TempDir::new().unwrap(); + let store = TaskStore::open(dir.path().join("tasks")).unwrap(); + let mut tasks = HashMap::new(); + for (id, state, timestamp) in [ + ("queued", TaskState::Queued, 1), + ("old", TaskState::Done, 2), + ("new", TaskState::Failed, 3), + ] { + let task = TaskRecord { + id: id.to_string(), + kind: TaskKind::Compact, + target: "experiment".to_string(), + state, + error: None, + detail: None, + enqueued_at: timestamp, + started_at: None, + finished_at: matches!(state, TaskState::Done | TaskState::Failed) + .then_some(timestamp), + }; + store.put(&task).unwrap(); + tasks.insert(id.to_string(), task); + } + + assert_eq!(prune_terminal_history(&store, &mut tasks, 1).unwrap(), 1); + assert!(tasks.contains_key("queued")); + assert!(tasks.contains_key("new")); + assert!(!tasks.contains_key("old")); + } } diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs new file mode 100644 index 0000000..b60926b --- /dev/null +++ b/crates/lance-context-master/src/task_store.rs @@ -0,0 +1,168 @@ +//! Durable scheduler task storage backed by RocksDB. + +use std::path::Path; + +use lance_context_api::TaskRecord; +use rocksdb::{Direction, IteratorMode, Options, WriteBatch, WriteOptions, DB}; + +const TASK_KEY_PREFIX: &[u8] = b"task/"; + +/// Single-process RocksDB store for scheduler task records. +/// +/// Every write uses RocksDB's WAL with `sync=true`: the scheduler never +/// acknowledges an enqueue or lifecycle transition before it is durable on +/// the master PVC. +pub struct TaskStore { + db: DB, +} + +impl TaskStore { + pub fn open(path: impl AsRef) -> lance::Result { + let path = path.as_ref(); + let text = path.to_string_lossy(); + if text.contains("://") { + return Err(lance::Error::io(format!( + "task DB path must be local, got '{text}'" + ))); + } + if let Some(parent) = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + std::fs::create_dir_all(parent).map_err(|err| { + lance::Error::io(format!( + "failed to create task DB parent '{}': {err}", + parent.display() + )) + })?; + } + + let mut options = Options::default(); + options.create_if_missing(true); + options.set_max_open_files(64); + options.set_keep_log_file_num(4); + let db = DB::open(&options, path).map_err(|err| { + lance::Error::io(format!( + "failed to open task DB '{}': {err}", + path.display() + )) + })?; + Ok(Self { db }) + } + + pub fn put(&self, task: &TaskRecord) -> lance::Result<()> { + let value = serde_json::to_vec(task).map_err(|err| { + lance::Error::io(format!("failed to encode task '{}': {err}", task.id)) + })?; + self.db + .put_opt(task_key(&task.id), value, &sync_write_options()) + .map_err(|err| lance::Error::io(format!("failed to persist task '{}': {err}", task.id))) + } + + pub fn list(&self) -> lance::Result> { + let mut tasks = Vec::new(); + for item in self + .db + .iterator(IteratorMode::From(TASK_KEY_PREFIX, Direction::Forward)) + { + let (key, value) = + item.map_err(|err| lance::Error::io(format!("failed to iterate task DB: {err}")))?; + if !key.starts_with(TASK_KEY_PREFIX) { + break; + } + let task = serde_json::from_slice(&value).map_err(|err| { + lance::Error::io(format!( + "failed to decode task record '{}': {err}", + String::from_utf8_lossy(&key) + )) + })?; + tasks.push(task); + } + Ok(tasks) + } + + pub fn delete_many<'a>(&self, ids: impl IntoIterator) -> lance::Result<()> { + let mut batch = WriteBatch::default(); + let mut count = 0usize; + for id in ids { + batch.delete(task_key(id)); + count += 1; + } + if count == 0 { + return Ok(()); + } + self.db + .write_opt(batch, &sync_write_options()) + .map_err(|err| lance::Error::io(format!("failed to prune task history: {err}"))) + } +} + +fn task_key(id: &str) -> Vec { + let mut key = Vec::with_capacity(TASK_KEY_PREFIX.len() + id.len()); + key.extend_from_slice(TASK_KEY_PREFIX); + key.extend_from_slice(id.as_bytes()); + key +} + +fn sync_write_options() -> WriteOptions { + let mut options = WriteOptions::default(); + options.set_sync(true); + options +} + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_api::{TaskKind, TaskState}; + use tempfile::TempDir; + + fn task(id: &str, state: TaskState, enqueued_at: i64) -> TaskRecord { + TaskRecord { + id: id.to_string(), + kind: TaskKind::Compact, + target: "experiment".to_string(), + state, + error: None, + detail: None, + enqueued_at, + started_at: None, + finished_at: None, + } + } + + #[test] + fn survives_reopen_and_deletes_in_batch() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("tasks"); + { + let store = TaskStore::open(&path).unwrap(); + store.put(&task("a", TaskState::Queued, 1)).unwrap(); + store.put(&task("b", TaskState::Done, 2)).unwrap(); + } + + let store = TaskStore::open(&path).unwrap(); + let mut tasks = store.list().unwrap(); + tasks.sort_by_key(|task| task.enqueued_at); + assert_eq!( + tasks + .iter() + .map(|task| task.id.as_str()) + .collect::>(), + ["a", "b"] + ); + + store.delete_many(["a"]).unwrap(); + let tasks = store.list().unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "b"); + } + + #[test] + fn rejects_object_store_uri() { + let error = match TaskStore::open("s3://bucket/tasks") { + Ok(_) => panic!("object-store URI should be rejected"), + Err(error) => error, + }; + assert!(error.to_string().contains("must be local")); + } +} diff --git a/deploy/kubernetes/README.md b/deploy/kubernetes/README.md new file mode 100644 index 0000000..01450a2 --- /dev/null +++ b/deploy/kubernetes/README.md @@ -0,0 +1,17 @@ +# Kubernetes deployment + +`master.yaml` deploys one `lance-context-master` process with a dedicated +ReadWriteOnce PVC for its RocksDB scheduler state. + +Before applying it: + +1. Replace `DATA_DIR` with the shared object-store prefix used by rollout + workers. +2. Set object-store credentials and `WORKER_ENDPOINTS` for your environment. +3. Replace the example image name and `UI_DIR` with paths from the published + master image. +4. Set `storageClassName` on the PVC when the cluster has no default class. + +Keep the master at one replica. RocksDB is an embedded single-process database; +the PVC must not be mounted by multiple active master pods. The `Recreate` +strategy ensures an old pod releases the volume before its replacement starts. diff --git a/deploy/kubernetes/master.yaml b/deploy/kubernetes/master.yaml new file mode 100644 index 0000000..165ca10 --- /dev/null +++ b/deploy/kubernetes/master.yaml @@ -0,0 +1,87 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: lance-context-master-data +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 10Gi +--- +apiVersion: v1 +kind: Service +metadata: + name: lance-context-master +spec: + selector: + app: lance-context-master + ports: + - name: http + port: 8090 + targetPort: http +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: lance-context-master +spec: + # RocksDB is single-process. Run one master against this PVC. + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: lance-context-master + template: + metadata: + labels: + app: lance-context-master + spec: + securityContext: + fsGroup: 1000 + containers: + - name: master + image: lance-context-master:latest + imagePullPolicy: IfNotPresent + env: + # Replace with the same shared object-store prefix used by workers. + - name: DATA_DIR + value: s3://replace-me/rollouts + - name: TASK_DB_PATH + value: /var/lib/lance-context-master/tasks.rocksdb + - name: TASK_HISTORY_LIMIT + value: "1000" + - name: MASTER_PORT + value: "8090" + # Adjust to the built UI asset path in the master image. + - name: UI_DIR + value: /app/ui + ports: + - name: http + containerPort: 8090 + readinessProbe: + httpGet: + path: /metrics + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /metrics + port: http + initialDelaySeconds: 15 + periodSeconds: 20 + resources: + requests: + cpu: 250m + memory: 512Mi + limits: + memory: 2Gi + volumeMounts: + - name: task-db + mountPath: /var/lib/lance-context-master + volumes: + - name: task-db + persistentVolumeClaim: + claimName: lance-context-master-data