[model-gateway] code clean up in tokenizer register step workflow (#16316)

sgl-model-gateway/src/core/metrics_aggregator.rs

@@ -16,7 +16,7 @@ pub fn aggregate_metrics(metric_packs: Vec<MetricPack>) -> anyhow::Result<String
     let mut expositions = vec![];
     for metric_pack in metric_packs {
         let metrics_text = &metric_pack.metrics_text;
-        // Hacky workaround since the parser do not understand `:`, should improve later
+        // openmetrics_parser doesn't handle colons in metric names; replace with underscores
         let metrics_text = metrics_text.replace(":", "_");
 
         let exposition = match openmetrics_parser::prometheus::parse_prometheus(&metrics_text) {

sgl-model-gateway/src/core/steps/tokenizer_registration.rs

@@ -41,13 +41,8 @@ pub struct ValidateTokenizerConfigStep;
 #[async_trait]
 impl StepExecutor for ValidateTokenizerConfigStep {
     async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
-        let config: Arc<TokenizerConfigRequest> = context
+        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
-            .get("tokenizer_config")
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("tokenizer_config".to_string()))?;
-
-        let app_context: Arc<AppContext> = context
-            .get("app_context")
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         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<StepResult> {
-        let config: Arc<TokenizerConfigRequest> = context
+        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
-            .get("tokenizer_config")
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("tokenizer_config".to_string()))?;
-
-        let app_context: Arc<AppContext> = context
-            .get("app_context")
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         info!(
             "Loading tokenizer '{}' (id: {}) from source: {}",

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<StepResult> {
-        let labels: Arc<HashMap<String, String>> = context
+        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
-            .get("labels")
+        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
-            .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".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()))?;
 
         for worker in workers.iter() {
             let model_id = worker.model_id().to_string();