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

[model-gateway] refactor cleanup WorkflowContext.get_or_err (#14890)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-11 06:48:30 -08:00
committed by GitHub
parent ca1144212a
commit 60560c07a2
7 changed files with 76 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs
+13 -39
View File
@@ -218,9 +218,7 @@ pub struct DiscoverModelsStep;
#[async_trait] #[async_trait]
impl StepExecutor for DiscoverModelsStep { impl StepExecutor for DiscoverModelsStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
// If no API key is provided, skip model discovery and use wildcard mode. // If no API key is provided, skip model discovery and use wildcard mode.
if config.api_key.as_ref().is_none_or(|k| k.is_empty()) { if config.api_key.as_ref().is_none_or(|k| k.is_empty()) {
@@ -316,15 +314,9 @@ pub struct CreateExternalWorkersStep;
#[async_trait] #[async_trait]
impl StepExecutor for CreateExternalWorkersStep { impl StepExecutor for CreateExternalWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; let model_cards: Arc<Vec<ModelCard>> = context.get_or_err("model_cards")?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let model_cards: Arc<Vec<ModelCard>> = context
.get("model_cards")
.ok_or_else(|| WorkflowError::ContextValueNotFound("model_cards".to_string()))?;
// Build configs from router settings // Build configs from router settings
let circuit_breaker_config = { let circuit_breaker_config = {
@@ -459,15 +451,9 @@ pub struct RegisterExternalWorkersStep;
#[async_trait] #[async_trait]
impl StepExecutor for RegisterExternalWorkersStep { impl StepExecutor for RegisterExternalWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let mut worker_ids = Vec::new(); let mut worker_ids = Vec::new();
for worker in workers.iter() { for worker in workers.iter() {
@@ -496,18 +482,10 @@ pub struct UpdateExternalPoliciesStep;
#[async_trait] #[async_trait]
impl StepExecutor for UpdateExternalPoliciesStep { impl StepExecutor for UpdateExternalPoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let labels: Arc<HashMap<String, String>> = context let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
.get("labels")
.ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let policy_hint = labels.get("policy").map(|s| s.as_str()); let policy_hint = labels.get("policy").map(|s| s.as_str());
@@ -539,12 +517,8 @@ pub struct ActivateExternalWorkersStep;
#[async_trait] #[async_trait]
impl StepExecutor for ActivateExternalWorkersStep { impl StepExecutor for ActivateExternalWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
for worker in workers.iter() { for worker in workers.iter() {
worker.set_healthy(true); worker.set_healthy(true);
sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs
+9 -27
View File
@@ -46,12 +46,8 @@ pub struct ConnectMcpServerStep;
#[async_trait] #[async_trait]
impl StepExecutor for ConnectMcpServerStep { impl StepExecutor for ConnectMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
.get("mcp_server_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Connecting to MCP server: {}", config_request.name); debug!("Connecting to MCP server: {}", config_request.name);
@@ -102,15 +98,9 @@ impl StepExecutor for DiscoverMcpInventoryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
use rmcp::{service::RunningService, RoleClient}; use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
.get("mcp_server_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context
.get("mcp_client")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
debug!( debug!(
"Discovering inventory for MCP server: {}", "Discovering inventory for MCP server: {}",
@@ -153,15 +143,9 @@ impl StepExecutor for RegisterMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
use rmcp::{service::RunningService, RoleClient}; use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
.get("mcp_server_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context
.get("mcp_client")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
debug!("Registering MCP server: {}", config_request.name); debug!("Registering MCP server: {}", config_request.name);
@@ -199,9 +183,7 @@ pub struct ValidateRegistrationStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateRegistrationStep { impl StepExecutor for ValidateRegistrationStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
.get("mcp_server_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?;
let client_registered = context let client_registered = context
.get::<RunningService<RoleClient, ()>>("mcp_client") .get::<RunningService<RoleClient, ()>>("mcp_client")
sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs
+13 -39
View File
@@ -75,9 +75,7 @@ pub struct ValidateDescriptorStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateDescriptorStep { impl StepExecutor for ValidateDescriptorStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let descriptor = &config_request.descriptor; let descriptor = &config_request.descriptor;
@@ -235,9 +233,7 @@ pub struct CalculateHashStep;
#[async_trait] #[async_trait]
impl StepExecutor for CalculateHashStep { impl StepExecutor for CalculateHashStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let file_path = &config_request.descriptor.file_path; let file_path = &config_request.descriptor.file_path;
@@ -295,15 +291,9 @@ pub struct CheckDuplicateStep;
#[async_trait] #[async_trait]
impl StepExecutor for CheckDuplicateStep { impl StepExecutor for CheckDuplicateStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let sha256_hash: Arc<[u8; 32]> = context
.get("sha256_hash")
.ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
debug!( debug!(
"Checking for duplicate SHA256 hash for module: {}", "Checking for duplicate SHA256 hash for module: {}",
@@ -349,9 +339,7 @@ pub struct LoadWasmBytesStep;
#[async_trait] #[async_trait]
impl StepExecutor for LoadWasmBytesStep { impl StepExecutor for LoadWasmBytesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let file_path = &config_request.descriptor.file_path; let file_path = &config_request.descriptor.file_path;
@@ -386,12 +374,8 @@ pub struct ValidateWasmComponentStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateWasmComponentStep { impl StepExecutor for ValidateWasmComponentStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config") let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let wasm_bytes: Arc<Vec<u8>> = context
.get("wasm_bytes")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?;
debug!( debug!(
"Validating WASM component format for module: {}", "Validating WASM component format for module: {}",
@@ -441,21 +425,11 @@ pub struct RegisterModuleStep;
#[async_trait] #[async_trait]
impl StepExecutor for RegisterModuleStep { impl StepExecutor for RegisterModuleStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
.get("wasm_module_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
let app_context: Arc<AppContext> = context let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?;
.get("app_context") let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let sha256_hash: Arc<[u8; 32]> = context
.get("sha256_hash")
.ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
let file_size_bytes: Arc<u64> = context
.get("file_size_bytes")
.ok_or_else(|| WorkflowError::ContextValueNotFound("file_size_bytes".to_string()))?;
let wasm_bytes: Arc<Vec<u8>> = context
.get("wasm_bytes")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?;
debug!( debug!(
"Registering WASM module in manager: {}", "Registering WASM module in manager: {}",
sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs
+4 -12
View File
@@ -41,12 +41,8 @@ pub struct FindModuleToRemoveStep;
impl StepExecutor for FindModuleToRemoveStep { impl StepExecutor for FindModuleToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> = let removal_request: Arc<WasmModuleRemovalRequest> =
context.get("wasm_module_removal_request").ok_or_else(|| { context.get_or_err("wasm_module_removal_request")?;
WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string()) let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
})?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Finding module to remove: {}", removal_request.module_uuid); debug!("Finding module to remove: {}", removal_request.module_uuid);
@@ -93,12 +89,8 @@ pub struct RemoveModuleStep;
impl StepExecutor for RemoveModuleStep { impl StepExecutor for RemoveModuleStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> = let removal_request: Arc<WasmModuleRemovalRequest> =
context.get("wasm_module_removal_request").ok_or_else(|| { context.get_or_err("wasm_module_removal_request")?;
WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string()) let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
})?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Removing WASM module: {}", removal_request.module_uuid); debug!("Removing WASM module: {}", removal_request.module_uuid);
sgl-model-gateway/src/core/workflow/steps/worker_registration.rs
+22 -66
View File
@@ -288,12 +288,8 @@ pub struct DetectConnectionModeStep;
#[async_trait] #[async_trait]
impl StepExecutor for DetectConnectionModeStep { impl StepExecutor for DetectConnectionModeStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!( debug!(
"Detecting connection mode for {} (timeout: {}s, max_attempts: {})", "Detecting connection mode for {} (timeout: {}s, max_attempts: {})",
@@ -346,12 +342,8 @@ pub struct DiscoverMetadataStep;
#[async_trait] #[async_trait]
impl StepExecutor for DiscoverMetadataStep { impl StepExecutor for DiscoverMetadataStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let connection_mode: Arc<ConnectionMode> = context
.get("connection_mode")
.ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
debug!( debug!(
"Discovering metadata for {} ({:?})", "Discovering metadata for {} ({:?})",
@@ -430,9 +422,7 @@ pub struct DiscoverDPInfoStep;
#[async_trait] #[async_trait]
impl StepExecutor for DiscoverDPInfoStep { impl StepExecutor for DiscoverDPInfoStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
if !config.dp_aware { if !config.dp_aware {
debug!( debug!(
@@ -471,18 +461,10 @@ pub struct CreateWorkerStep;
#[async_trait] #[async_trait]
impl StepExecutor for CreateWorkerStep { impl StepExecutor for CreateWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let app_context: Arc<AppContext> = context let discovered_labels: Arc<HashMap<String, String>> = context.get_or_err("discovered_labels")?;
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let connection_mode: Arc<ConnectionMode> = context
.get("connection_mode")
.ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
let discovered_labels: Arc<HashMap<String, String>> = context
.get("discovered_labels")
.ok_or_else(|| WorkflowError::ContextValueNotFound("discovered_labels".to_string()))?;
// Check if worker already exists // Check if worker already exists
if app_context if app_context
@@ -638,9 +620,7 @@ impl StepExecutor for CreateWorkerStep {
// Handle DP-aware vs non-DP-aware workers // Handle DP-aware vs non-DP-aware workers
if config.dp_aware { if config.dp_aware {
let dp_info: Arc<DpInfo> = context let dp_info: Arc<DpInfo> = context.get_or_err("dp_info")?;
.get("dp_info")
.ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?;
debug!( debug!(
"Creating {} DP-aware workers for {} (dp_size: {})", "Creating {} DP-aware workers for {} (dp_size: {})",
@@ -724,17 +704,11 @@ pub struct RegisterWorkerStep;
#[async_trait] #[async_trait]
impl StepExecutor for RegisterWorkerStep { impl StepExecutor for RegisterWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
if config.dp_aware { if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let mut worker_ids = Vec::new(); let mut worker_ids = Vec::new();
for worker in workers.iter() { for worker in workers.iter() {
@@ -748,9 +722,7 @@ impl StepExecutor for RegisterWorkerStep {
context.set("worker_ids", worker_ids); context.set("worker_ids", worker_ids);
} else { } else {
let worker: Arc<Arc<dyn Worker>> = context let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
let worker_id = app_context let worker_id = app_context
.worker_registry .worker_registry
@@ -773,22 +745,14 @@ pub struct UpdatePoliciesStep;
#[async_trait] #[async_trait]
impl StepExecutor for UpdatePoliciesStep { impl StepExecutor for UpdatePoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config") let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let labels: Arc<HashMap<String, String>> = context
.get("labels")
.ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let policy_hint = labels.get("policy").map(|s| s.as_str()); let policy_hint = labels.get("policy").map(|s| s.as_str());
if config.dp_aware { if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let model_id = workers[0].model_id().to_string(); let model_id = workers[0].model_id().to_string();
@@ -815,9 +779,7 @@ impl StepExecutor for UpdatePoliciesStep {
model_id model_id
); );
} else { } else {
let worker: Arc<Arc<dyn Worker>> = context let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
let model_id = worker.model_id().to_string(); let model_id = worker.model_id().to_string();
@@ -864,14 +826,10 @@ pub struct ActivateWorkerStep;
#[async_trait] #[async_trait]
impl StepExecutor for ActivateWorkerStep { impl StepExecutor for ActivateWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
if config.dp_aware { if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
for worker in workers.iter() { for worker in workers.iter() {
worker.set_healthy(true); worker.set_healthy(true);
@@ -883,9 +841,7 @@ impl StepExecutor for ActivateWorkerStep {
config.url config.url
); );
} else { } else {
let worker: Arc<Arc<dyn Worker>> = context let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
worker.set_healthy(true); worker.set_healthy(true);
sgl-model-gateway/src/core/workflow/steps/worker_removal.rs
+9 -27
View File
@@ -32,12 +32,8 @@ pub struct FindWorkersToRemoveStep;
#[async_trait] #[async_trait]
impl StepExecutor for FindWorkersToRemoveStep { impl StepExecutor for FindWorkersToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerRemovalRequest> = context let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
.get("removal_request") let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("removal_request".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!( debug!(
"Finding workers to remove for {} (dp_aware: {})", "Finding workers to remove for {} (dp_aware: {})",
@@ -109,12 +105,8 @@ pub struct RemoveFromPolicyRegistryStep;
#[async_trait] #[async_trait]
impl StepExecutor for RemoveFromPolicyRegistryStep { impl StepExecutor for RemoveFromPolicyRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.get("app_context") let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers_to_remove")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers_to_remove")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_remove".to_string()))?;
debug!( debug!(
"Removing {} worker(s) from policy registry", "Removing {} worker(s) from policy registry",
@@ -153,12 +145,8 @@ pub struct RemoveFromWorkerRegistryStep;
#[async_trait] #[async_trait]
impl StepExecutor for RemoveFromWorkerRegistryStep { impl StepExecutor for RemoveFromWorkerRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.get("app_context") let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let worker_urls: Arc<Vec<String>> = context
.get("worker_urls")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_urls".to_string()))?;
debug!( debug!(
"Removing {} worker(s) from worker registry", "Removing {} worker(s) from worker registry",
@@ -202,15 +190,9 @@ pub struct UpdateRemainingPoliciesStep;
#[async_trait] #[async_trait]
impl StepExecutor for UpdateRemainingPoliciesStep { impl StepExecutor for UpdateRemainingPoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
.get("app_context") let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
let affected_models: Arc<HashSet<String>> = context
.get("affected_models")
.ok_or_else(|| WorkflowError::ContextValueNotFound("affected_models".to_string()))?;
let worker_urls: Arc<Vec<String>> = context
.get("worker_urls")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_urls".to_string()))?;
debug!( debug!(
"Updating cache-aware policies for {} affected model(s)", "Updating cache-aware policies for {} affected model(s)",
sgl-model-gateway/src/core/workflow/types.rs
+6
View File
@@ -218,6 +218,12 @@ impl WorkflowContext {
.and_then(|v| v.clone().downcast::<T>().ok()) .and_then(|v| v.clone().downcast::<T>().ok())
} }
/// Retrieve a value from the context, returning an error if not found
pub fn get_or_err<T: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<T>, WorkflowError> {
self.get(key)
.ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string()))
}
/// Check if the context has any data that would be lost during serialization /// Check if the context has any data that would be lost during serialization
pub fn has_unserializable_data(&self) -> bool { pub fn has_unserializable_data(&self) -> bool {
!self.data.is_empty() !self.data.is_empty()