diff --git a/sgl-model-gateway/src/core/metrics_aggregator.rs b/sgl-model-gateway/src/core/metrics_aggregator.rs index 360736095..f39c96a52 100644 --- a/sgl-model-gateway/src/core/metrics_aggregator.rs +++ b/sgl-model-gateway/src/core/metrics_aggregator.rs @@ -16,7 +16,7 @@ pub fn aggregate_metrics(metric_packs: Vec) -> anyhow::Result WorkflowResult { - let config: Arc = context - .get("tokenizer_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("tokenizer_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("tokenizer_config")?; + let app_context: Arc = context.get_or_err("app_context")?; debug!( "Validating tokenizer config: name={}, source={}", @@ -93,13 +88,8 @@ pub struct LoadTokenizerStep; #[async_trait] impl StepExecutor for LoadTokenizerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context - .get("tokenizer_config") - .ok_or_else(|| WorkflowError::ContextValueNotFound("tokenizer_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("tokenizer_config")?; + let app_context: Arc = context.get_or_err("app_context")?; info!( "Loading tokenizer '{}' (id: {}) from source: {}", diff --git a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs index df6e37880..b3bfe4318 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs @@ -1,4 +1,4 @@ -//! Connection mode detection step. +//! Tokenizer registration step for local workers. use std::{collections::HashMap, sync::Arc}; @@ -18,15 +18,9 @@ pub struct RegisterTokenizerStep; #[async_trait] impl StepExecutor for RegisterTokenizerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - 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 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")?; for worker in workers.iter() { let model_id = worker.model_id().to_string();