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

[model-gateway] make StateStore trait async for external persistence (#16979)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 12:55:59 -08:00
committed by GitHub
parent 6e158e55b4
commit 6620548fd8
4 changed files with 167 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

171
sgl-model-gateway/src/workflow/engine.rs
View File

@@ -229,7 +229,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
/// to ensure all workflows are stopped. Note that this cancels workflows
/// at the state level; running steps may still complete.
pub async fn force_cancel_all(&self) -> usize {
let active_states = match self.state_store.list_active() {
let active_states = match self.state_store.list_active().await {
Ok(states) => states,
Err(e) => {
tracing::error!(error = ?e, "Failed to list active workflows for force cancel");
@@ -304,7 +304,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
loop {
tokio::select! {
_ = ticker.tick() => {
state_store.cleanup_old_workflows(ttl);
state_store.cleanup_old_workflows(ttl).await;
}
_ = shutdown_rx.changed() => {
tracing::info!("Cleanup task stopping due to shutdown");
@@ -373,7 +373,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.insert(step.id.clone(), StepState::default());
}
self.state_store.save(state)?;
self.state_store.save(state).await?;
self.event_bus
.publish(WorkflowEvent::WorkflowStarted {
@@ -421,7 +421,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
.collect();
loop {
if self.state_store.is_cancelled(instance_id)? {
if self.state_store.is_cancelled(instance_id).await? {
self.event_bus
.publish(WorkflowEvent::WorkflowCancelled { instance_id })
.await;
@@ -460,9 +460,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
"Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
};
self.state_store.update(instance_id, |s| {
s.status = WorkflowStatus::Failed;
})?;
self.state_store
.update(instance_id, |s| {
s.status = WorkflowStatus::Failed;
})
.await?;
self.event_bus
.publish(WorkflowEvent::WorkflowFailed {
instance_id,
@@ -498,37 +500,31 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
Err(_) => StepResult::Failure,
};
{
// Track whether we need to update state to Skipped after releasing lock
let needs_skip_update = {
let mut t = tracker.write();
t.running.remove(&step_id);
match result {
let needs_update = match result {
Ok(StepResult::Success) => {
t.completed.insert(step_id.clone());
false
}
Ok(StepResult::Skip) => {
t.skipped.insert(step_id.clone());
false
}
Ok(StepResult::Failure) | Err(_) => match step.on_failure {
FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => {
t.failed.insert(step_id.clone());
false
}
FailureAction::ContinueNextStep => {
if let Err(e) = engine.state_store.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step_id) {
step_state.status = StepStatus::Skipped;
}
}) {
tracing::warn!(
step_id = %step_id,
error = ?e,
"Failed to update step state to Skipped"
);
}
t.skipped.insert(step_id.clone());
true // Need to update state store after releasing lock
}
},
}
};
if let Err(e) = tx.try_send((step_id.clone(), signal)) {
use mpsc::error::TrySendError;
@@ -547,6 +543,27 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}
}
}
needs_update
};
// Perform async state update after releasing the tracker lock
if needs_skip_update {
if let Err(e) = 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
{
tracing::warn!(
step_id = %step_id,
error = ?e,
"Failed to update step state to Skipped"
);
}
}
});
}
@@ -577,9 +594,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
};
if let Some(ref step) = failed_step {
self.state_store.update(instance_id, |s| {
s.status = WorkflowStatus::Failed;
})?;
self.state_store
.update(instance_id, |s| {
s.status = WorkflowStatus::Failed;
})
.await?;
self.event_bus
.publish(WorkflowEvent::WorkflowFailed {
instance_id,
@@ -588,9 +607,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
})
.await;
} else {
self.state_store.update(instance_id, |s| {
s.status = WorkflowStatus::Completed;
})?;
self.state_store
.update(instance_id, |s| {
s.status = WorkflowStatus::Completed;
})
.await?;
let duration = start_time.elapsed();
self.event_bus
@@ -624,23 +645,25 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let mut backoff = Self::create_backoff(&retry_policy.backoff);
loop {
if self.state_store.is_cancelled(instance_id)? {
if self.state_store.is_cancelled(instance_id).await? {
return Err(WorkflowError::Cancelled(instance_id));
}
// Update step state
self.state_store.update(instance_id, |s| {
s.current_step = Some(step.id.clone());
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = if attempt == 1 {
StepStatus::Running
} else {
StepStatus::Retrying
};
step_state.attempt = attempt;
step_state.started_at = Some(Utc::now());
}
})?;
self.state_store
.update(instance_id, |s| {
s.current_step = Some(step.id.clone());
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = if attempt == 1 {
StepStatus::Running
} else {
StepStatus::Retrying
};
step_state.attempt = attempt;
step_state.started_at = Some(Utc::now());
}
})
.await?;
// Emit step started event
self.event_bus
@@ -651,7 +674,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
})
.await;
let mut context = self.state_store.get_context(instance_id)?;
let mut context = self.state_store.get_context(instance_id).await?;
// Execute step with timeout
let step_start = std::time::Instant::now();
@@ -659,19 +682,23 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let step_duration = step_start.elapsed();
self.state_store.update(instance_id, |s| {
s.context = context.clone();
})?;
self.state_store
.update(instance_id, |s| {
s.context = context.clone();
})
.await?;
match result {
Ok(Ok(StepResult::Success)) => {
// Step succeeded
self.state_store.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = StepStatus::Succeeded;
step_state.completed_at = Some(Utc::now());
}
})?;
self.state_store
.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = StepStatus::Succeeded;
step_state.completed_at = Some(Utc::now());
}
})
.await?;
self.event_bus
.publish(WorkflowEvent::StepSucceeded {
@@ -708,19 +735,21 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let will_retry = should_retry && attempt < max_attempts;
// Update step state
self.state_store.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = if will_retry {
StepStatus::Retrying
} else {
StepStatus::Failed
};
step_state.last_error = Some(error_msg.clone());
if !will_retry {
step_state.completed_at = Some(Utc::now());
self.state_store
.update(instance_id, |s| {
if let Some(step_state) = s.step_states.get_mut(&step.id) {
step_state.status = if will_retry {
StepStatus::Retrying
} else {
StepStatus::Failed
};
step_state.last_error = Some(error_msg.clone());
if !will_retry {
step_state.completed_at = Some(Utc::now());
}
}
}
})?;
})
.await?;
// Emit step failed event
self.event_bus
@@ -787,9 +816,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
/// Cancel a running workflow
pub async fn cancel_workflow(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.state_store.update(instance_id, |s| {
s.status = WorkflowStatus::Cancelled;
})?;
self.state_store
.update(instance_id, |s| {
s.status = WorkflowStatus::Cancelled;
})
.await?;
self.event_bus
.publish(WorkflowEvent::WorkflowCancelled { instance_id })
@@ -799,8 +830,11 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}
/// Get workflow status
pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
self.state_store.load(instance_id)
pub async fn get_status(
&self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowState<D>> {
self.state_store.load(instance_id).await
}
/// Wait for a workflow to complete with adaptive polling
@@ -829,6 +863,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
let state = self
.get_status(instance_id)
.await
.map_err(|e| format!("Failed to get workflow status: {:?}", e))?;
let result = match state.status {
@@ -857,7 +892,7 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
}
};
self.state_store.cleanup_if_terminal(instance_id);
self.state_store.cleanup_if_terminal(instance_id).await;
return result;
}
}

