chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 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]
impl StepExecutor for DiscoverModelsStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let config: Arc<WorkerConfigRequest> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let model_cards: Arc<Vec<ModelCard>> = context
.get("model_cards")
.ok_or_else(|| WorkflowError::ContextValueNotFound("model_cards".to_string()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let model_cards: Arc<Vec<ModelCard>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
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 workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
for worker in workers.iter() {
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]
impl StepExecutor for ConnectMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context
.get("mcp_server_config")
.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()))?;
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = 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<StepResult> {
use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context
.get("mcp_server_config")
.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()))?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context
.get("mcp_client")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = 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<StepResult> {
use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context
.get("mcp_server_config")
.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()))?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context
.get("mcp_client")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = 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<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context
.get("mcp_server_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_server_config".to_string()))?;
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?;
let client_registered = context
.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]
impl StepExecutor for ValidateDescriptorStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let config_request: Arc<WasmModuleConfigRequest> = 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<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let config_request: Arc<WasmModuleConfigRequest> = 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<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
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()))?;
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = 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<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
let config_request: Arc<WasmModuleConfigRequest> = 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<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.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()))?;
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
let wasm_bytes: Arc<Vec<u8>> = 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<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context
.get("wasm_module_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_module_config".to_string()))?;
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()))?;
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()))?;
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?;
let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
debug!(
"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 {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> =
context.get("wasm_module_removal_request").ok_or_else(|| {
WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string())
})?;
let app_context: Arc<AppContext> = 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<AppContext> = 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<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> =
context.get("wasm_module_removal_request").ok_or_else(|| {
WorkflowError::ContextValueNotFound("wasm_module_removal_request".to_string())
})?;
let app_context: Arc<AppContext> = 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<AppContext> = context.get_or_err("app_context")?;
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]
impl StepExecutor for DetectConnectionModeStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let connection_mode: Arc<ConnectionMode> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let config: Arc<WorkerConfigRequest> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
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()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let discovered_labels: Arc<HashMap<String, String>> = 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<DpInfo> = context
.get("dp_info")
.ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?;
let dp_info: Arc<DpInfo> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.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()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = context
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
let worker: Arc<Arc<dyn Worker>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
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 config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let policy_hint = labels.get("policy").map(|s| s.as_str());
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = context
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
let worker: Arc<Arc<dyn Worker>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context
.get("worker_config")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?;
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = context
.get("worker")
.ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?;
let worker: Arc<Arc<dyn Worker>> = context.get_or_err("worker")?;
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]
impl StepExecutor for FindWorkersToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerRemovalRequest> = context
.get("removal_request")
.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()))?;
let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
let app_context: Arc<AppContext> = 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<StepResult> {
let app_context: Arc<AppContext> = context
.get("app_context")
.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()))?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> = 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<StepResult> {
let app_context: Arc<AppContext> = context
.get("app_context")
.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()))?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let worker_urls: Arc<Vec<String>> = 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<StepResult> {
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
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()))?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
debug!(
"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())
}
/// 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
pub fn has_unserializable_data(&self) -> bool {
!self.data.is_empty()