From 60560c07a2c2a9fe424b316051a7f9054f9fc8ed Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 11 Dec 2025 22:48:30 +0800 Subject: [PATCH] [model-gateway] refactor cleanup WorkflowContext.get_or_err (#14890) --- .../steps/external_worker_registration.rs | 52 +++-------- .../core/workflow/steps/mcp_registration.rs | 36 ++------ .../steps/wasm_module_registration.rs | 52 +++-------- .../workflow/steps/wasm_module_removal.rs | 16 +--- .../workflow/steps/worker_registration.rs | 88 +++++-------------- .../src/core/workflow/steps/worker_removal.rs | 36 ++------ sgl-model-gateway/src/core/workflow/types.rs | 6 ++ 7 files changed, 76 insertions(+), 210 deletions(-) diff --git a/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs b/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs index 4ebbefc4d..cd54067dd 100644 --- a/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs @@ -218,9 +218,7 @@ pub struct DiscoverModelsStep; #[async_trait] impl StepExecutor for DiscoverModelsStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; // 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()) { @@ -316,15 +314,9 @@ pub struct CreateExternalWorkersStep; #[async_trait] impl StepExecutor for CreateExternalWorkersStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let model_cards: Arc> = context - .get("model_cards") - .ok_or_else(|| WorkflowError::ContextValueNotFound("model_cards".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let model_cards: Arc> = context.get_or_err("model_cards")?; // Build configs from router settings let circuit_breaker_config = { @@ -459,15 +451,9 @@ pub struct RegisterExternalWorkersStep; #[async_trait] impl StepExecutor for RegisterExternalWorkersStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let workers: Arc>> = context.get_or_err("workers")?; let mut worker_ids = Vec::new(); for worker in workers.iter() { @@ -496,18 +482,10 @@ pub struct UpdateExternalPoliciesStep; #[async_trait] impl StepExecutor for UpdateExternalPoliciesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let labels: Arc> = context - .get("labels") - .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let labels: Arc> = context.get_or_err("labels")?; + let app_context: Arc = context.get_or_err("app_context")?; + let workers: Arc>> = context.get_or_err("workers")?; let policy_hint = labels.get("policy").map(|s| s.as_str()); @@ -539,12 +517,8 @@ pub struct ActivateExternalWorkersStep; #[async_trait] impl StepExecutor for ActivateExternalWorkersStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let workers: Arc>> = context.get_or_err("workers")?; for worker in workers.iter() { worker.set_healthy(true); diff --git a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs b/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs index 94fbdd100..7952e37c2 100644 --- a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs @@ -46,12 +46,8 @@ pub struct ConnectMcpServerStep; #[async_trait] impl StepExecutor for ConnectMcpServerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("mcp_server_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let config_request: Arc = context.get_or_err("mcp_server_config")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!("Connecting to MCP server: {}", config_request.name); @@ -102,15 +98,9 @@ impl StepExecutor for DiscoverMcpInventoryStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { use rmcp::{service::RunningService, RoleClient}; - let config_request: Arc = context - .get("mcp_server_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let mcp_client: Arc> = context - .get("mcp_client") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?; + let config_request: Arc = context.get_or_err("mcp_server_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let mcp_client: Arc> = context.get_or_err("mcp_client")?; debug!( "Discovering inventory for MCP server: {}", @@ -153,15 +143,9 @@ impl StepExecutor for RegisterMcpServerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { use rmcp::{service::RunningService, RoleClient}; - let config_request: Arc = context - .get("mcp_server_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let mcp_client: Arc> = context - .get("mcp_client") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?; + let config_request: Arc = context.get_or_err("mcp_server_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let mcp_client: Arc> = context.get_or_err("mcp_client")?; debug!("Registering MCP server: {}", config_request.name); @@ -199,9 +183,7 @@ pub struct ValidateRegistrationStep; #[async_trait] impl StepExecutor for ValidateRegistrationStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("mcp_server_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?; + let config_request: Arc = context.get_or_err("mcp_server_config")?; let client_registered = context .get::>("mcp_client") diff --git a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs index 8da84b31e..2ee3082c2 100644 --- a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs @@ -75,9 +75,7 @@ pub struct ValidateDescriptorStep; #[async_trait] impl StepExecutor for ValidateDescriptorStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; let descriptor = &config_request.descriptor; @@ -235,9 +233,7 @@ pub struct CalculateHashStep; #[async_trait] impl StepExecutor for CalculateHashStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; let file_path = &config_request.descriptor.file_path; @@ -295,15 +291,9 @@ pub struct CheckDuplicateStep; #[async_trait] impl StepExecutor for CheckDuplicateStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; - let app_context: Arc = 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()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; debug!( "Checking for duplicate SHA256 hash for module: {}", @@ -349,9 +339,7 @@ pub struct LoadWasmBytesStep; #[async_trait] impl StepExecutor for LoadWasmBytesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; let file_path = &config_request.descriptor.file_path; @@ -386,12 +374,8 @@ pub struct ValidateWasmComponentStep; #[async_trait] impl StepExecutor for ValidateWasmComponentStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; - let wasm_bytes: Arc> = context - .get("wasm_bytes") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; + let wasm_bytes: Arc> = context.get_or_err("wasm_bytes")?; debug!( "Validating WASM component format for module: {}", @@ -441,21 +425,11 @@ pub struct RegisterModuleStep; #[async_trait] impl StepExecutor for RegisterModuleStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context - .get("wasm_module_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?; - let app_context: Arc = 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()))?; - let file_size_bytes: Arc = context - .get("file_size_bytes") - .ok_or_else(|| WorkflowError::ContextValueNotFound("file_size_bytes".to_string()))?; - let wasm_bytes: Arc> = context - .get("wasm_bytes") - .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?; + let config_request: Arc = context.get_or_err("wasm_module_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; + let file_size_bytes: Arc = context.get_or_err("file_size_bytes")?; + let wasm_bytes: Arc> = context.get_or_err("wasm_bytes")?; debug!( "Registering WASM module in manager: {}", diff --git a/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs b/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs index 9d2fafc27..9cd20c9d2 100644 --- a/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs +++ b/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs @@ -41,12 +41,8 @@ pub struct FindModuleToRemoveStep; impl StepExecutor for FindModuleToRemoveStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { let removal_request: Arc = - context.get("wasm_module_removal_request").ok_or_else(|| { - WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string()) - })?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + context.get_or_err("wasm_module_removal_request")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!("Finding module to remove: {}", removal_request.module_uuid); @@ -93,12 +89,8 @@ pub struct RemoveModuleStep; impl StepExecutor for RemoveModuleStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { let removal_request: Arc = - context.get("wasm_module_removal_request").ok_or_else(|| { - WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string()) - })?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + context.get_or_err("wasm_module_removal_request")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!("Removing WASM module: {}", removal_request.module_uuid); diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs b/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs index f1cace8ad..2fcb52f48 100644 --- a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs @@ -288,12 +288,8 @@ pub struct DetectConnectionModeStep; #[async_trait] impl StepExecutor for DetectConnectionModeStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!( "Detecting connection mode for {} (timeout: {}s, max_attempts: {})", @@ -346,12 +342,8 @@ pub struct DiscoverMetadataStep; #[async_trait] impl StepExecutor for DiscoverMetadataStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let connection_mode: Arc = context - .get("connection_mode") - .ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let connection_mode: Arc = context.get_or_err("connection_mode")?; debug!( "Discovering metadata for {} ({:?})", @@ -430,9 +422,7 @@ pub struct DiscoverDPInfoStep; #[async_trait] impl StepExecutor for DiscoverDPInfoStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; if !config.dp_aware { debug!( @@ -471,18 +461,10 @@ pub struct CreateWorkerStep; #[async_trait] impl StepExecutor for CreateWorkerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let connection_mode: Arc = context - .get("connection_mode") - .ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?; - let discovered_labels: Arc> = context - .get("discovered_labels") - .ok_or_else(|| WorkflowError::ContextValueNotFound("discovered_labels".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let connection_mode: Arc = context.get_or_err("connection_mode")?; + let discovered_labels: Arc> = context.get_or_err("discovered_labels")?; // Check if worker already exists if app_context @@ -638,9 +620,7 @@ impl StepExecutor for CreateWorkerStep { // Handle DP-aware vs non-DP-aware workers if config.dp_aware { - let dp_info: Arc = context - .get("dp_info") - .ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?; + let dp_info: Arc = context.get_or_err("dp_info")?; debug!( "Creating {} DP-aware workers for {} (dp_size: {})", @@ -724,17 +704,11 @@ pub struct RegisterWorkerStep; #[async_trait] impl StepExecutor for RegisterWorkerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; if config.dp_aware { - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let workers: Arc>> = context.get_or_err("workers")?; let mut worker_ids = Vec::new(); for worker in workers.iter() { @@ -748,9 +722,7 @@ impl StepExecutor for RegisterWorkerStep { context.set("worker_ids", worker_ids); } else { - let worker: Arc> = context - .get("worker") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; + let worker: Arc> = context.get_or_err("worker")?; let worker_id = app_context .worker_registry @@ -773,22 +745,14 @@ pub struct UpdatePoliciesStep; #[async_trait] impl StepExecutor for UpdatePoliciesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - let labels: Arc> = context - .get("labels") - .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; + let labels: Arc> = context.get_or_err("labels")?; + let app_context: Arc = context.get_or_err("app_context")?; let policy_hint = labels.get("policy").map(|s| s.as_str()); if config.dp_aware { - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let workers: Arc>> = context.get_or_err("workers")?; let model_id = workers[0].model_id().to_string(); @@ -815,9 +779,7 @@ impl StepExecutor for UpdatePoliciesStep { model_id ); } else { - let worker: Arc> = context - .get("worker") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; + let worker: Arc> = context.get_or_err("worker")?; let model_id = worker.model_id().to_string(); @@ -864,14 +826,10 @@ pub struct ActivateWorkerStep; #[async_trait] impl StepExecutor for ActivateWorkerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("worker_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let config: Arc = context.get_or_err("worker_config")?; if config.dp_aware { - let workers: Arc>> = context - .get("workers") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + let workers: Arc>> = context.get_or_err("workers")?; for worker in workers.iter() { worker.set_healthy(true); @@ -883,9 +841,7 @@ impl StepExecutor for ActivateWorkerStep { config.url ); } else { - let worker: Arc> = context - .get("worker") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; + let worker: Arc> = context.get_or_err("worker")?; worker.set_healthy(true); diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs b/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs index a26ab85ed..0f279e416 100644 --- a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs +++ b/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs @@ -32,12 +32,8 @@ pub struct FindWorkersToRemoveStep; #[async_trait] impl StepExecutor for FindWorkersToRemoveStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let request: Arc = context - .get("removal_request") - .ok_or_else(|| WorkflowError::ContextValueNotFound("removal_request".to_string()))?; - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let request: Arc = context.get_or_err("removal_request")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!( "Finding workers to remove for {} (dp_aware: {})", @@ -109,12 +105,8 @@ pub struct RemoveFromPolicyRegistryStep; #[async_trait] impl StepExecutor for RemoveFromPolicyRegistryStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let workers_to_remove: Arc>> = context - .get("workers_to_remove") - .ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_remove".to_string()))?; + let app_context: Arc = context.get_or_err("app_context")?; + let workers_to_remove: Arc>> = context.get_or_err("workers_to_remove")?; debug!( "Removing {} worker(s) from policy registry", @@ -153,12 +145,8 @@ pub struct RemoveFromWorkerRegistryStep; #[async_trait] impl StepExecutor for RemoveFromWorkerRegistryStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let worker_urls: Arc> = context - .get("worker_urls") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_urls".to_string()))?; + let app_context: Arc = context.get_or_err("app_context")?; + let worker_urls: Arc> = context.get_or_err("worker_urls")?; debug!( "Removing {} worker(s) from worker registry", @@ -202,15 +190,9 @@ pub struct UpdateRemainingPoliciesStep; #[async_trait] impl StepExecutor for UpdateRemainingPoliciesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let app_context: Arc = context - .get("app_context") - .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - let affected_models: Arc> = context - .get("affected_models") - .ok_or_else(|| WorkflowError::ContextValueNotFound("affected_models".to_string()))?; - let worker_urls: Arc> = context - .get("worker_urls") - .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_urls".to_string()))?; + let app_context: Arc = context.get_or_err("app_context")?; + let affected_models: Arc> = context.get_or_err("affected_models")?; + let worker_urls: Arc> = context.get_or_err("worker_urls")?; debug!( "Updating cache-aware policies for {} affected model(s)", diff --git a/sgl-model-gateway/src/core/workflow/types.rs b/sgl-model-gateway/src/core/workflow/types.rs index e7af5c654..f9214f09b 100644 --- a/sgl-model-gateway/src/core/workflow/types.rs +++ b/sgl-model-gateway/src/core/workflow/types.rs @@ -218,6 +218,12 @@ impl WorkflowContext { .and_then(|v| v.clone().downcast::().ok()) } + /// Retrieve a value from the context, returning an error if not found + pub fn get_or_err(&self, key: &str) -> Result, 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 pub fn has_unserializable_data(&self) -> bool { !self.data.is_empty()