fix: Load the tokenizer JSON once for chat and completions. (#2910)

Signed-off-by: Graham King <grahamk@nvidia.com>

Author: grahamking

lib/bindings/python/rust/llm/backend.rs

@@ -1,17 +1,5 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
 
 use super::*;
 use crate::llm::model_card::ModelDeploymentCard;
@@ -33,10 +21,7 @@ pub(crate) struct Backend {
 impl Backend {
     #[new]
     fn new(mdc: ModelDeploymentCard, endpoint: Endpoint) -> PyResult<Self> {
-        let runtime = pyo3_async_runtimes::tokio::get_runtime();
-        let backend = runtime
-            .block_on(llm_rs::backend::Backend::from_mdc(mdc.inner))
-            .map_err(to_pyerr)?;
+        let backend = llm_rs::backend::Backend::from_mdc(&mdc.inner);
         Ok(Self {
             inner: backend,
             endpoint,

lib/bindings/python/rust/llm/model_card.rs

@@ -16,14 +16,10 @@ impl ModelDeploymentCard {}
 impl ModelDeploymentCard {
     // Previously called "from_local_path"
     #[staticmethod]
-    fn load(path: String, model_name: String, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let mut card = RsModelDeploymentCard::load(&path, None)
-                .await
-                .map_err(to_pyerr)?;
-            card.set_name(&model_name);
-            Ok(ModelDeploymentCard { inner: card })
-        })
+    fn load(path: String, model_name: String) -> PyResult<ModelDeploymentCard> {
+        let mut card = RsModelDeploymentCard::load(&path, None).map_err(to_pyerr)?;
+        card.set_name(&model_name);
+        Ok(ModelDeploymentCard { inner: card })
     }
 
     #[staticmethod]

lib/bindings/python/rust/llm/preprocessor.rs

@@ -1,17 +1,5 @@
 // SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
 
 use super::*;
 use crate::llm::model_card::ModelDeploymentCard;
@@ -42,10 +30,7 @@ pub(crate) struct OAIChatPreprocessor {
 impl OAIChatPreprocessor {
     #[new]
     fn new(mdc: ModelDeploymentCard, current: Endpoint, next: Endpoint) -> PyResult<Self> {
-        let runtime = pyo3_async_runtimes::tokio::get_runtime();
-        let preprocessor = runtime
-            .block_on(OpenAIPreprocessor::new(mdc.inner.clone()))
-            .map_err(to_pyerr)?;
+        let preprocessor = OpenAIPreprocessor::new(mdc.inner.clone()).map_err(to_pyerr)?;
         Ok(Self {
             inner: preprocessor,
             current,

lib/llm/src/backend.rs

@@ -17,11 +17,11 @@
 
 use std::{collections::HashSet, sync::Arc};
 
-use anyhow::{Error, Result};
+use anyhow::Result;
 use futures::stream::{self, StreamExt};
 use tracing as log;
 
-use crate::model_card::{ModelDeploymentCard, TokenizerKind};
+use crate::model_card::ModelDeploymentCard;
 use dynamo_runtime::{
     pipeline::{
         AsyncEngineContextProvider, ManyOut, Operator, ResponseStream, ServerStreamingEngine,
@@ -66,30 +66,27 @@ struct DecoderUnfoldState {
 }
 
 impl Backend {
-    pub async fn from_tokenizer(tokenizer: HfTokenizer) -> Result<Arc<Self>> {
+    pub fn from_tokenizer(tokenizer: HfTokenizer) -> Arc<Self> {
         let tokenizer = HuggingFaceTokenizer::from_tokenizer(tokenizer);
         let tokenizer = Tokenizer::from(Arc::new(tokenizer));
 
-        Ok(Arc::new(Self {
+        Arc::new(Self {
             tokenizer: Some(tokenizer),
             validate_engine_decode: false,
-        }))
+        })
     }
 
-    pub async fn from_mdc(mdc: ModelDeploymentCard) -> Result<Arc<Self>> {
-        let tokenizer = match &mdc.tokenizer {
-            Some(TokenizerKind::HfTokenizerJson(file)) => {
-                HfTokenizer::from_file(file).map_err(Error::msg)?
-            }
-            Some(TokenizerKind::GGUF(t)) => *t.clone(),
-            None => {
-                return Ok(Arc::new(Self {
+    pub fn from_mdc(mdc: &ModelDeploymentCard) -> Arc<Self> {
+        match mdc.tokenizer_hf() {
+            Ok(tokenizer) => Self::from_tokenizer(tokenizer),
+            Err(err) => {
+                tracing::warn!(%err, "tokenizer_hf error converting ModelDeploymentCard to HF tokenizer");
+                Arc::new(Self {
                     tokenizer: None,
                     validate_engine_decode: false,
-                }));
+                })
             }
-        };
-        Self::from_tokenizer(tokenizer).await
+        }
     }
 
     fn decoder(

lib/llm/src/discovery/watcher.rs

@@ -176,13 +176,13 @@ impl ModelWatcher {
                 }
                 WatchEvent::Delete(kv) => match self.handle_delete(&kv).await {
                     Ok(Some(model_name)) => {
-                        tracing::info!("removed model {}", model_name);
+                        tracing::info!(model_name, "removed model");
                     }
                     Ok(None) => {
                         // There are other instances running this model, nothing to do
                     }
                     Err(e) => {
-                        tracing::error!("error removing model: {}", e);
+                        tracing::error!(error = %e, "error removing model");
                     }
                 },
             }
@@ -271,7 +271,7 @@ impl ModelWatcher {
                 Some(card)
             }
             Err(err) => {
-                tracing::info!(%err, "load_mdc did not complete");
+                tracing::info!(error = %err, "load_mdc did not complete");
                 None
             }
         };
@@ -308,6 +308,9 @@ impl ModelWatcher {
                 None
             };
 
+            // This is expensive, we are loading ~10MiB JSON, so only do it once
+            let tokenizer_hf = card.tokenizer_hf()?;
+
             // Add chat engine only if the model supports chat
             if model_entry.model_type.supports_chat() {
                 let chat_engine = entrypoint::build_routed_pipeline::<
@@ -319,18 +322,23 @@ impl ModelWatcher {
                     self.router_mode,
                     self.busy_threshold,
                     kv_chooser.clone(),
+                    tokenizer_hf.clone(),
                 )
                 .await?;
                 self.manager
                     .add_chat_completions_model(&model_entry.name, chat_engine)?;
+                tracing::info!("Chat completions is ready");
             }
 
             // Add completions engine only if the model supports completions
             if model_entry.model_type.supports_completions() {
                 let formatter = PromptFormatter::no_op();
                 let PromptFormatter::OAI(formatter) = formatter;
-                let preprocessor =
-                    OpenAIPreprocessor::new_with_formatter(card.clone(), formatter).await?;
+                let preprocessor = OpenAIPreprocessor::new_with_parts(
+                    card.clone(),
+                    formatter,
+                    tokenizer_hf.clone(),
+                )?;
                 let completions_engine = entrypoint::build_routed_pipeline_with_preprocessor::<
                     NvCreateCompletionRequest,
                     NvCreateCompletionResponse,
@@ -341,10 +349,12 @@ impl ModelWatcher {
                     self.busy_threshold,
                     kv_chooser,
                     preprocessor,
+                    tokenizer_hf,
                 )
                 .await?;
                 self.manager
                     .add_completions_model(&model_entry.name, completions_engine)?;
+                tracing::info!("Completions is ready");
             }
         } else if model_entry.model_input == ModelInput::Text
             && model_entry.model_type.supports_chat()
@@ -391,8 +401,8 @@ impl ModelWatcher {
                 ManyOut<Annotated<NvCreateEmbeddingResponse>>,
             >::new();
 
-            let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
-            let backend = Backend::from_mdc(card.clone()).await?.into_operator();
+            let preprocessor = OpenAIPreprocessor::new(card.clone())?.into_operator();
+            let backend = Backend::from_mdc(&card).into_operator();
 
             let router = PushRouter::<
                 PreprocessedEmbeddingRequest,

lib/llm/src/entrypoint/input/batch.rs

@@ -67,7 +67,9 @@ pub async fn run(
     let mut prepared_engine = common::prepare_engine(runtime, engine_config).await?;
 
     let pre_processor = if prepared_engine.has_tokenizer() {
-        Some(OpenAIPreprocessor::new(prepared_engine.card.take().unwrap()).await?)
+        Some(OpenAIPreprocessor::new(
+            prepared_engine.card.take().unwrap(),
+        )?)
     } else {
         None
     };

lib/llm/src/entrypoint/input/common.rs

@@ -11,7 +11,7 @@ use crate::{
     kv_router::{KvPushRouter, KvRouter},
     migration::Migration,
     model_card::ModelDeploymentCard,
-    preprocessor::OpenAIPreprocessor,
+    preprocessor::{OpenAIPreprocessor, prompt::PromptFormatter},
     protocols::common::llm_backend::{BackendOutput, LLMEngineOutput, PreprocessedRequest},
     request_template::RequestTemplate,
     types::{
@@ -131,10 +131,18 @@ pub async fn prepare_engine(
                 None
             };
 
+            let hf_tokenizer = card.tokenizer_hf()?;
             let chat_engine = entrypoint::build_routed_pipeline::<
                 NvCreateChatCompletionRequest,
                 NvCreateChatCompletionStreamResponse,
-            >(card, &client, router_mode, None, kv_chooser.clone())
+            >(
+                card,
+                &client,
+                router_mode,
+                None,
+                kv_chooser.clone(),
+                hf_tokenizer,
+            )
             .await?;
 
             let service_name = local_model.service_name().to_string();
@@ -167,7 +175,7 @@ pub async fn prepare_engine(
             let pipeline = build_pipeline::<
                 NvCreateChatCompletionRequest,
                 NvCreateChatCompletionStreamResponse,
-            >(model.card(), inner_engine)
+            >(model.card(), inner_engine, model.card().tokenizer_hf()?)
             .await?;
 
             let service_name = model.service_name().to_string();
@@ -186,6 +194,7 @@ pub async fn prepare_engine(
 pub async fn build_pipeline<Req, Resp>(
     card: &ModelDeploymentCard,
     engine: ExecutionContext,
+    hf_tokenizer: tokenizers::Tokenizer,
 ) -> anyhow::Result<Arc<ServiceFrontend<SingleIn<Req>, ManyOut<Annotated<Resp>>>>>
 where
     Req: Data,
@@ -198,10 +207,11 @@ where
         >,
 {
     let frontend = ServiceFrontend::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
-    let preprocessor = OpenAIPreprocessor::new((*card).clone())
-        .await?
-        .into_operator();
-    let backend = Backend::from_mdc((*card).clone()).await?.into_operator();
+    let PromptFormatter::OAI(formatter) = PromptFormatter::from_mdc(card)?;
+    let preprocessor =
+        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, hf_tokenizer.clone())?
+            .into_operator();
+    let backend = Backend::from_tokenizer(hf_tokenizer).into_operator();
     let engine = ServiceBackend::from_engine(engine);
 
     Ok(frontend
@@ -219,6 +229,7 @@ pub async fn build_routed_pipeline<Req, Resp>(
     router_mode: RouterMode,
     busy_threshold: Option<f64>,
     chooser: Option<Arc<KvRouter>>,
+    hf_tokenizer: tokenizers::Tokenizer,
 ) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
 where
     Req: Data,
@@ -230,14 +241,17 @@ where
             Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
         >,
 {
-    let preprocessor = OpenAIPreprocessor::new(card.clone()).await?;
+    let PromptFormatter::OAI(formatter) = PromptFormatter::from_mdc(card)?;
+    let preprocessor =
+        OpenAIPreprocessor::new_with_parts(card.clone(), formatter, hf_tokenizer.clone())?;
     build_routed_pipeline_with_preprocessor(
         card,
         client,
         router_mode,
         busy_threshold,
         chooser,
         preprocessor,
+        hf_tokenizer,
     )
     .await
 }
@@ -249,6 +263,7 @@ pub async fn build_routed_pipeline_with_preprocessor<Req, Resp>(
     busy_threshold: Option<f64>,
     chooser: Option<Arc<KvRouter>>,
     preprocessor: Arc<OpenAIPreprocessor>,
+    hf_tokenizer: tokenizers::Tokenizer,
 ) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
 where
     Req: Data,
@@ -262,8 +277,8 @@ where
 {
     let frontend = SegmentSource::<SingleIn<Req>, ManyOut<Annotated<Resp>>>::new();
     let preprocessor_op = preprocessor.into_operator();
-    let backend = Backend::from_mdc(card.clone()).await?.into_operator();
-    let migration = Migration::from_mdc(card.clone()).await?.into_operator();
+    let backend = Backend::from_tokenizer(hf_tokenizer).into_operator();
+    let migration = Migration::from_mdc(card).into_operator();
     let router =
         PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client_with_threshold(
             client.clone(),
@@ -312,14 +327,14 @@ mod tests {
     #[tokio::test]
     async fn test_build_chat_completions_pipeline_core_engine_succeeds() -> anyhow::Result<()> {
         // Create test model card
-        let card = ModelDeploymentCard::load(HF_PATH, None).await?;
+        let card = ModelDeploymentCard::load(HF_PATH, None)?;
         let engine = crate::engines::make_engine_core();
 
         // Build pipeline for chat completions
         let pipeline = build_pipeline::<
             NvCreateChatCompletionRequest,
             NvCreateChatCompletionStreamResponse,
-        >(&card, engine)
+        >(&card, engine, card.tokenizer_hf()?)
         .await?;
 
         // Verify pipeline was created
@@ -331,13 +346,16 @@ mod tests {
     #[tokio::test]
     async fn test_build_completions_pipeline_core_engine_succeeds() -> anyhow::Result<()> {
         // Create test model card
-        let card = ModelDeploymentCard::load(HF_PATH, None).await?;
+        let card = ModelDeploymentCard::load(HF_PATH, None)?;
         let engine = crate::engines::make_engine_core();
 
         // Build pipeline for completions
-        let pipeline =
-            build_pipeline::<NvCreateCompletionRequest, NvCreateCompletionResponse>(&card, engine)
-                .await?;
+        let pipeline = build_pipeline::<NvCreateCompletionRequest, NvCreateCompletionResponse>(
+            &card,
+            engine,
+            card.tokenizer_hf()?,
+        )
+        .await?;
 
         // Verify pipeline was created
         assert!(Arc::strong_count(&pipeline) >= 1);

lib/llm/src/entrypoint/input/endpoint.rs

@@ -73,9 +73,7 @@ pub async fn run(
                 SingleIn<PreprocessedRequest>,
                 ManyOut<Annotated<BackendOutput>>,
             >::new();
-            let backend = Backend::from_mdc(model.card().clone())
-                .await?
-                .into_operator();
+            let backend = Backend::from_mdc(model.card()).into_operator();
             let engine = ServiceBackend::from_engine(inner_engine);
             let pipeline = frontend
                 .link(backend.forward_edge())?