diff --git a/Cargo.lock b/Cargo.lock index 1f6e1a1..536b329 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5515,7 +5515,6 @@ dependencies = [ "lance-context-metrics", "metrics", "reqwest 0.12.28", - "rocksdb", "serde", "serde_json", "tempfile", @@ -6599,30 +6598,6 @@ 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" @@ -9175,16 +9150,6 @@ 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/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index ec93c5b..351c81b 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -26,7 +26,6 @@ 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 f8fa1dc..eef29f1 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -1,15 +1,6 @@ //! Command-line / environment configuration for the master control-plane. -use clap::{Parser, ValueEnum}; - -/// Durable scheduler state backend. -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] -pub enum TaskStoreBackend { - /// Single-master RocksDB on a dedicated local PVC. - Rocksdb, - /// Shared etcd state with lease-based task claims for HA masters. - Etcd, -} +use clap::Parser; /// Control-plane (master) process for lance-context rollout stores. #[derive(Debug, Clone, Parser)] @@ -65,29 +56,9 @@ pub struct MasterConfig { #[arg(long, env = "TASK_CONCURRENCY", default_value_t = 4)] pub task_concurrency: usize, - /// Scheduler state backend. `rocksdb` requires exactly one master and a - /// dedicated PVC; `etcd` supports multiple stateless master replicas. - #[arg( - long, - env = "TASK_STORE_BACKEND", - value_enum, - default_value_t = TaskStoreBackend::Rocksdb - )] - pub task_store_backend: TaskStoreBackend, - - /// 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. Used - /// only when `TASK_STORE_BACKEND=rocksdb`. - #[arg( - long, - env = "TASK_DB_PATH", - default_value = "./master-data/tasks.rocksdb" - )] - pub task_db_path: String, - - /// Comma-separated etcd v3 endpoints. Required when - /// `TASK_STORE_BACKEND=etcd`. + /// Comma-separated etcd v3 endpoints. Scheduler state (task queue, + /// lease-based claims, per-experiment write locks) lives in etcd so several + /// stateless master replicas can share one queue. Required. #[arg(long, env = "ETCD_ENDPOINTS", value_delimiter = ',')] pub etcd_endpoints: Vec, @@ -120,7 +91,7 @@ pub struct MasterConfig { #[arg(long, env = "ETCD_LEASE_TTL_SECS", default_value_t = 30)] pub etcd_lease_ttl_secs: i64, - /// Maximum number of terminal scheduler tasks retained in RocksDB and the + /// Maximum number of terminal scheduler tasks retained in etcd 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)] diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index cb1c53b..1713f4f 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -422,9 +422,9 @@ fn blob_filename(record: &lance_context_core::RolloutRecord) -> String { #[cfg(test)] mod tests { use super::*; - use crate::config::{MasterConfig, TaskStoreBackend}; + use crate::config::MasterConfig; use chrono::Utc; - use lance_context_core::{RolloutRecord, RolloutRegistry, RolloutStore}; + use lance_context_core::{generate_id, RolloutRecord, RolloutRegistry, RolloutStore}; use serde_json::json; use tempfile::TempDir; @@ -440,25 +440,25 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, - task_store_backend: TaskStoreBackend::Rocksdb, - task_db_path: dir - .path() - .join("master-tasks") - .to_string_lossy() - .to_string(), - etcd_endpoints: vec![], - etcd_prefix: "/test".to_string(), + etcd_endpoints: test_etcd_endpoints(), + etcd_prefix: format!("/lance-context/test/{}", generate_id()), etcd_username: None, etcd_password: None, etcd_ca_cert: None, etcd_client_cert: None, etcd_client_key: None, - etcd_lease_ttl_secs: 30, + etcd_lease_ttl_secs: 5, task_history_limit: 1_000, ui_dir: None, } } + fn test_etcd_endpoints() -> Vec { + std::env::var("ETCD_TEST_ENDPOINTS") + .map(|value| value.split(',').map(str::to_string).collect()) + .unwrap_or_default() + } + fn test_record(id: &str, with_blob: bool) -> RolloutRecord { let bytes = with_blob.then(|| b"master-blob".to_vec()); RolloutRecord { @@ -499,6 +499,7 @@ mod tests { /// Create N experiments via core, register them, scan, and assert the stats /// table + list endpoint reflect them. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn scan_populates_and_list_returns_experiments() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -550,6 +551,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn scan_sees_registry_commits_from_another_handle() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -570,6 +572,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn scan_reconciles_removed_experiments() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -594,6 +597,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn get_experiment_not_found() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -607,6 +611,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn fresh_detail_refreshes_only_requested_experiment() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -650,6 +655,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn records_endpoint_filters_pages_and_rejects_unknown_experiment() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); @@ -716,6 +722,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn blob_endpoint_sets_download_headers_and_returns_bytes() { let dir = TempDir::new().unwrap(); let state = MasterState::new(test_config(&dir)).await.unwrap(); diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index dd64cd1..7181662 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -1,8 +1,8 @@ //! Unified task scheduler. //! -//! Each master polls one durable queue with **bounded concurrency**. RocksDB -//! supports one master; etcd uses atomic lease-backed claims so multiple -//! stateless masters can drain the same queue. It executes three kinds of task +//! Each master polls one durable queue with **bounded concurrency**. The queue +//! lives in etcd, which uses atomic lease-backed claims so multiple stateless +//! masters can drain the same queue. It executes three kinds of task //! ([`TaskKind`]): //! //! - **Compact** — rewrites an experiment's base-table fragments. Two `Rewrite`s @@ -362,8 +362,9 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { #[cfg(test)] mod tests { use super::*; - use crate::config::{MasterConfig, TaskStoreBackend}; + use crate::config::MasterConfig; use lance_context_api::TaskState; + use lance_context_core::generate_id; use tempfile::TempDir; fn config(dir: &TempDir) -> MasterConfig { @@ -379,20 +380,16 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, - task_store_backend: TaskStoreBackend::Rocksdb, - task_db_path: dir - .path() - .join("master-tasks") - .to_string_lossy() - .to_string(), - etcd_endpoints: vec![], - etcd_prefix: "/test".to_string(), + etcd_endpoints: std::env::var("ETCD_TEST_ENDPOINTS") + .map(|value| value.split(',').map(str::to_string).collect()) + .unwrap_or_default(), + etcd_prefix: format!("/lance-context/test/{}", generate_id()), etcd_username: None, etcd_password: None, etcd_ca_cert: None, etcd_client_cert: None, etcd_client_key: None, - etcd_lease_ttl_secs: 30, + etcd_lease_ttl_secs: 5, task_history_limit: 1_000, ui_dir: None, } @@ -414,6 +411,7 @@ mod tests { /// Manual enqueue -> dispatcher compacts -> task reaches Done and the stats /// table records a compaction. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn manual_compaction_runs_and_updates_stats() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); @@ -466,6 +464,7 @@ mod tests { /// Manual enqueue of an `IndexId` task -> dispatcher builds the ZoneMap /// index -> task reaches Done with the expected detail summary. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn index_id_task_builds_index_and_reaches_done() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); @@ -499,6 +498,7 @@ mod tests { /// Enqueuing the same experiment twice while queued de-dupes to one task. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn enqueue_dedupes_queued_compactions() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); @@ -511,6 +511,7 @@ mod tests { /// A MergeWal task with no configured endpoints fails fast with a clear msg. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn merge_wal_without_endpoints_fails() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); @@ -525,6 +526,7 @@ mod tests { /// MergeWal fans out to every configured worker endpoint and sums the /// reclaimed counts. Uses a tiny in-process stub server per "worker". #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn merge_wal_broadcasts_and_sums_reclaimed() { use axum::{routing::post, Json, Router}; @@ -561,6 +563,7 @@ mod tests { /// `index_id` depending on a `compact` must start after compaction finishes, /// so the two never contend for the shared per-experiment base-table gate. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn dependent_task_runs_after_dependency_done() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); @@ -612,6 +615,7 @@ mod tests { /// A dependent whose dependency `Failed` is skipped (marked `Failed`) rather /// than run. #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn dependent_skipped_when_dependency_fails() { let dir = TempDir::new().unwrap(); let state = MasterState::new(config(&dir)).await.unwrap(); diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index cff0524..8eab409 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -26,8 +26,8 @@ pub struct MasterState { pub base_uri: String, /// Effective configuration. pub config: MasterConfig, - /// Durable scheduler queue. RocksDB is single-master; etcd provides shared - /// CAS/lease semantics for stateless HA master replicas. + /// Durable scheduler queue in etcd, providing shared CAS/lease semantics + /// for stateless HA master replicas. pub task_store: TaskStore, /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. pub http: reqwest::Client, @@ -79,9 +79,8 @@ impl MasterState { #[cfg(test)] mod tests { use super::*; - use crate::config::TaskStoreBackend; use lance_context_api::{TaskKind, TaskState}; - use lance_context_core::RolloutStore; + use lance_context_core::{generate_id, RolloutStore}; use tempfile::TempDir; fn test_config(dir: &TempDir) -> MasterConfig { @@ -96,26 +95,23 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, - task_store_backend: TaskStoreBackend::Rocksdb, - task_db_path: dir - .path() - .join("master-tasks") - .to_string_lossy() - .to_string(), - etcd_endpoints: vec![], - etcd_prefix: "/test".to_string(), + etcd_endpoints: std::env::var("ETCD_TEST_ENDPOINTS") + .map(|value| value.split(',').map(str::to_string).collect()) + .unwrap_or_default(), + etcd_prefix: format!("/lance-context/test/{}", generate_id()), etcd_username: None, etcd_password: None, etcd_ca_cert: None, etcd_client_cert: None, etcd_client_key: None, - etcd_lease_ttl_secs: 30, + etcd_lease_ttl_secs: 5, task_history_limit: 1_000, ui_dir: None, } } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn startup_backfills_legacy_rollout_datasets() { let dir = TempDir::new().unwrap(); let uri = dir.path().join("legacy.rollout.lance"); @@ -130,6 +126,7 @@ mod tests { } #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn startup_requeues_interrupted_local_tasks() { let dir = TempDir::new().unwrap(); let config = test_config(&dir); diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 4d8dd70..5479047 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -1,11 +1,13 @@ -//! Durable scheduler task storage. +//! Durable scheduler task storage backed by etcd. //! -//! RocksDB is the single-master backend. etcd adds compare-and-swap enqueue, -//! lease-backed task claims, and distributed target locks so several stateless -//! masters can safely drain one shared queue. +//! Scheduler state lives entirely in etcd: compare-and-swap enqueue, +//! lease-backed task claims, and distributed per-experiment write locks so +//! several stateless masters can safely drain one shared queue. Because claims +//! ride renewed leases, a master that disappears has its in-flight task requeued +//! after the lease expires — execution is at-least-once, so task +//! implementations must be idempotent. use std::collections::HashSet; -use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -16,72 +18,47 @@ use etcd_client::{ }; use lance_context_api::{TaskKind, TaskRecord, TaskState}; use lance_context_core::generate_id; -use rocksdb::{Direction, IteratorMode, Options, WriteBatch, WriteOptions, DB}; -use tokio::sync::{oneshot, Mutex}; +use tokio::sync::oneshot; use tokio::task::JoinHandle; -use crate::config::{MasterConfig, TaskStoreBackend}; +use crate::config::MasterConfig; -const TASK_KEY_PREFIX: &[u8] = b"task/"; const TASK_POLL_BATCH: usize = 256; #[derive(Clone)] pub struct TaskStore { - inner: Arc, + inner: Arc, history_limit: usize, } -enum TaskStoreInner { - Rocks(Box), - Etcd(Box), -} - -struct RocksTaskStore { - db: DB, - operation_lock: Mutex<()>, - claimed_tasks: Mutex>, - claimed_targets: Mutex>, -} - struct EtcdTaskStore { client: Client, prefix: String, lease_ttl: i64, } -/// Ownership of one running task. Dropping an etcd claim stops lease renewal; -/// etcd then removes its claim and target-lock keys, allowing recovery. +/// Ownership of one running task. Dropping the claim stops lease renewal; etcd +/// then removes its claim and target-lock keys, allowing recovery. pub struct TaskClaim { pub task: TaskRecord, backend: ClaimBackend, } -enum ClaimBackend { - Rocks { - target_locked: bool, - }, - Etcd { - token: String, - lease_id: i64, - claim_key: String, - target_key: Option, - keepalive: LeaseKeepalive, - }, +struct ClaimBackend { + token: String, + lease_id: i64, + claim_key: String, + target_key: Option, + keepalive: LeaseKeepalive, } -/// Short-lived distributed guard used to serialize stats-table writers. +/// Short-lived distributed guard used to serialize stats-table writers across +/// master replicas. pub struct CoordinationGuard { - backend: GuardBackend, -} - -enum GuardBackend { - Rocks, - Etcd { - token: String, - lease_id: i64, - key: String, - keepalive: LeaseKeepalive, - }, + token: String, + lease_id: i64, + key: String, + keepalive: LeaseKeepalive, } struct LeaseKeepalive { @@ -100,16 +77,8 @@ impl Drop for LeaseKeepalive { impl TaskStore { pub async fn open(config: &MasterConfig) -> lance::Result { - let inner = match config.task_store_backend { - TaskStoreBackend::Rocksdb => { - TaskStoreInner::Rocks(Box::new(RocksTaskStore::open(&config.task_db_path)?)) - } - TaskStoreBackend::Etcd => { - TaskStoreInner::Etcd(Box::new(EtcdTaskStore::connect(config).await?)) - } - }; let store = Self { - inner: Arc::new(inner), + inner: Arc::new(EtcdTaskStore::connect(config).await?), history_limit: config.task_history_limit.max(1), }; store.recover_orphaned().await?; @@ -117,61 +86,35 @@ impl TaskStore { Ok(store) } - #[must_use] - pub fn is_distributed(&self) -> bool { - matches!(self.inner.as_ref(), TaskStoreInner::Etcd(_)) - } - /// Atomically enqueue a task. Standalone Compact and IndexId tasks use an - /// etcd/local dedupe key while queued or running. Tasks with dependencies - /// are always distinct because they belong to a specific ordered chain. + /// etcd dedupe key while queued or running. Tasks with dependencies are + /// always distinct because they belong to a specific ordered chain. pub async fn enqueue( &self, kind: TaskKind, target: &str, depends_on: Vec, ) -> lance::Result { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.enqueue(kind, target, depends_on).await, - TaskStoreInner::Etcd(store) => store.enqueue(kind, target, depends_on).await, - } + self.inner.enqueue(kind, target, depends_on).await } pub async fn list(&self) -> lance::Result> { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.list(), - TaskStoreInner::Etcd(store) => store.list().await, - } + self.inner.list().await } pub async fn get(&self, id: &str) -> lance::Result> { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.get(id), - TaskStoreInner::Etcd(store) => store.get(id).await, - } + self.inner.get(id).await } pub async fn queue_depth(&self) -> lance::Result { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => Ok(store - .list()? - .into_iter() - .filter(|task| task.state == TaskState::Queued) - .count()), - TaskStoreInner::Etcd(store) => store.queue_depth().await, - } + self.inner.queue_depth().await } - /// Claim the oldest runnable task. In etcd mode the queued->running update, - /// claim lease, and per-experiment write lock are one transaction. + /// Claim the oldest runnable task. The queued->running update, claim lease, + /// and per-experiment write lock are one etcd transaction. pub async fn claim_next(&self) -> lance::Result> { - let (claim, dependency_failed) = match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.claim_next().await, - TaskStoreInner::Etcd(store) => { - store.recover_orphaned().await?; - store.claim_next().await - } - }?; + self.inner.recover_orphaned().await?; + let (claim, dependency_failed) = self.inner.claim_next().await?; if dependency_failed { if let Err(error) = self.prune_terminal_history().await { tracing::warn!(error = %error, "failed to prune task history"); @@ -186,10 +129,7 @@ impl TaskStore { claim: TaskClaim, outcome: Result, ) -> lance::Result<()> { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.finish(claim, outcome).await?, - TaskStoreInner::Etcd(store) => store.finish(claim, outcome).await?, - } + self.inner.finish(claim, outcome).await?; self.prune_terminal_history().await.map(|_| ()) } @@ -200,12 +140,7 @@ impl TaskStore { &self, name: &str, ) -> lance::Result> { - match self.inner.as_ref() { - TaskStoreInner::Rocks(_) => Ok(Some(CoordinationGuard { - backend: GuardBackend::Rocks, - })), - TaskStoreInner::Etcd(store) => store.try_coordination_lock(name).await, - } + self.inner.try_coordination_lock(name).await } pub async fn coordination_lock(&self, name: &str) -> lance::Result { @@ -218,32 +153,19 @@ impl TaskStore { } pub async fn release_coordination_lock(&self, guard: CoordinationGuard) -> lance::Result<()> { - match (self.inner.as_ref(), guard.backend) { - (TaskStoreInner::Rocks(_), GuardBackend::Rocks) => Ok(()), - ( - TaskStoreInner::Etcd(store), - GuardBackend::Etcd { - token, - lease_id, - key, - keepalive, - }, - ) => { - drop(keepalive); - store.delete_owned_key(&key, &token).await?; - store.revoke_lease(lease_id).await - } - _ => Err(lance::Error::io( - "coordination guard belongs to a different task backend", - )), - } + let CoordinationGuard { + token, + lease_id, + key, + keepalive, + } = guard; + drop(keepalive); + self.inner.delete_owned_key(&key, &token).await?; + self.inner.revoke_lease(lease_id).await } async fn recover_orphaned(&self) -> lance::Result { - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.recover_running().await, - TaskStoreInner::Etcd(store) => store.recover_orphaned().await, - } + self.inner.recover_orphaned().await } async fn prune_terminal_history(&self) -> lance::Result { @@ -268,211 +190,16 @@ impl TaskStore { if ids.is_empty() { return Ok(0); } - match self.inner.as_ref() { - TaskStoreInner::Rocks(store) => store.delete_many(ids.iter().map(String::as_str))?, - TaskStoreInner::Etcd(store) => store.delete_many(&ids).await?, - } + self.inner.delete_many(&ids).await?; Ok(ids.len()) } } -impl RocksTaskStore { - 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, - operation_lock: Mutex::new(()), - claimed_tasks: Mutex::new(HashSet::new()), - claimed_targets: Mutex::new(HashSet::new()), - }) - } - - async fn enqueue( - &self, - kind: TaskKind, - target: &str, - depends_on: Vec, - ) -> lance::Result { - let _operation = self.operation_lock.lock().await; - if should_dedupe(kind, &depends_on) { - if let Some(existing) = self.list()?.into_iter().find(|task| { - task.kind == kind - && task.target == target - && task.depends_on.is_empty() - && matches!(task.state, TaskState::Queued | TaskState::Running) - }) { - return Ok(existing); - } - } - let task = new_task(kind, target, depends_on); - self.put(&task)?; - Ok(task) - } - - fn put(&self, task: &TaskRecord) -> lance::Result<()> { - let value = encode_task(task)?; - 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))) - } - - fn get(&self, id: &str) -> lance::Result> { - self.db - .get(task_key(id)) - .map_err(|err| lance::Error::io(format!("failed to read task '{id}': {err}")))? - .map(|value| decode_task(&value, id)) - .transpose() - } - - 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; - } - tasks.push(decode_task(&value, &String::from_utf8_lossy(&key))?); - } - Ok(tasks) - } - - async fn claim_next(&self) -> lance::Result<(Option, bool)> { - let _operation = self.operation_lock.lock().await; - let claimed_tasks = self.claimed_tasks.lock().await; - let mut claimed_targets = self.claimed_targets.lock().await; - let tasks = self.list()?; - let mut queued = tasks - .iter() - .filter(|task| task.state == TaskState::Queued && !claimed_tasks.contains(&task.id)) - .cloned() - .collect::>(); - queued.sort_by_key(|task| task.enqueued_at); - let mut dependency_failed = false; - let mut runnable = None; - for mut task in queued { - match dependency_status(&task, |id| { - tasks - .iter() - .find(|candidate| candidate.id == id) - .map(|task| task.state) - }) { - DependencyStatus::Ready => {} - DependencyStatus::Waiting => continue, - DependencyStatus::Failed(dependency) => { - fail_dependency(&mut task, &dependency); - self.put(&task)?; - dependency_failed = true; - continue; - } - } - if requires_target_lock(task.kind) && claimed_targets.contains(&task.target) { - continue; - } - runnable = Some(task); - break; - } - let Some(mut task) = runnable else { - return Ok((None, dependency_failed)); - }; - let target_locked = requires_target_lock(task.kind); - task.state = TaskState::Running; - task.started_at = Some(now_ms()); - self.put(&task)?; - if target_locked { - claimed_targets.insert(task.target.clone()); - } - drop(claimed_targets); - drop(claimed_tasks); - self.claimed_tasks.lock().await.insert(task.id.clone()); - Ok(( - Some(TaskClaim { - task, - backend: ClaimBackend::Rocks { target_locked }, - }), - dependency_failed, - )) - } - - async fn finish(&self, claim: TaskClaim, outcome: Result) -> lance::Result<()> { - let ClaimBackend::Rocks { target_locked } = claim.backend else { - return Err(lance::Error::io("task claim belongs to etcd")); - }; - let _operation = self.operation_lock.lock().await; - let mut task = claim.task; - apply_outcome(&mut task, outcome); - let result = self.put(&task); - self.claimed_tasks.lock().await.remove(&task.id); - if target_locked { - self.claimed_targets.lock().await.remove(&task.target); - } - result - } - - async fn recover_running(&self) -> lance::Result { - let _operation = self.operation_lock.lock().await; - let mut recovered = 0; - for mut task in self.list()? { - if task.state == TaskState::Running { - requeue(&mut task); - self.put(&task)?; - recovered += 1; - } - } - Ok(recovered) - } - - fn delete_many<'a>(&self, ids: impl IntoIterator) -> lance::Result<()> { - let mut batch = WriteBatch::default(); - let mut count = 0; - 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}"))) - } -} - impl EtcdTaskStore { async fn connect(config: &MasterConfig) -> lance::Result { if config.etcd_endpoints.is_empty() { return Err(lance::Error::io( - "ETCD_ENDPOINTS is required when TASK_STORE_BACKEND=etcd", + "ETCD_ENDPOINTS is required to run the master", )); } if config.etcd_lease_ttl_secs < 5 { @@ -709,7 +436,7 @@ impl EtcdTaskStore { return Ok(( Some(TaskClaim { task, - backend: ClaimBackend::Etcd { + backend: ClaimBackend { token, lease_id, claim_key, @@ -734,16 +461,13 @@ impl EtcdTaskStore { } async fn finish(&self, claim: TaskClaim, outcome: Result) -> lance::Result<()> { - let ClaimBackend::Etcd { + let ClaimBackend { token, lease_id, claim_key, target_key, keepalive, - } = claim.backend - else { - return Err(lance::Error::io("task claim belongs to RocksDB")); - }; + } = claim.backend; let mut task = claim.task; apply_outcome(&mut task, outcome); let mut operations = vec![ @@ -846,12 +570,10 @@ impl EtcdTaskStore { } let keepalive = self.start_keepalive(lease_id).await?; Ok(Some(CoordinationGuard { - backend: GuardBackend::Etcd { - token, - lease_id, - key, - keepalive, - }, + token, + lease_id, + key, + keepalive, })) } @@ -1076,27 +798,6 @@ enum DependencyStatus { Failed(String), } -fn dependency_status( - task: &TaskRecord, - mut lookup: impl FnMut(&str) -> Option, -) -> DependencyStatus { - let mut waiting = false; - for dependency in &task.depends_on { - match lookup(dependency) { - Some(TaskState::Done) => {} - Some(TaskState::Failed) | None => { - return DependencyStatus::Failed(dependency.clone()); - } - Some(TaskState::Queued | TaskState::Running) => waiting = true, - } - } - if waiting { - DependencyStatus::Waiting - } else { - DependencyStatus::Ready - } -} - fn fail_dependency(task: &mut TaskRecord, dependency: &str) { apply_outcome( task, @@ -1154,19 +855,6 @@ fn decode_task(value: &[u8], key: &str) -> lance::Result { .map_err(|err| lance::Error::io(format!("failed to decode task '{key}': {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 -} - fn etcd_error(action: &'static str) -> impl FnOnce(etcd_client::Error) -> lance::Error { move |err| lance::Error::io(format!("failed to {action}: {err}")) } @@ -1188,8 +876,6 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, - task_store_backend: TaskStoreBackend::Rocksdb, - task_db_path: dir.path().join("tasks").to_string_lossy().to_string(), etcd_endpoints: vec![], etcd_prefix: "/test".to_string(), etcd_username: None, @@ -1204,132 +890,13 @@ mod tests { } #[tokio::test] - async fn rocksdb_survives_reopen_and_dedupes() { - let dir = TempDir::new().unwrap(); - let cfg = config(&dir); - let first = TaskStore::open(&cfg).await.unwrap(); - let a = first - .enqueue(TaskKind::Compact, "experiment", Vec::new()) - .await - .unwrap(); - let b = first - .enqueue(TaskKind::Compact, "experiment", Vec::new()) - .await - .unwrap(); - assert_eq!(a.id, b.id); - drop(first); - - let reopened = TaskStore::open(&cfg).await.unwrap(); - assert_eq!(reopened.get(&a.id).await.unwrap().unwrap(), a); - } - - #[tokio::test] - async fn rocksdb_claims_and_finishes() { + async fn connect_requires_etcd_endpoints() { let dir = TempDir::new().unwrap(); - let store = TaskStore::open(&config(&dir)).await.unwrap(); - let task = store - .enqueue(TaskKind::IndexId, "experiment", Vec::new()) - .await - .unwrap(); - let claim = store.claim_next().await.unwrap().unwrap(); - assert_eq!(claim.task.id, task.id); - assert_eq!(claim.task.state, TaskState::Running); - store - .finish(claim, Ok("indexed".to_string())) - .await - .unwrap(); - let finished = store.get(&task.id).await.unwrap().unwrap(); - assert_eq!(finished.state, TaskState::Done); - assert_eq!(finished.detail.as_deref(), Some("indexed")); - } - - #[tokio::test] - async fn rocksdb_waits_for_dependencies_and_propagates_failure() { - let dir = TempDir::new().unwrap(); - let store = TaskStore::open(&config(&dir)).await.unwrap(); - - let prerequisite = store - .enqueue(TaskKind::MergeWal, "experiment", Vec::new()) - .await - .unwrap(); - let dependent = store - .enqueue( - TaskKind::Compact, - "experiment", - vec![prerequisite.id.clone()], - ) - .await - .unwrap(); - let claim = store.claim_next().await.unwrap().unwrap(); - assert_eq!(claim.task.id, prerequisite.id); - assert!( - store.claim_next().await.unwrap().is_none(), - "dependent task must wait while its prerequisite is running" - ); - store.finish(claim, Ok("merged".to_string())).await.unwrap(); - let dependent_claim = store.claim_next().await.unwrap().unwrap(); - assert_eq!(dependent_claim.task.id, dependent.id); - store - .finish(dependent_claim, Ok("compacted".to_string())) - .await - .unwrap(); - - let failed_prerequisite = store - .enqueue(TaskKind::MergeWal, "failed-chain", Vec::new()) - .await - .unwrap(); - let skipped = store - .enqueue( - TaskKind::IndexId, - "failed-chain", - vec![failed_prerequisite.id.clone()], - ) - .await - .unwrap(); - let claim = store.claim_next().await.unwrap().unwrap(); - assert_eq!(claim.task.id, failed_prerequisite.id); - store - .finish(claim, Err("merge failed".to_string())) - .await - .unwrap(); - assert!(store.claim_next().await.unwrap().is_none()); - let skipped = store.get(&skipped.id).await.unwrap().unwrap(); - assert_eq!(skipped.state, TaskState::Failed); - assert!(skipped - .error - .as_deref() - .is_some_and(|error| error.contains("dependency"))); - } - - #[tokio::test] - async fn dependent_tasks_do_not_participate_in_standalone_dedupe() { - let dir = TempDir::new().unwrap(); - let store = TaskStore::open(&config(&dir)).await.unwrap(); - let prerequisite = store - .enqueue(TaskKind::MergeWal, "experiment", Vec::new()) - .await - .unwrap(); - let dependent = store - .enqueue(TaskKind::Compact, "experiment", vec![prerequisite.id]) - .await - .unwrap(); - let standalone = store - .enqueue(TaskKind::Compact, "experiment", Vec::new()) - .await - .unwrap(); - assert_ne!(dependent.id, standalone.id); - } - - #[tokio::test] - async fn rejects_object_store_path_for_rocksdb() { - let dir = TempDir::new().unwrap(); - let mut cfg = config(&dir); - cfg.task_db_path = "s3://bucket/tasks".to_string(); - let error = match TaskStore::open(&cfg).await { - Ok(_) => panic!("object-store URI should be rejected"), + let error = match TaskStore::open(&config(&dir)).await { + Ok(_) => panic!("etcd backend must require endpoints"), Err(error) => error, }; - assert!(error.to_string().contains("must be local")); + assert!(error.to_string().contains("ETCD_ENDPOINTS is required")); } #[tokio::test] @@ -1339,7 +906,6 @@ mod tests { .expect("ETCD_TEST_ENDPOINTS must point to a test etcd"); let dir = TempDir::new().unwrap(); let mut cfg = config(&dir); - cfg.task_store_backend = TaskStoreBackend::Etcd; cfg.etcd_endpoints = endpoint.split(',').map(str::to_string).collect(); cfg.etcd_prefix = format!("/lance-context/test/{}", generate_id()); cfg.etcd_lease_ttl_secs = 5; diff --git a/deploy/kubernetes/README.md b/deploy/kubernetes/README.md index d9f350e..dc37c25 100644 --- a/deploy/kubernetes/README.md +++ b/deploy/kubernetes/README.md @@ -1,11 +1,7 @@ # Kubernetes deployment -Two scheduler persistence modes are supported: - -- `master.yaml`: one master with scheduler state in RocksDB on a dedicated - ReadWriteOnce PVC. -- `master-etcd.yaml`: three stateless master replicas sharing scheduler state - through a separate etcd cluster. +`master.yaml` runs several stateless master replicas that share scheduler state +through a separate etcd cluster. Before applying it: @@ -14,24 +10,16 @@ Before applying it: 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. For RocksDB, set `storageClassName` on the PVC when the cluster has no - default class. -5. For etcd, replace `ETCD_ENDPOINTS` and create the referenced auth/TLS - Secrets. Remove the auth variables or TLS volume/configuration when the - application etcd cluster does not use them. - -Do not apply both manifests at once. - -## RocksDB mode - -Keep the master at one replica. RocksDB is an embedded single-process database, -and its PVC must not be mounted by multiple active master pods. The `Recreate` -strategy ensures an old pod releases the volume before its replacement starts. +4. Replace `ETCD_ENDPOINTS` and create the referenced auth/TLS Secrets. Remove + the auth variables or TLS volume/configuration when the application etcd + cluster does not use them. -## etcd mode +## Scheduler state -Use a separate application etcd cluster, not the Kubernetes control-plane -etcd. Each etcd member owns its own persistent volume; master pods have no PVC. +Scheduler state lives entirely in etcd, so the master is stateless and runs with +multiple replicas. Use a separate application etcd cluster, not the Kubernetes +control-plane etcd. Each etcd member owns its own persistent volume; master pods +have no PVC. Task enqueue de-duplication, queued-to-running claims, and per-experiment Compact/IndexId locks use etcd transactions. Claims and locks are attached to diff --git a/deploy/kubernetes/master-etcd.yaml b/deploy/kubernetes/master-etcd.yaml deleted file mode 100644 index 4c1e52c..0000000 --- a/deploy/kubernetes/master-etcd.yaml +++ /dev/null @@ -1,97 +0,0 @@ -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: - replicas: 3 - selector: - matchLabels: - app: lance-context-master - template: - metadata: - labels: - app: lance-context-master - spec: - 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_STORE_BACKEND - value: etcd - # Point this at a separate application etcd cluster, not the - # Kubernetes control-plane etcd. - - name: ETCD_ENDPOINTS - value: https://etcd-0.etcd:2379,https://etcd-1.etcd:2379,https://etcd-2.etcd:2379 - - name: ETCD_PREFIX - value: /lance-context/master - - name: ETCD_LEASE_TTL_SECS - value: "30" - - name: ETCD_USERNAME - valueFrom: - secretKeyRef: - name: lance-context-etcd - key: username - optional: true - - name: ETCD_PASSWORD - valueFrom: - secretKeyRef: - name: lance-context-etcd - key: password - optional: true - - name: ETCD_CA_CERT - value: /etc/lance-context/etcd/ca.crt - - name: ETCD_CLIENT_CERT - value: /etc/lance-context/etcd/tls.crt - - name: ETCD_CLIENT_KEY - value: /etc/lance-context/etcd/tls.key - - name: TASK_HISTORY_LIMIT - value: "1000" - - name: MASTER_PORT - value: "8090" - - 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: etcd-tls - mountPath: /etc/lance-context/etcd - readOnly: true - volumes: - - name: etcd-tls - secret: - secretName: lance-context-etcd-client-tls diff --git a/deploy/kubernetes/master.yaml b/deploy/kubernetes/master.yaml index d3ea7eb..a88ba52 100644 --- a/deploy/kubernetes/master.yaml +++ b/deploy/kubernetes/master.yaml @@ -1,15 +1,4 @@ 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 @@ -26,10 +15,7 @@ kind: Deployment metadata: name: lance-context-master spec: - # RocksDB is single-process. Run one master against this PVC. - replicas: 1 - strategy: - type: Recreate + replicas: 3 selector: matchLabels: app: lance-context-master @@ -38,8 +24,6 @@ spec: labels: app: lance-context-master spec: - securityContext: - fsGroup: 1000 containers: - name: master image: lance-context-master:latest @@ -48,15 +32,36 @@ spec: # Replace with the same shared object-store prefix used by workers. - name: DATA_DIR value: s3://replace-me/rollouts - - name: TASK_STORE_BACKEND - value: rocksdb - - name: TASK_DB_PATH - value: /var/lib/lance-context-master/tasks.rocksdb + # Point this at a separate application etcd cluster, not the + # Kubernetes control-plane etcd. + - name: ETCD_ENDPOINTS + value: https://etcd-0.etcd:2379,https://etcd-1.etcd:2379,https://etcd-2.etcd:2379 + - name: ETCD_PREFIX + value: /lance-context/master + - name: ETCD_LEASE_TTL_SECS + value: "30" + - name: ETCD_USERNAME + valueFrom: + secretKeyRef: + name: lance-context-etcd + key: username + optional: true + - name: ETCD_PASSWORD + valueFrom: + secretKeyRef: + name: lance-context-etcd + key: password + optional: true + - name: ETCD_CA_CERT + value: /etc/lance-context/etcd/ca.crt + - name: ETCD_CLIENT_CERT + value: /etc/lance-context/etcd/tls.crt + - name: ETCD_CLIENT_KEY + value: /etc/lance-context/etcd/tls.key - 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: @@ -81,9 +86,10 @@ spec: limits: memory: 2Gi volumeMounts: - - name: task-db - mountPath: /var/lib/lance-context-master + - name: etcd-tls + mountPath: /etc/lance-context/etcd + readOnly: true volumes: - - name: task-db - persistentVolumeClaim: - claimName: lance-context-master-data + - name: etcd-tls + secret: + secretName: lance-context-etcd-client-tls