[model-gateway] code clean up in tokenizer register step workflow (#16316)
This commit is contained in:
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user