diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index a0c687e3..5bcd488e 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -31,7 +31,10 @@ pub(crate) async fn run_merge_wal( state: &Arc, claim: &TaskClaim, ) -> Result { - if state.config.worker_endpoints.is_empty() { + if state.config.worker_endpoints.is_empty() + && !(state.config.append.local(&claim.task.target) + && state.config.merge_rollout.owned(&claim.task.target)) + { return Err("no worker endpoints configured".into()); } let coordinator = state.task_store.merge_coordinator(); diff --git a/crates/lance-context-master/src/rollout_append.rs b/crates/lance-context-master/src/rollout_append.rs index ae765a09..cfaea9de 100644 --- a/crates/lance-context-master/src/rollout_append.rs +++ b/crates/lance-context-master/src/rollout_append.rs @@ -1,4 +1,4 @@ -//! Workers stage files concurrently; the master publishes bounded groups under +//! Workers or the resident master stage files concurrently; the master publishes bounded groups under //! the existing table claim, idle watchdog, and manifest-version fence. use crate::state::MasterState; use futures::{stream, FutureExt, StreamExt}; @@ -23,6 +23,17 @@ pub struct AppendConfig { pub rollout_append_max_generations: usize, #[arg(long, env = "ROLLOUT_APPEND_MAX_BYTES", default_value_t = 67_108_864)] pub rollout_append_max_bytes: usize, + /// Read and encode WAL in this master for opted-in owned append targets. + /// Uses the ordinary durable task queue; creates no Kubernetes Jobs. + #[arg(long, env = "ROLLOUT_APPEND_LOCAL", default_value_t = false)] + pub rollout_append_local: bool, + /// Shared Arrow-buffer budget across every local append task in this master. + #[arg( + long, + env = "ROLLOUT_APPEND_LOCAL_MEMORY_BYTES", + default_value_t = 2_147_483_648 + )] + pub rollout_append_local_memory_bytes: usize, } impl Default for AppendConfig { fn default() -> Self { @@ -31,10 +42,53 @@ impl Default for AppendConfig { rollout_append_concurrency: 4, rollout_append_max_generations: 64, rollout_append_max_bytes: 67_108_864, + rollout_append_local: false, + rollout_append_local_memory_bytes: 2_147_483_648, } } } impl AppendConfig { + pub(crate) fn local(&self, target: &str) -> bool { + self.rollout_append_local && self.enabled(target) + } + + pub(crate) fn validate_local(&self, config: &crate::config::MasterConfig) -> Result<()> { + if !self.rollout_append_local { + return Ok(()); + } + self.validate()?; + if config.catchup.enabled + || config.catchup.pod_template.is_some() + || config.catchup.target.is_some() + { + return Err( + "resident local append cannot be combined with Kubernetes catch-up mode".into(), + ); + } + if self.rollout_append_targets.is_empty() + || config.catchup.shards.is_empty() + || config.catchup.shards.len() > 256 + || config.catchup.shards.iter().any(|s| s.is_empty()) + || self.rollout_append_local_memory_bytes < self.rollout_append_max_bytes + || config.merge_wal_concurrency == 0 + || config.catchup.slice_secs == 0 + { + return Err("local append requires targets, stable WAL shards, a sufficient shared memory budget, and a separate merge task pool".into()); + } + for target in &self.rollout_append_targets { + if target != "*" + && (target.starts_with("generic:") + || !config.merge_rollout.owned(target) + || config.merge_rollout.draining(target)) + { + return Err(format!( + "local append requires a non-draining owned rollout: {target}" + )); + } + } + Ok(()) + } + pub fn enabled(&self, target: &str) -> bool { !target.starts_with("generic:") && self @@ -54,6 +108,25 @@ impl AppendConfig { } } +/// One instance per resident master, shared across tables and passes. +#[derive(Clone)] +pub(crate) struct LocalStaging { + pub(crate) budget: Arc, + session: Arc, +} + +impl LocalStaging { + pub(crate) fn new(memory_bytes: usize) -> Self { + Self { + budget: lance_context_core::MergeMemoryBudget::new(memory_bytes), + session: lance_context_core::RolloutStore::build_session( + 96 * 1024 * 1024, + 32 * 1024 * 1024, + ), + } + } +} + #[derive(Deserialize)] struct Capabilities { #[serde(default)] @@ -129,14 +202,14 @@ pub(crate) async fn run(state: &Arc, target: &str) -> Result = state .config .worker_endpoints .iter() - .filter(|_| !dedicated) + .filter(|_| !local_execution) .cloned() .map(|endpoint| { let client = state.http.clone(); @@ -191,9 +264,12 @@ pub(crate) async fn run(state: &Arc, target: &str) -> Result = plans.into_iter().enumerate().map(|(i, plan)| { let client = state.http.clone(); @@ -248,8 +326,8 @@ async fn run_pass( let uri = state.rollout_uri(&target); let local = local.clone(); async move { - let result = if let Some((budget, session)) = local { - lance_context_core::rollout_append::stage(&uri, plan.clone(), budget, Some(session)) + let result = if let Some(local) = local { + lance_context_core::rollout_append::stage(&uri, plan.clone(), local.budget, Some(local.session)) .await.map_err(|e| e.to_string()) } else { let (endpoint, fallback) = routes.expect("ordinary master requires staging endpoints"); @@ -297,15 +375,16 @@ async fn run_pass( if !group.is_empty() { reclaimed += coordinator.commit(group).await.map_err(|e| e.to_string())?; } - // Every speculative reader has now dropped its reservation. Retry an - // indivisible oversized generation once alone, without Job-level backoff. + // This pass's speculative readers have dropped their reservations. Retry an + // indivisible oversized generation serially within this pass. Other resident + // tasks still share the budget, so contention can remain retryable. for plan in memory_retries { - let (budget, session) = local.as_ref().unwrap(); + let local = local.as_ref().unwrap(); match lance_context_core::rollout_append::stage( &state.rollout_uri(target), plan, - budget.clone(), - Some(session.clone()), + local.budget.clone(), + Some(local.session.clone()), ) .await { @@ -320,7 +399,7 @@ async fn run_pass( } metrics::counter!("master_rollout_append_generations_reclaimed_total") .increment(reclaimed as u64); - tracing::info!(%target, reclaimed, failed = errors.len(), version = coordinator.version(), "parallel rollout append completed"); + tracing::info!(%target, reclaimed, local = local.is_some(), failed = errors.len(), version = coordinator.version(), "parallel rollout append completed"); if !errors.is_empty() { return Err(format!( "staged append reclaimed {reclaimed}; {} shard(s) failed: {}", @@ -350,6 +429,177 @@ mod tests { budget: Arc, } + #[test] + fn resident_merge_rejects_unowned_targets_and_unbounded_resources() { + let mut config = crate::config::MasterConfig::parse_from(["test"]); + config.append.rollout_append_local = true; + config.append.rollout_append_targets = vec!["hot".into()]; + config.catchup.shards = vec!["historical-worker".into()]; + assert!(config.append.validate_local(&config).is_err()); + config.merge_rollout.owned_targets = vec!["hot".into()]; + assert!(config.append.validate_local(&config).is_ok()); + config.merge_wal_concurrency = 0; + assert!(config.append.validate_local(&config).is_err()); + config.merge_wal_concurrency = 2; + config.append.rollout_append_local_memory_bytes = 1; + assert!(config.append.validate_local(&config).is_err()); + config.append.rollout_append_local_memory_bytes = 2 * 1024 * 1024 * 1024; + config.merge_rollout.drain_targets = vec!["hot".into()]; + assert!(config.append.validate_local(&config).is_err()); + config.merge_rollout.drain_targets.clear(); + config.catchup.enabled = true; + assert!(config.append.validate_local(&config).is_err()); + config.catchup.enabled = false; + config.catchup.slice_secs = 0; + assert!(config.append.validate_local(&config).is_err()); + } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_replicas_merge_from_queue_with_no_workers_or_jobs() { + use crate::task_store::TaskKinds; + use lance_context_api::TaskKind; + + let dir = tempfile::tempdir().unwrap(); + for target in ["hot", "cold"] { + let uri = dir.path().join(format!("{target}.rollout.lance")); + for shard in ["a", "b"] { + let writer = RolloutStore::open_with_options( + uri.to_str().unwrap(), + RolloutStoreOptions { + shard_id: Some(shard.into()), + merge_after_generations: Some(0), + ..Default::default() + }, + ) + .await + .unwrap(); + let dto = serde_json::from_value( + json!({"id":shard, "rollout_id":"r", "content":format!("{target}-{shard}")}), + ) + .unwrap(); + writer + .add(&[lance_context_core::rollout_record_from_add_request(&dto)]) + .await + .unwrap(); + writer.flush().await.unwrap(); + } + } + let mut config = crate::config::MasterConfig::parse_from([ + "test", + "--data-dir", + dir.path().to_str().unwrap(), + ]); + config.etcd.etcd_endpoints = std::env::var("ETCD_TEST_ENDPOINTS") + .unwrap() + .split(',') + .map(str::to_owned) + .collect(); + config.etcd.etcd_prefix = + format!("/resident-merge-test/{}", lance_context_core::generate_id()); + config.worker_endpoints.clear(); + config.merge_rollout.owned_targets = vec!["hot".into(), "cold".into()]; + config.append.rollout_append_targets = config.merge_rollout.owned_targets.clone(); + config.append.rollout_append_local = true; + config.append.rollout_append_max_bytes = 1024 * 1024; + config.append.rollout_append_local_memory_bytes = 2 * 1024 * 1024; + config.catchup.shards = vec!["a".into(), "b".into()]; + let first = MasterState::new(config.clone()).await.unwrap(); + let second = MasterState::new(config).await.unwrap(); + assert!(first.config.catchup.target.is_none()); + assert!(first.config.catchup.pod_template.is_none()); + assert!(!first.config.catchup.enabled); + let queued = first + .task_store + .enqueue(TaskKind::MergeWal, "hot", Vec::new()) + .await + .unwrap(); + let duplicate = second + .task_store + .enqueue(TaskKind::MergeWal, "hot", Vec::new()) + .await + .unwrap(); + assert_eq!(queued.id, duplicate.id); + let claim = first + .task_store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + assert!(second + .task_store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .is_none()); + + // Exhaust the first process's shared budget. Its local task must wait; + // another master can still service a different table from the same queue. + let budget = first.local_append.as_ref().unwrap().budget.clone(); + let held = budget.reserve(budget.limit()).await; + let runner = { + let first = first.clone(); + tokio::spawn(async move { + let result = crate::merge_execution::run_merge_wal(&first, &claim).await; + assert!(result.as_ref().unwrap().contains("2 generations")); + first.task_store.finish(claim, result).await.unwrap(); + }) + }; + second + .task_store + .enqueue(TaskKind::MergeWal, "cold", Vec::new()) + .await + .unwrap(); + let other = second + .task_store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + assert_eq!(other.task.target, "cold"); + let result = tokio::time::timeout( + Duration::from_secs(30), + crate::merge_execution::run_merge_wal(&second, &other), + ) + .await + .unwrap(); + assert!(result.as_ref().unwrap().contains("2 generations")); + second.task_store.finish(other, result).await.unwrap(); + assert!( + !runner.is_finished(), + "the first master must respect its shared memory budget" + ); + drop(held); + tokio::time::timeout(Duration::from_secs(30), runner) + .await + .unwrap() + .unwrap(); + assert_eq!(budget.reserved(), 0); + assert_eq!(second.local_append.as_ref().unwrap().budget.reserved(), 0); + for target in ["hot", "cold"] { + let uri = first.rollout_uri(target); + assert_eq!( + lance::Dataset::open(&uri) + .await + .unwrap() + .count_rows(None) + .await + .unwrap(), + 2 + ); + let store = + RolloutStore::open_existing_with_options(&uri, RolloutStoreOptions::default()) + .await + .unwrap(); + for id in ["a", "b"] { + assert_eq!( + store.get_by_id(id).await.unwrap().unwrap().content.unwrap(), + format!("{target}-{id}") + ); + } + } + } + async fn caps(Path(worker): Path) -> Json { Json( json!({"rollout_append_protocol": 1, "shard_name": worker, "owned_targets": ["hot"], "drain_targets": []}), diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 2f3fe7e8..5c8106d5 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -2025,6 +2025,97 @@ mod tests { worker.abort(); } + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn resident_merge_waiting_for_memory_keeps_general_tasks_running() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.append.rollout_append_local = true; + cfg.append.rollout_append_targets = vec!["exp".into()]; + cfg.catchup.shards = vec!["historical".into()]; + cfg.merge_wal_concurrency = 1; + cfg.task_concurrency = 1; + let state = MasterState::new(cfg).await.unwrap(); + let hot = RolloutStore::open_with_options( + &state.rollout_uri("exp"), + lance_context_core::RolloutStoreOptions { + shard_id: Some("historical".into()), + merge_after_generations: Some(0), + ..Default::default() + }, + ) + .await + .unwrap(); + hot.add(&[rollout_record("hot-row")]).await.unwrap(); + hot.flush().await.unwrap(); + let mut cold = RolloutStore::open(&state.rollout_uri("other")) + .await + .unwrap(); + cold.add(&[rollout_record("cold-row")]).await.unwrap(); + cold.cleanup_own_shard().await.unwrap(); + let budget = state.local_append.as_ref().unwrap().budget.clone(); + let held = budget.reserve(budget.limit()).await; + let merge = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let merge_poller = spawn_pool_poller( + state.clone(), + Arc::new(Semaphore::new(1)), + TaskKinds::MERGE_WAL, + false, + ); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + if state + .task_store + .get(&merge.id) + .await + .unwrap() + .unwrap() + .state + == TaskState::Running + { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .unwrap(); + let index = enqueue(&state, TaskKind::IndexId, "other").await.unwrap(); + let general = spawn_pool_poller( + state.clone(), + Arc::new(Semaphore::new(1)), + TaskKinds::GENERAL.without_compact(), + false, + ); + let indexed = await_terminal(&state, &index.id).await; + assert_eq!(indexed.state, TaskState::Done, "{indexed:?}"); + assert_eq!( + state + .task_store + .get(&merge.id) + .await + .unwrap() + .unwrap() + .state, + TaskState::Running + ); + drop(held); + let merged = await_terminal(&state, &merge.id).await; + assert_eq!(merged.state, TaskState::Done, "{merged:?}"); + assert_eq!(budget.reserved(), 0); + assert_eq!( + lance::Dataset::open(&state.rollout_uri("exp")) + .await + .unwrap() + .count_rows(None) + .await + .unwrap(), + 1 + ); + merge_poller.abort(); + general.abort(); + } + /// Enqueuing the same experiment twice while queued de-dupes to one task. #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index b13cee85..9779c912 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -107,6 +107,8 @@ pub struct MasterState { pub admission: Arc, /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. pub http: reqwest::Client, + /// Explicit resident-master staging shares one byte budget and cache session. + pub(crate) local_append: Option, /// Process-wide compaction permits shared by scheduler and retirement work. pub(crate) compaction_permits: Arc, /// Consecutive `_stats` maintenance failures, for alerting. @@ -136,6 +138,10 @@ impl MasterState { pub async fn new(config: MasterConfig) -> lance::Result> { config.merge_rollout.validate().map_err(lance::Error::io)?; config.catchup.validate().map_err(lance::Error::io)?; + config + .append + .validate_local(&config) + .map_err(lance::Error::io)?; if config.maintenance.maintenance_timeout_secs == 0 || config.maintenance.maintenance_idle_timeout_secs == 0 || config.maintenance.maintenance_drain_timeout_secs == 0 @@ -210,6 +216,11 @@ impl MasterState { if let Some(guard) = init_guard { task_store.release_coordination_lock(guard).await?; } + let local_append = config.append.rollout_append_local.then(|| { + crate::rollout_append::LocalStaging::new( + config.append.rollout_append_local_memory_bytes, + ) + }); let state = Arc::new(Self { registry, generic_registry, @@ -225,6 +236,7 @@ impl MasterState { task_store, admission: Arc::new(crate::admission::Admission::default()), http: worker_http_client().map_err(|error| lance::Error::io(error.to_string()))?, + local_append, compaction_permits: Arc::new(Semaphore::new(compaction_concurrency)), stats_maintenance_failures: std::sync::atomic::AtomicU64::new(0), stats_last_reclaimed_version: std::sync::atomic::AtomicU64::new(0), diff --git a/crates/lance-context-master/src/wal_tail.rs b/crates/lance-context-master/src/wal_tail.rs index 9d60f018..2cce2b5b 100644 --- a/crates/lance-context-master/src/wal_tail.rs +++ b/crates/lance-context-master/src/wal_tail.rs @@ -9,7 +9,7 @@ use std::{sync::Arc, time::Duration}; #[derive(Clone, Debug, clap::Args)] pub struct WalTailConfig { /// Revisit low-count WAL tails after this interval; 0 disables this sweep. - /// Requires ordinary automatic WAL merge and worker endpoints as well. + /// Requires automatic WAL merge and either workers or resident local append. #[arg(long, env = "MERGE_WAL_TAIL_INTERVAL_SECS", default_value_t = 0)] pub wal_tail_interval_secs: u64, /// At most this many candidates per fleet-wide 30-second batch (hard cap 64). @@ -38,7 +38,7 @@ struct Cursor { fn enabled(config: &MasterConfig) -> bool { config.wal_tail.wal_tail_interval_secs > 0 && config.merge_wal_interval_secs > 0 - && !config.worker_endpoints.is_empty() + && (!config.worker_endpoints.is_empty() || config.append.rollout_append_local) && config.wal_tail.wal_tail_stats_max_age_secs > 0 } @@ -57,6 +57,8 @@ fn candidates( && row.scanned_at <= now && now.saturating_sub(row.scanned_at) <= age && !config.merge_rollout.draining(&row.name) + && (!config.worker_endpoints.is_empty() + || (config.append.local(&row.name) && config.merge_rollout.owned(&row.name))) && after.is_none_or(|after| row.name.as_str() > after) }) .map(|row| row.name.clone()) @@ -246,6 +248,30 @@ mod tests { } } + #[test] + fn local_master_services_owned_tails_without_worker_endpoints() { + let mut cfg = config(); + cfg.worker_endpoints.clear(); + cfg.append.rollout_append_local = true; + cfg.append.rollout_append_targets = vec!["*".into()]; + cfg.merge_rollout.owned_targets = vec!["hot".into(), "generic:other".into()]; + let now = 2_000_000; + assert!(enabled(&cfg)); + assert_eq!( + candidates( + &cfg, + &[ + row("hot", 1, now), + row("legacy", 1, now), + row("generic:other", 1, now) + ], + None, + now + ), + ["hot"] + ); + } + #[test] fn only_fresh_low_count_tails_enter_the_bounded_rotation() { let mut cfg = config(); diff --git a/docs/rollout-parallel-append.md b/docs/rollout-parallel-append.md index 691735b2..d543692d 100644 --- a/docs/rollout-parallel-append.md +++ b/docs/rollout-parallel-append.md @@ -1,8 +1,9 @@ # Parallel rollout WAL append Rollout rows are immutable and their IDs are never reused for different records. -This opt-in path moves WAL reading, encoding and immutable file uploads to workers, -then publishes several workers' files in one base-table transaction on the master. +This opt-in path reads WAL, encodes and uploads immutable files concurrently, +then publishes the files together in one base-table transaction on the master. +Staging can run on workers or directly inside resident masters. Generic/context stores retain their existing keyed merge behavior. Staging requires Lance V2 files, whose dictionary encodings are local to each file. @@ -40,8 +41,53 @@ and perform parallel staging with their own CPU and shared catch-up memory budge they do not send their data work back to ordinary workers. A speculative read which cannot grow its reservation retries once alone after other readers finish. Both normal master merges and catch-up executions use the same publisher under -the existing maintenance ownership protocol. The ordinary master never reads WAL -payloads locally, including when no remote worker is available. +the existing maintenance ownership protocol. Remote mode never falls back to reading WAL payloads locally after a worker failure. +Resident local execution is an explicit configuration described below. + +## Resident master execution + +Set `ROLLOUT_APPEND_LOCAL=true` on regular master replicas to perform the same +parallel staging inside each resident process. Keep the existing shared etcd +prefix and storage root. No per-table executor, Kubernetes Job template, or Job +creation permission is needed. `CATCHUP_ENABLED`, `CATCHUP_POD_TEMPLATE` and +`CATCHUP_TARGET` cannot be combined with this mode. + +Configure: + +| Setting | Purpose | +| --- | --- | +| `ROLLOUT_APPEND_TARGETS` | Explicit owned, non-draining append-only rollout targets, or `*` for owned rollouts | +| `CATCHUP_SHARDS` | Stable writer identities, including historical workers; required even without live endpoints | +| `MERGE_WAL_CONCURRENCY` | Positive per-master table concurrency, independent of ordinary task slots | +| `ROLLOUT_APPEND_CONCURRENCY` | Parallel shard staging per active table, default 4 | +| `ROLLOUT_APPEND_LOCAL_MEMORY_BYTES` | Shared per-master Arrow reservation budget, default 2 GiB; at least the per-batch byte limit | +| `CATCHUP_SLICE_SECS` | Positive soft pass-admission duration; default 1800 s, at most 16 passes per task | + +Existing automatic pending sweeps, worker demand notifications and the low-count +WAL tail sweep feed the ordinary durable queue. Manual clients use the ordinary +MergeWal task API. Queue deduplication and leases distribute tables across master +replicas; the existing manifest fence protects publisher handoff. A resident +master does not claim the ingestion worker's live WAL epoch. + +The operator can increase the regular master replica count with the same settings. +This increases concurrent table capacity; one table still has one publisher and +uses parallel file staging internally. Queue leases preserve unfinished work when +a replica exits. Ordinary tasks keep their separate slots while merge waits on +its shared memory budget, although CPU and object-store bandwidth remain shared. + +The memory setting accounts for Arrow buffers, not total RSS. Lance encoding, +per-table coordinators, compaction and the rest of the process need headroom. One +indivisible oversized generation can exceed the nominal byte limit while holding +the entire reservation budget. Measure RSS, CPU, merge completion and pending +counts before increasing concurrency or replicas. Resident staging reuses one +96 MiB index / 32 MiB metadata cache session per process. + +Before moving an existing target, qualify the previous publisher's retirement or +complete the existing recovery fence. Persistent external-publisher guards remain +in force; enabling local execution does not erase them. All ingestion writers +must already honor owned-target fencing. To roll back local execution, drain the +resident work normally and disable `ROLLOUT_APPEND_LOCAL`; retain append/ownership +settings and compatible staging workers for remote execution. ## Commit and recovery