chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat(workflow): add scheduled/delayed steps and conditional branching (#16980)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 15:06:12 -08:00
committed by GitHub
parent 6620548fd8
commit e0ac559ae1
3 changed files with 1025 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/workflow/definition.rs
+109 -5
View File
@@ -7,9 +7,11 @@ use std::{
time::Duration,
};
use chrono::{DateTime, Utc};
use super::{
executor::StepExecutor,
types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId},
types::{FailureAction, RetryPolicy, StepId, WorkflowContext, WorkflowData, WorkflowId},
};
/// Errors that can occur during workflow validation
@@ -32,7 +34,17 @@ pub struct StepDefinition<D: WorkflowData> {
pub retry_policy: Option<RetryPolicy>,
pub timeout: Option<Duration>,
pub on_failure: FailureAction,
/// Dependencies that must ALL complete before this step runs
pub depends_on: Vec<StepId>,
/// Dependencies where ANY completing triggers this step (used with depends_on)
pub depends_on_any: Vec<StepId>,
/// Delay before starting the step (after dependencies satisfied)
pub delay: Option<Duration>,
/// Run step at or after this time (after dependencies satisfied)
pub scheduled_at: Option<DateTime<Utc>>,
/// Condition to evaluate; if false, step is skipped
#[allow(clippy::type_complexity)]
pub run_if: Option<Arc<dyn Fn(&WorkflowContext<D>) -> bool + Send + Sync>>,
}
impl<D: WorkflowData> fmt::Debug for StepDefinition<D> {
@@ -44,6 +56,10 @@ impl<D: WorkflowData> fmt::Debug for StepDefinition<D> {
.field("timeout", &self.timeout)
.field("on_failure", &self.on_failure)
.field("depends_on", &self.depends_on)
.field("depends_on_any", &self.depends_on_any)
.field("delay", &self.delay)
.field("scheduled_at", &self.scheduled_at)
.field("run_if", &self.run_if.as_ref().map(|_| "<condition>"))
.finish_non_exhaustive()
}
}
@@ -62,6 +78,10 @@ impl<D: WorkflowData> StepDefinition<D> {
timeout: None,
on_failure: FailureAction::FailWorkflow,
depends_on: Vec::new(),
depends_on_any: Vec::new(),
delay: None,
scheduled_at: None,
run_if: None,
}
}
@@ -81,12 +101,74 @@ impl<D: WorkflowData> StepDefinition<D> {
}
/// Set dependencies for this step.
/// The step will only run after all specified dependencies have completed successfully.
/// The step will only run after ALL specified dependencies have completed successfully.
/// Empty slice means no dependencies - step can run immediately in parallel with others.
pub fn depends_on(mut self, deps: &[&str]) -> Self {
self.depends_on = deps.iter().map(|s| StepId::new(*s)).collect();
self
}
/// Set "any of" dependencies for this step.
/// The step will run when ANY of these dependencies complete (in addition to
/// all `depends_on` dependencies).
///
/// **Combined semantics** (when both `depends_on` and `depends_on_any` are set):
/// - ALL `depends_on` must complete successfully, AND
/// - AT LEAST ONE `depends_on_any` must complete successfully
///
/// **Failure handling**:
/// - For `depends_on`: if ANY fails, this step is blocked
/// - For `depends_on_any`: only blocked if ALL fail (since we only need one)
///
/// **Skipped dependencies**: A skipped dependency (e.g., via `run_if`) counts as
/// "completed" for dependency satisfaction purposes.
pub fn depends_on_any(mut self, deps: &[&str]) -> Self {
self.depends_on_any = deps.iter().map(|s| StepId::new(*s)).collect();
self
}
/// Set a delay before starting the step (after dependencies are satisfied).
///
/// If both `delay` and `scheduled_at` are set, the step will wait for the
/// scheduled time AND THEN wait for the delay duration (they stack).
pub fn with_delay(mut self, delay: Duration) -> Self {
self.delay = Some(delay);
self
}
/// Schedule the step to run at or after the specified time.
/// The step will wait until this time even if dependencies are satisfied earlier.
///
/// If the scheduled time is in the past when the step becomes ready, it will
/// proceed immediately (a debug log is emitted).
///
/// If both `delay` and `scheduled_at` are set, the step will wait for the
/// scheduled time AND THEN wait for the delay duration (they stack).
pub fn scheduled_at(mut self, time: DateTime<Utc>) -> Self {
self.scheduled_at = Some(time);
self
}
/// Set a condition for running this step.
/// If the condition returns false, the step is skipped.
///
/// **Skipped step semantics**: When a step is skipped:
/// - It is marked as `StepStatus::Skipped`
/// - Downstream steps that depend on it (via `depends_on`) will consider it satisfied
/// - This allows conditional branches without blocking the workflow
///
/// **Error handling**: If the context cannot be retrieved to evaluate the condition,
/// the step will **fail** (not proceed blindly). This is a safety measure.
///
/// **Note**: The condition closure is not serializable, so workflows with `run_if`
/// cannot be persisted and resumed from external storage.
pub fn run_if<F>(mut self, condition: F) -> Self
where
F: Fn(&WorkflowContext<D>) -> bool + Send + Sync + 'static,
{
self.run_if = Some(Arc::new(condition));
self
}
}
/// Complete workflow definition
@@ -166,7 +248,7 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
let steps_map: HashMap<&StepId, &StepDefinition<D>> =
self.steps.iter().map(|s| (&s.id, s)).collect();
// Check all dependencies exist
// Check all dependencies exist (both depends_on and depends_on_any)
for step in &self.steps {
for dep in &step.depends_on {
if !steps_map.contains_key(dep) {
@@ -176,9 +258,17 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
});
}
}
for dep in &step.depends_on_any {
if !steps_map.contains_key(dep) {
return Err(ValidationError::MissingDependency {
step: step.id.clone(),
dependency: dep.clone(),
});
}
}
}
// Check for cycles using DFS
// Check for cycles using DFS (considers both dependency types)
let mut visited = HashSet::new();
let mut rec_stack = HashSet::new();
@@ -191,6 +281,7 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
}
// Build reverse dependency map: for each step, which steps depend on it?
// Include both depends_on and depends_on_any
self.reverse_deps.clear();
for (idx, step) in self.steps.iter().enumerate() {
for dep_id in &step.depends_on {
@@ -199,14 +290,21 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
.or_default()
.push(idx);
}
for dep_id in &step.depends_on_any {
self.reverse_deps
.entry(dep_id.clone())
.or_default()
.push(idx);
}
}
// Cache indices of steps with no dependencies (can start immediately)
// A step with only depends_on_any still needs at least one to complete
self.initial_step_indices = self
.steps
.iter()
.enumerate()
.filter(|(_, s)| s.depends_on.is_empty())
.filter(|(_, s)| s.depends_on.is_empty() && s.depends_on_any.is_empty())
.map(|(i, _)| i)
.collect();
@@ -231,12 +329,18 @@ impl<D: WorkflowData> WorkflowDefinition<D> {
rec_stack.insert(step_id);
// O(1) lookup instead of linear search
// Check both depends_on and depends_on_any for cycles
if let Some(step) = steps_map.get(step_id) {
for dep in &step.depends_on {
if Self::has_cycle(dep, steps_map, visited, rec_stack) {
return true;
}
}
for dep in &step.depends_on_any {
if Self::has_cycle(dep, steps_map, visited, rec_stack) {
return true;
}
}
}
rec_stack.remove(step_id);
sgl-model-gateway/src/workflow/engine.rs
+278 -21
View File
@@ -35,6 +35,9 @@ struct StepTracker {
failed: HashSet<StepId>,
skipped: HashSet<StepId>,
running: HashSet<StepId>,
/// Steps waiting for delay/scheduled_at: maps step INDEX to ready time
/// Using index instead of StepId for O(w) iteration in the main loop
waiting_until: HashMap<usize, std::time::Instant>,
}
impl StepTracker {
@@ -42,22 +45,59 @@ impl StepTracker {
self.completed.len() + self.failed.len() + self.skipped.len()
}
fn is_step_processable(&self, step_id: &StepId) -> bool {
fn is_step_processable(&self, step_id: &StepId, step_idx: usize) -> bool {
!self.completed.contains(step_id)
&& !self.failed.contains(step_id)
&& !self.skipped.contains(step_id)
&& !self.running.contains(step_id)
&& !self.waiting_until.contains_key(&step_idx)
}
/// Get indices of waiting steps that are now ready to run - O(w) where w = waiting count
fn get_ready_waiting_indices(&self) -> Vec<usize> {
let now = std::time::Instant::now();
self.waiting_until
.iter()
.filter(|(_, &ready_at)| now >= ready_at)
.map(|(&idx, _)| idx)
.collect()
}
/// Mark a step as waiting until a specific time (by index)
fn set_waiting(&mut self, step_idx: usize, ready_at: std::time::Instant) {
self.waiting_until.insert(step_idx, ready_at);
}
/// Clear waiting status for a step (it's now ready to run)
fn clear_waiting(&mut self, step_idx: usize) {
self.waiting_until.remove(&step_idx);
}
/// Check if ALL dependencies are satisfied (completed or skipped)
fn are_dependencies_satisfied(&self, depends_on: &[StepId]) -> bool {
depends_on
.iter()
.all(|dep| self.completed.contains(dep) || self.skipped.contains(dep))
}
/// Check if ANY dependency is satisfied (completed or skipped)
/// Returns true if the list is empty (no "any" dependencies)
fn is_any_dependency_satisfied(&self, depends_on_any: &[StepId]) -> bool {
depends_on_any.is_empty()
|| depends_on_any
.iter()
.any(|dep| self.completed.contains(dep) || self.skipped.contains(dep))
}
/// Check if ANY dependency has failed (for depends_on - blocks if any fail)
fn has_failed_dependency(&self, depends_on: &[StepId]) -> bool {
depends_on.iter().any(|dep| self.failed.contains(dep))
}
/// Check if ALL dependencies have failed (for depends_on_any - blocks only if all fail)
fn have_all_any_deps_failed(&self, depends_on_any: &[StepId]) -> bool {
!depends_on_any.is_empty() && depends_on_any.iter().all(|dep| self.failed.contains(dep))
}
}
/// Fixed backoff that returns the same delay every time
@@ -264,6 +304,35 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
self.active_workflows.fetch_sub(1, Ordering::Release);
}
/// Calculate how long a step needs to wait based on delay and/or scheduled_at.
/// Returns None if no waiting is needed.
fn calculate_wait_duration(step: &StepDefinition<D>) -> Option<Duration> {
let now = Utc::now();
// Calculate wait time for scheduled_at
let schedule_wait = step.scheduled_at.and_then(|scheduled_time| {
if now < scheduled_time {
(scheduled_time - now).to_std().ok()
} else {
// Scheduled time is in the past
tracing::debug!(
step_id = %step.id,
scheduled_time = %scheduled_time,
"Step scheduled_at is in the past, proceeding immediately"
);
None
}
});
// Combine delay and schedule_wait (both apply if both are set)
match (step.delay, schedule_wait) {
(Some(delay), Some(schedule)) => Some(delay + schedule),
(Some(delay), None) => Some(delay),
(None, Some(schedule)) => Some(schedule),
(None, None) => None,
}
}
/// Create a guard that decrements active_workflows on drop.
/// This ensures the count is decremented even if a task panics.
fn active_workflow_guard(&self) -> ActiveWorkflowGuard {
@@ -428,22 +497,95 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
return Ok(());
}
// Find ready steps from pending_check (not all steps)
let (ready_step_indices, total_processed, running_count) = {
// Phase 0: Drain any pending completion signals to ensure dependents are added
// This prevents race conditions where tasks finish but their signals aren't processed
while let Ok((step_id, result)) = rx.try_recv() {
tracing::debug!(
step_id = %step_id,
result = ?result,
"Step completed (startup drain)"
);
if matches!(result, StepResult::Success | StepResult::Skip) {
for &dep_idx in definition.get_dependent_indices(&step_id) {
pending_check.push_back(dep_idx);
}
}
}
// Phase 1: Check waiting steps + deps-ready steps + blocked detection
// Single lock acquisition for all read operations
let (
newly_ready_from_wait,
deps_ready_indices,
total_processed,
current_running,
current_waiting,
) = {
let t = tracker.read();
// Only check steps in pending_check, not all steps
let ready: Vec<usize> = pending_check
// O(w) iteration over waiting_until keys instead of O(n) over all steps
let wait_ready: Vec<usize> = t.get_ready_waiting_indices();
// Check pending_check for dependency-satisfied steps
let deps_ready: Vec<usize> = pending_check
.drain(..)
.filter(|&idx| {
let step = &definition.steps[idx];
t.is_step_processable(&step.id)
t.is_step_processable(&step.id, idx)
&& t.are_dependencies_satisfied(&step.depends_on)
&& t.is_any_dependency_satisfied(&step.depends_on_any)
&& !t.has_failed_dependency(&step.depends_on)
&& !t.have_all_any_deps_failed(&step.depends_on_any)
})
.collect();
(ready, t.total_processed(), t.running.len())
(
wait_ready,
deps_ready,
t.total_processed(),
t.running.len(),
t.waiting_until.len(),
)
};
// Phase 2: Process waiting and deps-ready steps, update waiting_until
// Single write lock for all mutations
// Returns (ready_to_launch, steps_added_to_waiting)
let (ready_to_launch, steps_added_to_waiting) = {
let now = std::time::Instant::now();
let mut t = tracker.write();
let mut added_to_waiting = 0usize;
// Clear waiting status for ready steps
for &idx in &newly_ready_from_wait {
t.clear_waiting(idx);
}
// Collect steps ready to launch
let mut ready = newly_ready_from_wait;
for idx in deps_ready_indices {
let step = &definition.steps[idx];
let wait_duration = Self::calculate_wait_duration(step);
if let Some(duration) = wait_duration {
if duration > Duration::ZERO {
// Step needs to wait - add to waiting_until
let ready_at = now + duration;
tracing::debug!(
step_id = %step.id,
wait_ms = duration.as_millis(),
"Step waiting for delay/schedule"
);
t.set_waiting(idx, ready_at);
added_to_waiting += 1;
continue;
}
}
ready.push(idx);
}
(ready, added_to_waiting)
};
// Check if we're done
@@ -451,13 +593,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
break;
}
// Handle blocked workflow (no ready steps, none running, but work remains)
if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() {
// Handle blocked workflow (no ready steps, none running/waiting, but work remains)
// Use current_running/current_waiting from Phase 1, adjusted for Phase 2 changes.
// current_running may be stale-high (tasks completed since), but that's safe -
// we'd just do an extra loop iteration. current_waiting needs adjustment for
// steps we just added to waiting in Phase 2.
let effective_waiting = current_waiting + steps_added_to_waiting;
if ready_to_launch.is_empty()
&& current_running == 0
&& effective_waiting == 0
&& pending_check.is_empty()
{
let failed_step = tracker.read().failed.iter().next().cloned();
let error_message = if failed_step.is_some() {
"Workflow failed due to step dependency failure".to_string()
// Use &'static str to avoid allocation in common error paths
let error_message: &'static str = if failed_step.is_some() {
"Workflow failed due to step dependency failure"
} else {
"Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
"Workflow deadlocked: no steps ready and none running"
};
self.state_store
@@ -470,19 +622,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
instance_id,
failed_step: failed_step
.unwrap_or_else(|| StepId::new("internal_scheduler")),
error: error_message,
error: error_message.to_string(),
})
.await;
return Ok(());
}
// Launch ready steps in parallel
let mut tasks_launched = 0;
for step_idx in ready_step_indices {
let step = &definition.steps[step_idx];
tracker.write().running.insert(step.id.clone());
tasks_launched += 1;
let tasks_launched = ready_to_launch.len();
if tasks_launched > 0 {
let mut t = tracker.write();
for &idx in &ready_to_launch {
t.running.insert(definition.steps[idx].id.clone());
}
}
for step_idx in ready_to_launch {
let step = &definition.steps[step_idx];
let engine = self.clone_for_execution();
let def = Arc::clone(&definition);
let step_id = step.id.clone();
@@ -491,6 +647,76 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
tokio::spawn(async move {
let step = &def.steps[step_idx];
// Evaluate run_if condition if present
if let Some(ref condition) = step.run_if {
match engine.state_store.get_context(instance_id).await {
Ok(ctx) => {
if !condition(&ctx) {
// Condition is false - skip this step
tracing::debug!(
step_id = %step_id,
"Step skipped due to run_if condition"
);
// Update tracker and send skip signal
{
let mut t = tracker.write();
t.running.remove(&step_id);
t.skipped.insert(step_id.clone());
}
// Update state store
let _ = engine
.state_store
.update(instance_id, |s| {
if let Some(step_state) =
s.step_states.get_mut(&step_id)
{
step_state.status = StepStatus::Skipped;
}
})
.await;
// Send skip signal
let _ = tx.try_send((step_id, StepResult::Skip));
return;
}
}
Err(e) => {
// Failed to get context - fail the step rather than proceeding blindly
tracing::error!(
step_id = %step_id,
error = ?e,
"Failed to get context for run_if evaluation, failing step"
);
// Update tracker
{
let mut t = tracker.write();
t.running.remove(&step_id);
t.failed.insert(step_id.clone());
}
// Update state store
let _ = engine
.state_store
.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step_id) {
step_state.status = StepStatus::Failed;
step_state.last_error =
Some(format!("run_if context error: {e}"));
}
})
.await;
// Send failure signal
let _ = tx.try_send((step_id, StepResult::Failure));
return;
}
}
}
let result = engine
.execute_step_with_retry(instance_id, step, &def)
.await;
@@ -568,8 +794,17 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
});
}
let should_wait = tasks_launched > 0 || !tracker.read().running.is_empty();
if should_wait {
// Single lock read for both checks
let (has_running, has_waiting) = {
let t = tracker.read();
(
tasks_launched > 0 || !t.running.is_empty(),
!t.waiting_until.is_empty(),
)
};
if has_running {
// Wait for a step to complete; Phase 0 will drain any others
if let Some((completed_step_id, result)) = rx.recv().await {
tracing::debug!(
step_id = %completed_step_id,
@@ -585,6 +820,28 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}
}
}
} else if has_waiting {
// No running tasks, but some steps are waiting for their delay/schedule
// Calculate how long to sleep based on the soonest waiting step
let sleep_duration = {
let t = tracker.read();
let now = std::time::Instant::now();
t.waiting_until
.values()
.filter_map(|&ready_at| {
if ready_at > now {
Some(ready_at - now)
} else {
None
}
})
.min()
.unwrap_or(Duration::from_millis(10))
};
// Sleep until the next step is ready (with a cap to allow cancellation checks)
let capped_sleep = sleep_duration.min(Duration::from_millis(100));
tokio::time::sleep(capped_sleep).await;
}
}
@@ -603,7 +860,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.publish(WorkflowEvent::WorkflowFailed {
instance_id,
failed_step: step.clone(),
error: "One or more steps failed".to_string(),
error: "One or more steps failed".into(),
})
.await;
} else {