dynamo-llm 1.5.1

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::path::PathBuf;
use std::sync::Arc;

use crate::{
    discovery::{ModelManager, ModelWatcher},
    engines::StreamingEngineAdapter,
    entrypoint::{EngineConfig, RouterConfig, input::common},
    grpc::service::kserve,
    http::service::metrics::Metrics,
    local_model::runtime_config::TokenizerBackend,
    namespace::NamespaceFilter,
    types::openai::{
        chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse},
        completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
    },
};
use dynamo_runtime::DistributedRuntime;

/// Build and run an KServe gRPC service
pub async fn run(
    distributed_runtime: DistributedRuntime,
    engine_config: EngineConfig,
) -> anyhow::Result<()> {
    let mut grpc_service_builder = kserve::KserveService::builder()
        .port(engine_config.local_model().http_port()) // [WIP] generalize port..
        .metrics_prefix(engine_config.local_model().metrics_prefix())
        .http_cancel_token(Some(distributed_runtime.primary_token()))
        .with_request_template(engine_config.local_model().request_template());

    // Set HTTP metrics port if provided (for parallel test execution)
    if let Some(http_metrics_port) = engine_config.local_model().http_metrics_port() {
        grpc_service_builder = grpc_service_builder.http_metrics_port(http_metrics_port);
    }

    let grpc_service = match engine_config {
        EngineConfig::Dynamic {
            ref model,
            ref prefill_load_estimator,
            ..
        } => {
            let grpc_service = grpc_service_builder.build()?;
            let router_config = model.router_config();
            let migration_limit = model.migration_limit();
            let migration_max_seq_len = model.migration_max_seq_len();
            // Listen for models registering themselves, add them to gRPC service
            let namespace_filter = NamespaceFilter::from_namespace_and_prefix(
                model.namespace(),
                model.namespace_prefix(),
            );
            let local_model_path =
                (!model.path().as_os_str().is_empty()).then(|| model.path().to_path_buf());
            run_watcher(
                distributed_runtime.clone(),
                grpc_service.state().manager_clone(),
                router_config.clone(),
                migration_limit,
                migration_max_seq_len,
                namespace_filter,
                prefill_load_estimator.clone(),
                local_model_path,
                model.metrics_prefix(),
                model.runtime_config().tokenizer_backend,
                model.runtime_config().tokenizer_fallback_enabled,
            )
            .await?;
            grpc_service
        }
        EngineConfig::InProcessText { engine, model, .. } => {
            let grpc_service = grpc_service_builder.build()?;
            let engine = Arc::new(StreamingEngineAdapter::new(engine));
            let manager = grpc_service.model_manager();
            let checksum = model.card().mdcsum();
            manager.add_completions_model(model.service_name(), checksum, engine.clone())?;
            manager.add_chat_completions_model(model.service_name(), checksum, engine)?;
            grpc_service
        }
        EngineConfig::InProcessTokens {
            engine: inner_engine,
            model,
            ..
        } => {
            let grpc_service = grpc_service_builder.build()?;
            let manager = grpc_service.model_manager();
            let checksum = model.card().mdcsum();

            let tokenizer = model.card().tokenizer()?;
            let chat_pipeline = common::build_pipeline::<
                NvCreateChatCompletionRequest,
                NvCreateChatCompletionStreamResponse,
            >(model.card(), inner_engine.clone(), tokenizer.clone())
            .await?;
            manager.add_chat_completions_model(model.service_name(), checksum, chat_pipeline)?;

            let cmpl_pipeline = common::build_pipeline::<
                NvCreateCompletionRequest,
                NvCreateCompletionResponse,
            >(model.card(), inner_engine, tokenizer)
            .await?;
            manager.add_completions_model(model.service_name(), checksum, cmpl_pipeline)?;
            grpc_service
        }
    };

    // Run both HTTP (for metrics) and gRPC servers concurrently
    let http_service = grpc_service.http_service().clone();
    let shutdown_token = distributed_runtime.primary_token();

    // Wait for both servers to complete, propagating the first error if any occurs
    // Both tasks should run indefinitely until cancelled by the shutdown token
    let join_result = tokio::try_join!(
        grpc_service.run(shutdown_token.clone()),
        http_service.run(shutdown_token)
    );

    // Initiate runtime shutdown if either server exits, including bind
    // failures, for both discovery-backed and in-process engines.
    distributed_runtime.shutdown();

    join_result?;
    Ok(())
}

