diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index 48281ac0d..259615a10 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/sgl-model-gateway/src/workflow/engine.rs @@ -229,7 +229,7 @@ impl + 'static> WorkflowEngine { /// 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 + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { .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 + 'static> WorkflowEngine { .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 + 'static> WorkflowEngine { "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 + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { } } } + + 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 + 'static> WorkflowEngine { }; 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 + 'static> WorkflowEngine { }) .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 + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { }) .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 + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { /// 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 + 'static> WorkflowEngine { } /// Get workflow status - pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { - self.state_store.load(instance_id) + pub async fn get_status( + &self, + instance_id: WorkflowInstanceId, + ) -> WorkflowResult> { + self.state_store.load(instance_id).await } /// Wait for a workflow to complete with adaptive polling @@ -829,6 +863,7 @@ impl + 'static> WorkflowEngine { 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 + 'static> WorkflowEngine { } }; - self.state_store.cleanup_if_terminal(instance_id); + self.state_store.cleanup_if_terminal(instance_id).await; return result; } } diff --git a/sgl-model-gateway/src/workflow/state.rs b/sgl-model-gateway/src/workflow/state.rs index 5a2d898a0..127159949 100644 --- a/sgl-model-gateway/src/workflow/state.rs +++ b/sgl-model-gateway/src/workflow/state.rs @@ -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: Send + Sync + Clone { /// Save workflow state - fn save(&self, state: WorkflowState) -> WorkflowResult<()>; + async fn save(&self, state: WorkflowState) -> WorkflowResult<()>; /// Load workflow state by instance ID - fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult>; + async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult>; /// Update workflow state using a closure - fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> + async fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> where - F: FnOnce(&mut WorkflowState); + F: FnOnce(&mut WorkflowState) + 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>>; + async fn list_active(&self) -> WorkflowResult>>; /// List all workflows - fn list_all(&self) -> WorkflowResult>>; + async fn list_all(&self) -> WorkflowResult>>; /// Check if workflow is cancelled without loading full state - fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult; + async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult; /// 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>; + async fn get_context( + &self, + instance_id: WorkflowInstanceId, + ) -> WorkflowResult>; /// 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 Default for InMemoryStore { } } +#[async_trait] impl StateStore for InMemoryStore { - fn save(&self, state: WorkflowState) -> WorkflowResult<()> { + async fn save(&self, state: WorkflowState) -> WorkflowResult<()> { self.states.write().insert(state.instance_id, state); Ok(()) } - fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { + async fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { self.states .read() .get(&instance_id) @@ -98,28 +104,9 @@ impl StateStore for InMemoryStore { .ok_or(WorkflowError::NotFound(instance_id)) } - fn list_active(&self) -> WorkflowResult>> { - let states = self.states.read(); - Ok(states - .values() - .filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending)) - .cloned() - .collect()) - } - - fn list_all(&self) -> WorkflowResult>> { - 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(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> + async fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> where - F: FnOnce(&mut WorkflowState), + F: FnOnce(&mut WorkflowState) + Send, { let mut states = self.states.write(); let state = states @@ -130,15 +117,26 @@ impl StateStore for InMemoryStore { Ok(()) } - fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { - 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 { + async fn list_active(&self) -> WorkflowResult>> { + 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>> { + let states = self.states.read(); + Ok(states.values().cloned().collect()) + } + + async fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult { self.states .read() .get(&instance_id) @@ -146,7 +144,7 @@ impl StateStore for InMemoryStore { .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 StateStore for InMemoryStore { removed_count } - fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool { + async fn get_context( + &self, + instance_id: WorkflowInstanceId, + ) -> WorkflowResult> { + 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!( diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index 09e1d6fda..825b55c87 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -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 { diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs index c77c2cc5e..c1b6b08ed 100644 --- a/sgl-model-gateway/tests/workflow_test.rs +++ b/sgl-model-gateway/tests/workflow_test.rs @@ -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