service_v2.rs 4.73 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// 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::metrics;
use super::ModelManager;
18
use super::RouteDoc;
19
use crate::request_template::RequestTemplate;
20
use anyhow::Result;
21
use derive_builder::Builder;
22
use tokio::task::JoinHandle;
23
24
25
26
27
28
29
use tokio_util::sync::CancellationToken;

#[derive(Clone)]
pub struct HttpService {
    models: ModelManager,
    router: axum::Router,
    port: u16,
30
    host: String,
31
    route_docs: Vec<RouteDoc>,
32
33
34
}

#[derive(Clone, Builder)]
35
#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
36
37
38
39
pub struct HttpServiceConfig {
    #[builder(default = "8787")]
    port: u16,

40
41
42
    #[builder(setter(into), default = "String::from(\"0.0.0.0\")")]
    host: String,

43
44
45
46
47
48
49
    // #[builder(default)]
    // custom: Vec<axum::Router>
    #[builder(default = "true")]
    enable_chat_endpoints: bool,

    #[builder(default = "true")]
    enable_cmpl_endpoints: bool,
50

51
52
53
    #[builder(default = "false")]
    enable_embeddings_endpoints: bool,

54
55
    #[builder(default = "None")]
    request_template: Option<RequestTemplate>,
56
57
58
59
60
61
62
63
64
65
66
}

impl HttpService {
    pub fn builder() -> HttpServiceConfigBuilder {
        HttpServiceConfigBuilder::default()
    }

    pub fn model_manager(&self) -> &ModelManager {
        &self.models
    }

67
68
69
70
71
72
    pub async fn spawn(&self, cancel_token: CancellationToken) -> JoinHandle<Result<()>> {
        let this = self.clone();
        tokio::spawn(async move { this.run(cancel_token).await })
    }

    pub async fn run(&self, cancel_token: CancellationToken) -> Result<()> {
73
        let address = format!("{}:{}", self.host, self.port);
74
75
76
77
78
79
80
81
82
        tracing::info!(address, "Starting HTTP service on: {address}");

        let listener = tokio::net::TcpListener::bind(address.as_str())
            .await
            .unwrap_or_else(|_| panic!("could not bind to address: {address}"));

        let router = self.router.clone();
        let observer = cancel_token.child_token();

83
        axum::serve(listener, router)
84
85
            .with_graceful_shutdown(observer.cancelled_owned())
            .await
86
87
88
            .inspect_err(|_| cancel_token.cancel())?;

        Ok(())
89
    }
90
91
92
93
94

    /// Documentation of exposed HTTP endpoints
    pub fn route_docs(&self) -> &[RouteDoc] {
        &self.route_docs
    }
95
96
97
98
}

impl HttpServiceConfigBuilder {
    pub fn build(self) -> Result<HttpService, anyhow::Error> {
99
        let config: HttpServiceConfig = self.build_internal()?;
100
101
102
103
104
105
106
107

        let model_manager = ModelManager::new();

        // enable prometheus metrics
        let registry = metrics::Registry::new();
        model_manager.metrics().register(&registry)?;

        let mut router = axum::Router::new();
108

109
110
111
112
113
114
115
116
        let mut all_docs = Vec::new();

        let mut routes = vec![
            metrics::router(registry, None),
            super::openai::list_models_router(model_manager.state(), None),
        ];

        if config.enable_chat_endpoints {
117
            routes.push(super::openai::chat_completions_router(
118
                model_manager.state(),
119
                config.request_template,
120
121
122
123
124
                None,
            ));
        }

        if config.enable_cmpl_endpoints {
125
            routes.push(super::openai::completions_router(
126
127
128
129
130
                model_manager.state(),
                None,
            ));
        }

131
132
133
134
135
136
137
        if config.enable_embeddings_endpoints {
            routes.push(super::openai::embeddings_router(
                model_manager.state(),
                None,
            ));
        }

138
139
140
141
142
143
144
145
146
147
148
149
150
151
        // for (route_docs, route) in routes.into_iter().chain(self.routes.into_iter()) {
        //     router = router.merge(route);
        //     all_docs.extend(route_docs);
        // }

        for (route_docs, route) in routes.into_iter() {
            router = router.merge(route);
            all_docs.extend(route_docs);
        }

        Ok(HttpService {
            models: model_manager,
            router,
            port: config.port,
152
            host: config.host,
153
            route_docs: all_docs,
154
155
        })
    }
156
157
158
159
160

    pub fn with_request_template(mut self, request_template: Option<RequestTemplate>) -> Self {
        self.request_template = Some(request_template);
        self
    }
161
}