/// Spawns a task that watches for new models in store,
/// and registers them with the ModelManager so that the HTTP service can use them.
#[allow(clippy::too_many_arguments)]
async fn run_watcher(
    runtime: DistributedRuntime,
    model_manager: Arc<ModelManager>,
    router_config: RouterConfig,
    migration_limit: u32,
    migration_max_seq_len: Option<u32>,
    namespace_filter: NamespaceFilter,
    prefill_load_estimator: Option<Arc<dyn dynamo_kv_router::PrefillLoadEstimator>>,
    local_model_path: Option<PathBuf>,
    metrics_prefix: Option<String>,
    tokenizer_backend: Option<TokenizerBackend>,
    tokenizer_fallback_enabled: Option<bool>,
) -> anyhow::Result<()> {
    // Start the LoRA allocation controller when LoRA serving is enabled (mirrors http.rs;
    // additionally gated on DYN_LORA_ALLOCATION_ENABLED inside start_lora_controller). Without
    // this, gRPC LoRA deployments would never get dynamic placement (N3).
    if crate::lora::lora_serving_enabled() {
        let _controller_handle = model_manager.start_lora_controller(runtime.primary_token());
    }

    // Create metrics for migration tracking (not exposed via /metrics in gRPC mode)
    let metrics = Arc::new(Metrics::new_with_prefix(metrics_prefix));
    let mut watch_obj = ModelWatcher::new(
        runtime.clone(),
        model_manager,
        router_config,
        migration_limit,
        migration_max_seq_len,
        None,
        prefill_load_estimator,
        metrics,
    );
    watch_obj.set_local_model_path(local_model_path);
    watch_obj.set_tokenizer_backend(tokenizer_backend);
    watch_obj.set_tokenizer_fallback_enabled(tokenizer_fallback_enabled);
    tracing::debug!("Waiting for remote model");
    let discovery = runtime.discovery();
    let discovery_stream = discovery
        .list_and_watch(
            dynamo_runtime::discovery::DiscoveryQuery::AllModels,
            Some(runtime.primary_token()),
        )
        .await?;

    // [gluo NOTE] This is different from http::run_watcher where it alters the HTTP service
    // endpoint being exposed, gRPC doesn't have the same concept as the KServe service
    // only has one kind of inference endpoint.

    // Pass the discovery stream to the watcher
    let watch_obj = Arc::new(watch_obj);
    let _watcher_task = tokio::spawn(async move {
        watch_obj.watch(discovery_stream, namespace_filter).await;
    });

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{engines::make_echo_engine, local_model::LocalModelBuilder};
    use dynamo_runtime::{Runtime, distributed::DistributedConfig};
    use std::time::Duration;

    #[tokio::test]
    async fn metrics_bind_failure_shuts_down_dynamic_and_in_process_runtimes() {
        for dynamic in [true, false] {
            let occupied = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let model = Box::new(
                LocalModelBuilder::default()
                    .model_name(Some("bind-failure".to_string()))
                    .http_port(0)
                    .http_metrics_port(Some(occupied.local_addr().unwrap().port()))
                    .build()
                    .await
                    .unwrap(),
            );
            let engine_config = if dynamic {
                EngineConfig::Dynamic {
                    model,
                    chat_engine_factory: None,
                    prefill_load_estimator: None,
                }
            } else {
                EngineConfig::InProcessText {
                    engine: make_echo_engine(),
                    model,
                }
            };
            let drt = DistributedRuntime::new(
                Runtime::from_current().unwrap(),
                DistributedConfig::process_local(),
            )
            .await
            .unwrap();
            let shutdown = drt.primary_token();

            let error = tokio::time::timeout(Duration::from_secs(5), run(drt, engine_config))
                .await
                .expect("gRPC run must return after a metrics-port bind failure")
                .expect_err("the occupied metrics port must prevent server startup");
            assert!(
                error.to_string().contains("already in use"),
                "expected a metrics bind error (dynamic={dynamic}), got {error:#}"
            );
            tokio::time::timeout(Duration::from_secs(5), shutdown.cancelled())
                .await
                .expect("metrics bind failure must initiate runtime shutdown");
        }
    }
}