dynamo_llm/http/service/
service_v2.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::metrics;
17use super::ModelManager;
18use anyhow::Result;
19use derive_builder::Builder;
20use tokio::task::JoinHandle;
21use tokio_util::sync::CancellationToken;
22
23#[derive(Clone)]
24pub struct HttpService {
25    models: ModelManager,
26    router: axum::Router,
27    port: u16,
28    host: String,
29}
30
31#[derive(Clone, Builder)]
32#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
33pub struct HttpServiceConfig {
34    #[builder(default = "8787")]
35    port: u16,
36
37    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
38    host: String,
39
40    // #[builder(default)]
41    // custom: Vec<axum::Router>
42    #[builder(default = "true")]
43    enable_chat_endpoints: bool,
44
45    #[builder(default = "true")]
46    enable_cmpl_endpoints: bool,
47}
48
49impl HttpService {
50    pub fn builder() -> HttpServiceConfigBuilder {
51        HttpServiceConfigBuilder::default()
52    }
53
54    pub fn model_manager(&self) -> &ModelManager {
55        &self.models
56    }
57
58    pub async fn spawn(&self, cancel_token: CancellationToken) -> JoinHandle<Result<()>> {
59        let this = self.clone();
60        tokio::spawn(async move { this.run(cancel_token).await })
61    }
62
63    pub async fn run(&self, cancel_token: CancellationToken) -> Result<()> {
64        let address = format!("{}:{}", self.host, self.port);
65        tracing::info!(address, "Starting HTTP service on: {address}");
66
67        let listener = tokio::net::TcpListener::bind(address.as_str())
68            .await
69            .unwrap_or_else(|_| panic!("could not bind to address: {address}"));
70
71        let router = self.router.clone();
72        let observer = cancel_token.child_token();
73
74        axum::serve(listener, router)
75            .with_graceful_shutdown(observer.cancelled_owned())
76            .await
77            .inspect_err(|_| cancel_token.cancel())?;
78
79        Ok(())
80    }
81}
82
83impl HttpServiceConfigBuilder {
84    pub fn build(self) -> Result<HttpService, anyhow::Error> {
85        let config = self.build_internal()?;
86
87        let model_manager = ModelManager::new();
88
89        // enable prometheus metrics
90        let registry = metrics::Registry::new();
91        model_manager.metrics().register(&registry)?;
92
93        let mut router = axum::Router::new();
94        let mut all_docs = Vec::new();
95
96        let mut routes = vec![
97            metrics::router(registry, None),
98            super::openai::list_models_router(model_manager.state(), None),
99        ];
100
101        if config.enable_chat_endpoints {
102            routes.push(super::openai::chat_completions_router(
103                model_manager.state(),
104                None,
105            ));
106        }
107
108        if config.enable_cmpl_endpoints {
109            routes.push(super::openai::completions_router(
110                model_manager.state(),
111                None,
112            ));
113        }
114
115        // for (route_docs, route) in routes.into_iter().chain(self.routes.into_iter()) {
116        //     router = router.merge(route);
117        //     all_docs.extend(route_docs);
118        // }
119
120        for (route_docs, route) in routes.into_iter() {
121            router = router.merge(route);
122            all_docs.extend(route_docs);
123        }
124
125        Ok(HttpService {
126            models: model_manager,
127            router,
128            port: config.port,
129            host: config.host,
130        })
131    }
132}
dynamo_llm/http/service/service_v2.rs

dynamo_llm/http/service/
service_v2.rs
