Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions crates/lance-context-api/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -980,13 +980,21 @@ pub struct TaskRecord {
/// When the task reached a terminal state, Unix milliseconds, if finished.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub finished_at: Option<i64>,
/// Task ids that must reach `Done` before this task may run. The scheduler
/// defers a task while any dependency is still `Queued`/`Running`, and marks
/// it `Failed` if any dependency `Failed`. Empty for standalone tasks.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub depends_on: Vec<String>,
}

/// Request body for `POST /api/v1/tasks`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnqueueTaskRequest {
pub kind: TaskKind,
pub target: String,
/// Optional task ids this task must wait for (see [`TaskRecord::depends_on`]).
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub depends_on: Vec<String>,
}

/// Response for `GET /api/v1/tasks`.
Expand Down Expand Up @@ -1096,6 +1104,7 @@ mod tests {
enqueued_at: 1_700_000_000_000,
started_at: None,
finished_at: None,
depends_on: Vec::new(),
};
let json = serde_json::to_string(&queued).unwrap();
// snake_case tags for the enums.
Expand All @@ -1105,6 +1114,7 @@ mod tests {
assert!(!json.contains("error"));
assert!(!json.contains("started_at"));
assert!(!json.contains("finished_at"));
assert!(!json.contains("depends_on"));
let back: TaskRecord = serde_json::from_str(&json).unwrap();
assert_eq!(back, queued);
}
Expand All @@ -1121,6 +1131,7 @@ mod tests {
enqueued_at: 1,
started_at: Some(2),
finished_at: Some(3),
depends_on: Vec::new(),
};
let back: TaskRecord =
serde_json::from_str(&serde_json::to_string(&done).unwrap()).unwrap();
Expand Down
2 changes: 1 addition & 1 deletion crates/lance-context-master/src/routes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,7 @@ pub async fn enqueue_task(
req.target
)));
}
let task = scheduler::enqueue(&state, req.kind, &req.target)
let task = scheduler::enqueue_with_deps(&state, req.kind, &req.target, req.depends_on)
.await
.map_err(MasterError::from_lance)?;
Ok((StatusCode::ACCEPTED, Json(task)))
Expand Down
156 changes: 153 additions & 3 deletions crates/lance-context-master/src/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,9 +68,22 @@ pub async fn enqueue(
state: &Arc<MasterState>,
kind: TaskKind,
target: &str,
) -> lance::Result<TaskRecord> {
enqueue_with_deps(state, kind, target, Vec::new()).await
}

