Code Rabbit feedback

Author: grahamking

components/frontend/src/dynamo/frontend/main.py

@@ -128,12 +128,18 @@ def parse_args():
         help="Path to model directory on disk (e.g., /tmp/model_cache/lama3.2_1B/)",
     )
 
-    return parser.parse_args()
+    flags = parser.parse_args()
+
+    if flags.static_endpoint and (not flags.model_name or not flags.model_path):
+        parser.error("--static-endpoint requires both --model-name and --model-path")
+
+    return flags
 
 
 async def async_main():
     flags = parse_args()
     is_static = bool(flags.static_endpoint)  # true if the string has a value
+
     runtime = DistributedRuntime(asyncio.get_running_loop(), is_static)
 
     if flags.router_mode == "kv":

launch/dynamo-run/src/flags.rs

@@ -171,7 +171,13 @@ impl Flags {
                 }
             }
             Output::Static(_) => {
-                if self.model_name.is_none() || self.model_name.is_none() {
+                if self.model_name.is_none()
+                    || self
+                        .model_path_pos
+                        .as_ref()
+                        .or(self.model_path_flag.as_ref())
+                        .is_none()
+                {
                     anyhow::bail!(
                         "out=dyn://<path> requires --model-name and --model-path, which are the name and path on disk of the model we expect to serve."
                     );

lib/bindings/python/examples/hello_world/server_sglang_static.py

@@ -62,6 +62,7 @@ async def generate(self, request):
             if finish_reason:
                 # Don't forward the stop token
                 out = {"token_ids": [], "finish_reason": finish_reason["type"]}
+                next_total_toks = num_output_tokens_so_far
             else:
                 next_total_toks = len(res["output_ids"])
                 out = {"token_ids": res["output_ids"][num_output_tokens_so_far:]}

lib/llm/src/discovery/watcher.rs

@@ -18,7 +18,7 @@ use dynamo_runtime::{
 
 use crate::{
     backend::Backend,
-    engines,
+    entrypoint,
     kv_router::KvRouterConfig,
     model_type::ModelType,
     preprocessor::{OpenAIPreprocessor, PreprocessedEmbeddingRequest},
@@ -216,19 +216,21 @@ impl ModelWatcher {
                     None
                 };
 
-                let chat_engine = engines::build_chat_completions(
-                    &card,
-                    &client,
-                    self.router_mode,
-                    kv_chooser.clone(),
-                )
-                .await?;
+                let chat_engine =
+                    entrypoint::build_routed_pipeline::<
+                        NvCreateChatCompletionRequest,
+                        NvCreateChatCompletionStreamResponse,
+                    >(&card, &client, self.router_mode, kv_chooser.clone())
+                    .await?;
                 self.manager
                     .add_chat_completions_model(&model_entry.name, chat_engine)?;
 
                 let completions_engine =
-                    engines::build_completions(&card, &client, self.router_mode, kv_chooser)
-                        .await?;
+                    entrypoint::build_routed_pipeline::<
+                        NvCreateCompletionRequest,
+                        NvCreateCompletionResponse,
+                    >(&card, &client, self.router_mode, kv_chooser)
+                    .await?;
                 self.manager
                     .add_completions_model(&model_entry.name, completions_engine)?;
             }

lib/llm/src/engines.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 std::env;
 use std::sync::Arc;
@@ -21,32 +9,17 @@ use std::time::Duration;
 use async_stream::stream;
 use async_trait::async_trait;
 
-use dynamo_runtime::component::Client;
 use dynamo_runtime::engine::{AsyncEngine, AsyncEngineContextProvider, ResponseStream};
-use dynamo_runtime::pipeline::Operator as _;
-use dynamo_runtime::pipeline::PushRouter;
-use dynamo_runtime::pipeline::RouterMode;
-use dynamo_runtime::pipeline::SegmentSource;
-use dynamo_runtime::pipeline::ServiceBackend;
-use dynamo_runtime::pipeline::Source as _;
 use dynamo_runtime::pipeline::{Error, ManyOut, SingleIn};
 use dynamo_runtime::protocols::annotated::Annotated;
 
-use crate::backend::Backend;
 use crate::backend::ExecutionContext;
-use crate::kv_router::KvPushRouter;
-use crate::kv_router::KvRouter;
-use crate::migration::Migration;
-use crate::model_card::ModelDeploymentCard;
-use crate::preprocessor::OpenAIPreprocessor;
 use crate::preprocessor::PreprocessedRequest;
 use crate::protocols::common::llm_backend::LLMEngineOutput;
 use crate::protocols::openai::{
     chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse},
     completions::{prompt_to_string, NvCreateCompletionRequest, NvCreateCompletionResponse},
 };
-use crate::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine;
-use crate::types::openai::completions::OpenAICompletionsStreamingEngine;
 use crate::types::openai::embeddings::NvCreateEmbeddingRequest;
 use crate::types::openai::embeddings::NvCreateEmbeddingResponse;
 
@@ -436,89 +409,3 @@ impl
         self.0.handle_chat(req).await
     }
 }
-
-pub async fn build_chat_completions(
-    card: &ModelDeploymentCard,
-    client: &Client,
-    router_mode: RouterMode,
-    chooser: Option<Arc<KvRouter>>,
-) -> anyhow::Result<OpenAIChatCompletionsStreamingEngine> {
-    let frontend = SegmentSource::<
-        SingleIn<NvCreateChatCompletionRequest>,
-        ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
-    >::new();
-    let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
-    let backend = Backend::from_mdc(card.clone()).await?.into_operator();
-    let migration = Migration::from_mdc(card.clone()).await?.into_operator();
-    let router = PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client(
-        client.clone(),
-        router_mode,
-    )
-    .await?;
-    let service_backend = match router_mode {
-        RouterMode::Random | RouterMode::RoundRobin | RouterMode::Direct(_) => {
-            ServiceBackend::from_engine(Arc::new(router))
-        }
-        RouterMode::KV => {
-            let Some(chooser) = chooser else {
-                anyhow::bail!("RouterMode::KV requires KVRouter to not be null");
-            };
-            let kv_push_router = KvPushRouter::new(router, chooser);
-            ServiceBackend::from_engine(Arc::new(kv_push_router))
-        }
-    };
-
-    let chat_engine = frontend
-        .link(preprocessor.forward_edge())?
-        .link(backend.forward_edge())?
-        .link(migration.forward_edge())?
-        .link(service_backend)?
-        .link(migration.backward_edge())?
-        .link(backend.backward_edge())?
-        .link(preprocessor.backward_edge())?
-        .link(frontend)?;
-    Ok(chat_engine)
-}
-
-pub async fn build_completions(
-    card: &ModelDeploymentCard,
-    client: &Client,
-    router_mode: RouterMode,
-    chooser: Option<Arc<KvRouter>>,
-) -> anyhow::Result<OpenAICompletionsStreamingEngine> {
-    let frontend = SegmentSource::<
-        SingleIn<NvCreateCompletionRequest>,
-        ManyOut<Annotated<NvCreateCompletionResponse>>,
-    >::new();
-    let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
-    let backend = Backend::from_mdc(card.clone()).await?.into_operator();
-    let migration = Migration::from_mdc(card.clone()).await?.into_operator();
-    let router = PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client(
-        client.clone(),
-        router_mode,
-    )
-    .await?;
-    let service_backend = match router_mode {
-        RouterMode::Random | RouterMode::RoundRobin | RouterMode::Direct(_) => {
-            ServiceBackend::from_engine(Arc::new(router))
-        }
-        RouterMode::KV => {
-            let Some(chooser) = chooser else {
-                anyhow::bail!("RouterMode::KV requires KVRouter to not be null");
-            };
-            let kv_push_router = KvPushRouter::new(router, chooser);
-            ServiceBackend::from_engine(Arc::new(kv_push_router))
-        }
-    };
-
-    let completions_engine = frontend
-        .link(preprocessor.forward_edge())?
-        .link(backend.forward_edge())?
-        .link(migration.forward_edge())?
-        .link(service_backend)?
-        .link(migration.backward_edge())?
-        .link(backend.backward_edge())?
-        .link(preprocessor.backward_edge())?
-        .link(frontend)?;
-    Ok(completions_engine)
-}

lib/llm/src/entrypoint.rs

@@ -6,6 +6,7 @@
 //! - Connect it to an Input
 
 pub mod input;
+pub use input::build_routed_pipeline;
 
 use std::sync::Arc;
 

lib/llm/src/entrypoint/input.rs

@@ -16,6 +16,7 @@ use std::{
 
 pub mod batch;
 mod common;
+pub use common::build_routed_pipeline;
 pub mod endpoint;
 pub mod http;
 pub mod text;

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

@@ -6,11 +6,13 @@ use std::pin::Pin;
 use crate::{
     backend::{Backend, ExecutionContext},
     discovery::{ModelManager, ModelWatcher, MODEL_ROOT_PATH},
-    engines::{self, StreamingEngineAdapter},
-    entrypoint::EngineConfig,
+    engines::StreamingEngineAdapter,
+    entrypoint::{self, EngineConfig},
+    kv_router::{KvPushRouter, KvRouter},
+    migration::Migration,
     model_card::ModelDeploymentCard,
     preprocessor::OpenAIPreprocessor,
-    protocols::common::llm_backend::{BackendOutput, PreprocessedRequest},
+    protocols::common::llm_backend::{BackendOutput, LLMEngineOutput, PreprocessedRequest},
     request_template::RequestTemplate,
     types::{
         openai::chat_completions::{
@@ -21,10 +23,12 @@ use crate::{
     },
 };
 use dynamo_runtime::{
+    component::Client,
     distributed::DistributedConfig,
     engine::{AsyncEngineStream, Data},
     pipeline::{
-        Context, ManyOut, Operator, RouterMode, ServiceBackend, ServiceFrontend, SingleIn, Source,
+        Context, ManyOut, Operator, PushRouter, RouterMode, SegmentSource, ServiceBackend,
+        ServiceEngine, ServiceFrontend, SingleIn, Source,
     },
     DistributedRuntime, Runtime,
 };
@@ -124,9 +128,11 @@ pub async fn prepare_engine(
                 None
             };
 
-            let chat_engine =
-                engines::build_chat_completions(card, &client, router_mode, kv_chooser.clone())
-                    .await?;
+            let chat_engine = entrypoint::build_routed_pipeline::<
+                NvCreateChatCompletionRequest,
+                NvCreateChatCompletionStreamResponse,
+            >(card, &client, router_mode, kv_chooser.clone())
+            .await?;
 
             let service_name = local_model.service_name().to_string();
             tracing::info!("Static connecting to {service_name}");
@@ -204,6 +210,56 @@ where
         .link(frontend)?)
 }
 
+pub async fn build_routed_pipeline<Req, Resp>(
+    card: &ModelDeploymentCard,
+    client: &Client,
+    router_mode: RouterMode,
+    chooser: Option<Arc<KvRouter>>,
+) -> anyhow::Result<ServiceEngine<SingleIn<Req>, ManyOut<Annotated<Resp>>>>
+where
+    Req: Data,
+    Resp: Data,
+    OpenAIPreprocessor: Operator<
+        Context<Req>,
+        Pin<Box<dyn AsyncEngineStream<Annotated<Resp>>>>,
+        Context<PreprocessedRequest>,
+        Pin<Box<dyn AsyncEngineStream<Annotated<BackendOutput>>>>,
+    >,
+{
+    let frontend = SegmentSource::<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 migration = Migration::from_mdc(card.clone()).await?.into_operator();
+    let router = PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client(
+        client.clone(),
+        router_mode,
+    )
+    .await?;
+    let service_backend = match router_mode {
+        RouterMode::Random | RouterMode::RoundRobin | RouterMode::Direct(_) => {
+            ServiceBackend::from_engine(Arc::new(router))
+        }
+        RouterMode::KV => {
+            let Some(chooser) = chooser else {
+                anyhow::bail!("RouterMode::KV requires KVRouter to not be null");
+            };
+            let kv_push_router = KvPushRouter::new(router, chooser);
+            ServiceBackend::from_engine(Arc::new(kv_push_router))
+        }
+    };
+
+    let engine = frontend
+        .link(preprocessor.forward_edge())?
+        .link(backend.forward_edge())?
+        .link(migration.forward_edge())?
+        .link(service_backend)?
+        .link(migration.backward_edge())?
+        .link(backend.backward_edge())?
+        .link(preprocessor.backward_edge())?
+        .link(frontend)?;
+    Ok(engine)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -5,8 +5,8 @@ use std::sync::Arc;
 
 use crate::{
     discovery::{ModelManager, ModelWatcher, MODEL_ROOT_PATH},
-    engines::{self, StreamingEngineAdapter},
-    entrypoint::{input::common, EngineConfig},
+    engines::StreamingEngineAdapter,
+    entrypoint::{self, input::common, EngineConfig},
     http::service::service_v2,
     kv_router::KvRouterConfig,
     types::openai::{
@@ -78,13 +78,18 @@ pub async fn run(runtime: Runtime, engine_config: EngineConfig) -> anyhow::Resul
                 None
             };
 
-            let chat_engine =
-                engines::build_chat_completions(card, &client, router_mode, kv_chooser.clone())
-                    .await?;
+            let chat_engine = entrypoint::build_routed_pipeline::<
+                NvCreateChatCompletionRequest,
+                NvCreateChatCompletionStreamResponse,
+            >(card, &client, router_mode, kv_chooser.clone())
+            .await?;
             manager.add_chat_completions_model(local_model.display_name(), chat_engine)?;
 
-            let completions_engine =
-                engines::build_completions(card, &client, router_mode, kv_chooser).await?;
+            let completions_engine = entrypoint::build_routed_pipeline::<
+                NvCreateCompletionRequest,
+                NvCreateCompletionResponse,
+            >(card, &client, router_mode, kv_chooser)
+            .await?;
             manager.add_completions_model(local_model.display_name(), completions_engine)?;
         }
         EngineConfig::StaticFull { engine, model, .. } => {