95
sgl-model-gateway/src/workflow/state.rs
View File

@@ -2,6 +2,7 @@
use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration};
use async_trait::async_trait;
use parking_lot::RwLock;
use super::types::{
@@ -13,39 +14,43 @@ use super::types::{
///
/// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis).
/// The default implementation is `InMemoryStore` which keeps state in memory.
#[async_trait]
pub trait StateStore<D: WorkflowData>: Send + Sync + Clone {
/// Save workflow state
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>;
async fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>;
/// Load workflow state by instance ID
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>;
async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>;
/// Update workflow state using a closure
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
async fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where
F: FnOnce(&mut WorkflowState<D>);
F: FnOnce(&mut WorkflowState<D>) + Send;
/// Delete workflow state
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>;
async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>;
/// List all active workflows (Running or Pending)
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
async fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// List all workflows
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
async fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// Check if workflow is cancelled without loading full state
fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>;
async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>;
/// Clean up old completed/failed/cancelled workflows beyond a time threshold
fn cleanup_old_workflows(&self, ttl: Duration) -> usize;
async fn cleanup_old_workflows(&self, ttl: Duration) -> usize;
/// Get just the workflow context without cloning the entire state
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>>;
async fn get_context(
&self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowContext<D>>;
/// Clean up a specific workflow immediately if it's in a terminal state
/// Returns true if the workflow was removed, false otherwise
fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool;
async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool;
}
/// In-memory state storage for workflow instances
@@ -84,13 +89,14 @@ impl<D: WorkflowData> Default for InMemoryStore<D> {
}
}
#[async_trait]
impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> {
async fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> {
self.states.write().insert(state.instance_id, state);
Ok(())
}
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
self.states
.read()
.get(&instance_id)
@@ -98,28 +104,9 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
.ok_or(WorkflowError::NotFound(instance_id))
}
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states.write().remove(&instance_id);
Ok(())
}
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
async fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where
F: FnOnce(&mut WorkflowState<D>),
F: FnOnce(&mut WorkflowState<D>) + Send,
{
let mut states = self.states.write();
let state = states
@@ -130,15 +117,26 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
Ok(())
}
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>> {
self.states
.read()
.get(&instance_id)
.map(|s| s.context.clone())
.ok_or(WorkflowError::NotFound(instance_id))
async fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states.write().remove(&instance_id);
Ok(())
}
fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> {
async fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
async fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> {
self.states
.read()
.get(&instance_id)
@@ -146,7 +144,7 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
.ok_or(WorkflowError::NotFound(instance_id))
}
fn cleanup_old_workflows(&self, ttl: Duration) -> usize {
async fn cleanup_old_workflows(&self, ttl: Duration) -> usize {
let now = chrono::Utc::now();
let mut states = self.states.write();
let initial_count = states.len();
@@ -179,7 +177,18 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
removed_count
}
fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
async fn get_context(
&self,
instance_id: WorkflowInstanceId,
) -> WorkflowResult<WorkflowContext<D>> {
self.states
.read()
.get(&instance_id)
.map(|s| s.context.clone())
.ok_or(WorkflowError::NotFound(instance_id))
}
async fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
let mut states = self.states.write();
if let Some(state) = states.get(&instance_id) {
if matches!(

1
sgl-model-gateway/tests/wasm_test.rs
View File

@@ -710,6 +710,7 @@ async fn test_wasm_module_execution() {
let state = engines
.wasm_registration
.get_status(instance_id)
.await
.expect("Failed to get workflow status");
match state.status {

22
sgl-model-gateway/tests/workflow_test.rs
View File

@@ -104,7 +104,7 @@ async fn test_simple_workflow_execution() {
sleep(Duration::from_millis(100)).await;
// Check status
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
assert_eq!(state.step_states.len(), 2);
}
@@ -149,7 +149,7 @@ async fn test_workflow_with_retry() {
sleep(Duration::from_millis(500)).await;
// Check that step was retried and eventually succeeded
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
let step_state = state.step_states.get(&StepId::new("retry_step")).unwrap();
@@ -200,7 +200,7 @@ async fn test_workflow_failure_after_max_retries() {
sleep(Duration::from_millis(500)).await;
// Check that workflow failed
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Failed);
let step_state = state.step_states.get(&StepId::new("failing_step")).unwrap();
@@ -257,7 +257,7 @@ async fn test_workflow_continue_on_failure() {
sleep(Duration::from_millis(500)).await;
// Workflow should complete despite first step failing
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
// First step should be skipped
@@ -341,7 +341,7 @@ async fn test_workflow_context_sharing() {
sleep(Duration::from_millis(100)).await;
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
}
@@ -434,7 +434,7 @@ async fn test_parallel_execution_no_dependencies() {
// Wait for completion - give enough time for async scheduling
for _ in 0..50 {
sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running {
break;
}
@@ -442,7 +442,7 @@ async fn test_parallel_execution_no_dependencies() {
let overall_duration = overall_start.elapsed();
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
// Check that all steps completed
@@ -528,13 +528,13 @@ async fn test_dag_with_dependencies() {
// Poll until workflow completes (or timeout)
for _ in 0..50 {
sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running {
break;
}
}
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Completed);
// Verify step C started after both A and B finished
@@ -617,13 +617,13 @@ async fn test_dag_dependency_failure_blocks_dependents() {
// Poll until workflow completes (or timeout)
for _ in 0..50 {
sleep(Duration::from_millis(50)).await;
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
if state.status != WorkflowStatus::Running {
break;
}
}
let state = engine.get_status(instance_id).unwrap();
let state = engine.get_status(instance_id).await.unwrap();
assert_eq!(state.status, WorkflowStatus::Failed);
// Step B should not have executed