Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions lib/llm/src/discovery/model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ use serde::Serialize;
use super::ModelManagerError;
use super::worker_monitor::LoadThresholdConfig;
use super::worker_set::WorkerSet;
use crate::preprocessor::OpenAIPreprocessor;
use crate::protocols::openai::ParsingOptions;

use crate::types::{
Expand Down Expand Up @@ -689,6 +690,21 @@ impl Model {
.ok_or_else(|| self.engine_error(self.has_chat_engine()))
}

/// `None` when chat is served by a Python engine factory or the model has no Rust
/// tokenizer.
pub fn get_chat_preprocessor(&self) -> Option<Arc<OpenAIPreprocessor>> {
self.select_worker_set_with(|ws| ws.chat_preprocessor.clone())
}

/// Either pipeline's preprocessor; both are built on the model's single tokenizer.
pub fn get_preprocessor(&self) -> Option<Arc<OpenAIPreprocessor>> {
self.select_worker_set_with(|ws| {
ws.completions_preprocessor
.clone()
.or_else(|| ws.chat_preprocessor.clone())
})
}

pub fn get_completions_engine(
&self,
) -> Result<OpenAICompletionsStreamingEngine, ModelManagerError> {
Expand Down
32 changes: 32 additions & 0 deletions lib/llm/src/discovery/model_manager.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::preprocessor::OpenAIPreprocessor;
use std::{
collections::{HashMap, HashSet},
sync::{
Expand Down Expand Up @@ -700,6 +701,37 @@ impl ModelManager {
.get_pooling_engine()
}

pub fn get_chat_preprocessor(&self, model: &str) -> Option<Arc<OpenAIPreprocessor>> {
self.models.get(model)?.get_chat_preprocessor()
}

/// Whichever pipeline exists; both share the model's tokenizer.
pub fn get_preprocessor(&self, model: &str) -> Option<Arc<OpenAIPreprocessor>> {
self.models.get(model)?.get_preprocessor()
}

/// Attach preprocessors for an in-process model. Discovery-backed models get theirs
/// from the watcher as it builds their pipelines.
pub fn add_model_preprocessors(
&self,
model: &str,
card_checksum: &str,
chat: Option<Arc<OpenAIPreprocessor>>,
completions: Option<Arc<OpenAIPreprocessor>>,
) -> Result<(), ModelManagerError> {
let model_entry = self.get_or_create_model(model);
let namespace = format!("__local_preprocessors_{}", model);
let mut ws = WorkerSet::new(
namespace.clone(),
card_checksum.to_string(),
Self::aggregated_local_card(),
);
ws.chat_preprocessor = chat;
ws.completions_preprocessor = completions;
model_entry.add_worker_set(namespace, Arc::new(ws));
Ok(())
}

pub fn get_completions_engine(
&self,
model: &str,
Expand Down
2 changes: 2 additions & 0 deletions lib/llm/src/discovery/watcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1824,6 +1824,7 @@ impl ModelWatcher {
let preprocessor =
OpenAIPreprocessor::new_with_parts(card.clone(), formatter, tk.clone())
.context("OpenAIPreprocessor.new_with_parts")?;
worker_set.chat_preprocessor = Some(preprocessor.clone());
Some(
routing
.build_pipeline::<
Expand Down Expand Up @@ -1866,6 +1867,7 @@ impl ModelWatcher {
let preprocessor =
OpenAIPreprocessor::new_with_parts(card.clone(), formatter, tk.clone())
.context("OpenAIPreprocessor::new_with_parts")?;
worker_set.completions_preprocessor = Some(preprocessor.clone());
let routing = preprocessed_routing.as_ref().ok_or_else(|| {
anyhow::anyhow!("completions pipeline requires preprocessed routing")
})?;
Expand Down
10 changes: 10 additions & 0 deletions lib/llm/src/discovery/worker_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use crate::{
discovery::{KvWorkerMonitor, allocator::AllocatorTrimOnDrop},
kv_router::{EncoderRouter, KvRouter, PrefillRouter},
model_card::ModelDeploymentCard,
preprocessor::OpenAIPreprocessor,
types::{
RealtimeBidirectionalEngine,
generic::tensor::TensorStreamingEngine,
Expand Down Expand Up @@ -143,6 +144,11 @@ pub struct WorkerSet {
card: ModelDeploymentCard,

// Engines — each WorkerSet owns its own pipelines
/// Retained so callers can reuse this model's template and tokenizer instead of
/// building a second copy that can drift from the pipeline's.
pub(crate) chat_preprocessor: Option<Arc<OpenAIPreprocessor>>,
pub(crate) completions_preprocessor: Option<Arc<OpenAIPreprocessor>>,

pub(crate) chat_engine: Option<OpenAIChatCompletionsStreamingEngine>,
pub(crate) completions_engine: Option<OpenAICompletionsStreamingEngine>,
pub(crate) embeddings_engine: Option<OpenAIEmbeddingsStreamingEngine>,
Expand Down Expand Up @@ -185,6 +191,8 @@ impl WorkerSet {
endpoint_id: None,
mdcsum,
card,
chat_preprocessor: None,
completions_preprocessor: None,
chat_engine: None,
completions_engine: None,
embeddings_engine: None,
Expand Down Expand Up @@ -391,6 +399,8 @@ impl WorkerSet {
endpoint_id: self.endpoint_id.clone(),
mdcsum,
card,
chat_preprocessor: self.chat_preprocessor.clone(),
completions_preprocessor: self.completions_preprocessor.clone(),
chat_engine: lora_context_engine(&self.chat_engine, &lora_name),
completions_engine: lora_context_engine(&self.completions_engine, &lora_name),
embeddings_engine: lora_context_engine(&self.embeddings_engine, &lora_name),
Expand Down
1 change: 1 addition & 0 deletions lib/llm/src/http/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ pub mod metrics;
pub mod openapi_docs;
pub mod realtime;
pub mod service_v2;
pub mod tokenize;

pub use axum;
pub use frontend_extension::{
Expand Down
12 changes: 12 additions & 0 deletions lib/llm/src/http/service/openapi_docs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,8 @@ fn generate_summary_for_path(path: &str) -> String {
"/v1/embeddings" => "Create embeddings".to_string(),
"/v1/responses" => "Create response".to_string(),
"/v1/models" => "List available models".to_string(),
"/tokenize" => "Tokenize a prompt or chat conversation".to_string(),
"/detokenize" => "Detokenize a list of token ids".to_string(),
"/health" => "Health check".to_string(),
"/live" => "Liveness check".to_string(),
"/metrics" => "Prometheus metrics".to_string(),
Expand Down Expand Up @@ -353,6 +355,16 @@ fn generate_description_for_path(path: &str) -> String {
"Lists the currently available models and provides basic information about each."
.to_string()
}
"/tokenize" => {
"Returns the token ids for a prompt (`prompt`) or a chat conversation (`messages`, \
rendered through the model's chat template) without running inference. \
Compatible with vLLM's tokenize API."
.to_string()
}
"/detokenize" => {
"Returns the text for a list of token ids. Compatible with vLLM's detokenize API."
.to_string()
}
"/health" => {
"Returns the health status of the service. Used for readiness probes."
.to_string()
Expand Down
9 changes: 9 additions & 0 deletions lib/llm/src/http/service/service_v2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -929,6 +929,10 @@ static HTTP_SVC_FILES_PATH_ENV: &str = "DYN_HTTP_SVC_FILES_PATH";
static HTTP_SVC_BATCHES_PATH_ENV: &str = "DYN_HTTP_SVC_BATCHES_PATH";
/// Environment variable to set the anthropic messages endpoint path (default: `/v1/messages`)
static HTTP_SVC_ANTHROPIC_PATH_ENV: &str = "DYN_HTTP_SVC_ANTHROPIC_PATH";
/// Environment variable to set the tokenize endpoint path (default: `/tokenize`)
static HTTP_SVC_TOKENIZE_PATH_ENV: &str = "DYN_HTTP_SVC_TOKENIZE_PATH";
/// Environment variable to set the detokenize endpoint path (default: `/detokenize`)
static HTTP_SVC_DETOKENIZE_PATH_ENV: &str = "DYN_HTTP_SVC_DETOKENIZE_PATH";
/// Environment variable to enable the experimental vLLM-compatible
/// `/inference/v1/generate` endpoint. Truthy value opts in; disabled by default.
pub(super) static VLLM_ENABLE_INFERENCE_V1_GENERATE_ENV: &str =
Expand Down Expand Up @@ -1102,6 +1106,11 @@ impl HttpServiceConfigBuilder {
},
super::health::health_check_router(state.clone(), var(HTTP_SVC_HEALTH_PATH_ENV).ok()),
super::health::live_check_router(state.clone(), var(HTTP_SVC_LIVE_PATH_ENV).ok()),
super::tokenize::tokenize_router(state.clone(), var(HTTP_SVC_TOKENIZE_PATH_ENV).ok()),
super::tokenize::detokenize_router(
state.clone(),
var(HTTP_SVC_DETOKENIZE_PATH_ENV).ok(),
),
];
if admin_api_enabled {
system_routes.push(super::busy_threshold::busy_threshold_router(
Expand Down
Loading
Loading