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

This commit is contained in:
Chang Su
2026-01-02 16:05:10 -08:00
committed by GitHub
parent c4edcac6d7
commit 24c91001cf
3 changed files with 9 additions and 25 deletions
@@ -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) {
@@ -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: {}",
@@ -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();