/// Like [`enqueue`] but the task waits for `depends_on` (task ids) to reach
/// `Done` before it runs. De-dup is skipped when dependencies are present: a
/// dependent task is part of an ordered chain and must not collapse into an
/// unrelated in-flight task for the same target.
pub async fn enqueue_with_deps(
state: &Arc<MasterState>,
kind: TaskKind,
target: &str,
depends_on: Vec<String>,
) -> lance::Result<TaskRecord> {
let mut tasks = state.tasks.lock().await;
if matches!(kind, TaskKind::Compact | TaskKind::IndexId) {
if depends_on.is_empty() && matches!(kind, TaskKind::Compact | TaskKind::IndexId) {
if let Some(existing) = tasks.values().find(|t| {
t.kind == kind
&& t.target == target
Expand All @@ -90,6 +103,7 @@ pub async fn enqueue(
enqueued_at: now_ms(),
started_at: None,
finished_at: None,
depends_on,
};
state.task_store.put(&record)?;
tasks.insert(record.id.clone(), record.clone());
Expand Down Expand Up @@ -126,6 +140,38 @@ async fn update_task(
Ok(())
}

/// Readiness of a task with respect to its declared dependencies.
enum DepStatus {
/// All dependencies reached `Done` (or the task has none): ok to run.
Ready,
/// At least one dependency is still `Queued`/`Running`: defer and re-check.
Waiting,
/// At least one dependency `Failed` (or vanished): the dependent cannot run.
/// Carries the id of the offending dependency for the error message.
DepFailed(String),
}

/// Inspect a task's `depends_on` against the current task map. A missing
/// dependency is treated as failed (its record was never created or was
/// dropped), so a dependent never waits forever on a ghost id.
async fn dep_status(state: &Arc<MasterState>, id: &str) -> DepStatus {
let tasks = state.tasks.lock().await;
let Some(task) = tasks.get(id) else {
return DepStatus::Ready; // vanished; run_task will no-op
};
for dep in &task.depends_on {
match tasks.get(dep) {
Some(d) => match d.state {
TaskState::Done => {}
TaskState::Failed => return DepStatus::DepFailed(dep.clone()),
TaskState::Queued | TaskState::Running => return DepStatus::Waiting,
},
None => return DepStatus::DepFailed(dep.clone()),
}
}
DepStatus::Ready
}

/// Execute one task by id: dispatch on its kind, recording lifecycle timestamps
/// and the terminal state.
async fn run_task(state: &Arc<MasterState>, id: String) {
Expand Down Expand Up @@ -435,6 +481,36 @@ pub fn spawn_scheduler(state: &Arc<MasterState>) -> JoinHandle<()> {
let dispatch_state = state.clone();
tokio::spawn(async move {
while let Some(id) = rx.recv().await {
// Resolve dependencies before taking a slot, so a task waiting on a
// chain never occupies a permit that its own dependency needs to
// finish (which would deadlock at concurrency == 1).
match dep_status(&dispatch_state, &id).await {
DepStatus::Ready => {}
DepStatus::DepFailed(dep) => {
// Skip: mark the dependent Failed and move on.
update_task(&dispatch_state, &id, |t| {
t.state = TaskState::Failed;
t.finished_at = Some(now_ms());
t.error = Some(format!("dependency {dep} did not complete"));
})
.await
.unwrap_or_else(|error| {
tracing::error!(task = %id, error = %error, "failed to persist skipped task");
});
metrics::gauge!("master_task_queue_depth").decrement(1.0);
continue;
}
DepStatus::Waiting => {
// Re-check later without busy-looping. Re-send the id after a
// short delay; the queue is unbounded so this always succeeds.
let tx = dispatch_state.task_tx.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(200)).await;
let _ = tx.send(id);
});
continue;
}
}
// Acquire a slot before spawning so at most `concurrency` run at once.
let permit = sem
.clone()
Expand Down Expand Up @@ -645,8 +721,82 @@ mod tests {
worker.abort();
}

/// Minimal rollout record builder for tests (the core struct has no
/// `Default`).
/// A dependent task runs only after its dependency reaches `Done`: an
/// `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]
async fn dependent_task_runs_after_dependency_done() {
let dir = TempDir::new().unwrap();
let state = MasterState::new(config(&dir)).await.unwrap();
let worker = spawn_scheduler(&state);

let name = "exp";
let uri = state.rollout_uri(name);
{
let mut store = RolloutStore::open(&uri).await.unwrap();
for i in 0..4 {
let rec = rollout_record(&format!("r{i}"));
store.add(&[rec]).await.unwrap();
store.cleanup_own_shard().await.unwrap();
}
}
state
.registry
.write()
.await
.upsert(name, &uri)
.await
.unwrap();
crate::scanner::scan_once(&state).await.unwrap();

let compact = enqueue(&state, TaskKind::Compact, name).await.unwrap();
let index = enqueue_with_deps(&state, TaskKind::IndexId, name, vec![compact.id.clone()])
.await
.unwrap();

let compact_final = await_terminal(&state, &compact.id).await;
assert_eq!(
compact_final.state,
TaskState::Done,
"got {compact_final:?}"
);
let index_final = await_terminal(&state, &index.id).await;
assert_eq!(index_final.state, TaskState::Done, "got {index_final:?}");
// The dependent could not have started before the dependency finished.
assert!(
index_final.started_at.unwrap() >= compact_final.finished_at.unwrap(),
"index started {:?} before compact finished {:?}",
index_final.started_at,
compact_final.finished_at
);

worker.abort();
}

/// A dependent whose dependency `Failed` is skipped (marked `Failed`) rather
/// than run.
#[tokio::test]
async fn dependent_skipped_when_dependency_fails() {
let dir = TempDir::new().unwrap();
let state = MasterState::new(config(&dir)).await.unwrap();
let worker = spawn_scheduler(&state);

// MergeWal with no worker endpoints fails; its dependent must be skipped.
let merge = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap();
let dependent = enqueue_with_deps(&state, TaskKind::Compact, "exp", vec![merge.id.clone()])
.await
.unwrap();

let merge_final = await_terminal(&state, &merge.id).await;
assert_eq!(merge_final.state, TaskState::Failed);
let dep_final = await_terminal(&state, &dependent.id).await;
assert_eq!(dep_final.state, TaskState::Failed, "got {dep_final:?}");
assert!(dep_final.error.unwrap().contains("dependency"));

worker.abort();
}

/// Minimal rollout record builder for tests (the core struct has no /// `Default`).
pub fn rollout_record(id: &str) -> lance_context_core::RolloutRecord {
use chrono::TimeZone;
lance_context_core::RolloutRecord {
Expand Down
2 changes: 2 additions & 0 deletions crates/lance-context-master/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ mod tests {
enqueued_at: 1,
started_at: (task_state == TaskState::Running).then_some(2),
finished_at: (task_state == TaskState::Done).then_some(3),
depends_on: Vec::new(),
};
state.task_store.put(&task).unwrap();
}
Expand Down Expand Up @@ -265,6 +266,7 @@ mod tests {
started_at: None,
finished_at: matches!(state, TaskState::Done | TaskState::Failed)
.then_some(timestamp),
depends_on: Vec::new(),
};
store.put(&task).unwrap();
tasks.insert(id.to_string(), task);
Expand Down
1 change: 1 addition & 0 deletions crates/lance-context-master/src/task_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ mod tests {
enqueued_at,
started_at: None,
finished_at: None,
depends_on: Vec::new(),
}
}

Expand Down
42 changes: 42 additions & 0 deletions crates/lance-context-master/ui/package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/lance-context-master/ui/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
"@tanstack/react-query": "^5.59.0",
"react": "^18.3.1",
"react-dom": "^18.3.1",
"react-router-dom": "^6.30.4",
"zustand": "^4.5.5"
},
"devDependencies": {
Expand Down
Loading
Loading