From 24c91001cf99ba642be791e099d358f4dfe955f5 Mon Sep 17 00:00:00 2001
From: Chang Su <chang.s.su@oracle.com>
Date: Fri, 2 Jan 2026 16:05:10 -0800
Subject: [PATCH] [model-gateway] code clean up in tokenizer register step
 workflow (#16316)

---
 .../src/core/metrics_aggregator.rs             |  2 +-
 .../src/core/steps/tokenizer_registration.rs   | 18 ++++--------------
 .../steps/worker/local/register_tokenizer.rs   | 14 ++++----------
 3 files changed, 9 insertions(+), 25 deletions(-)

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<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) {
diff --git a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs
index 31412ca4d..bbd0218ce 100644
--- a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs
+++ b/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
-            .get("tokenizer_config")
-            .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()))?;
+        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
+        let app_context: Arc<AppContext> = 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<StepResult> {
-        let config: Arc<TokenizerConfigRequest> = context
-            .get("tokenizer_config")
-            .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()))?;
+        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
+        let app_context: Arc<AppContext> = 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<StepResult> {
-        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 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")?;
 
         for worker in workers.iter() {
             let model_id = worker.model_id().to_string();