router.rs 7.9 KB
Newer Older
1
2
// gRPC Router Implementation

3
4
5
6
7
8
use std::sync::Arc;

use async_trait::async_trait;
use axum::{
    body::Body,
    extract::Request,
9
    http::{HeaderMap, StatusCode},
10
11
    response::{IntoResponse, Response},
};
12
use tracing::debug;
13

14
15
16
17
18
19
20
use super::{context::SharedComponents, pipeline::RequestPipeline};
use crate::{
    config::types::RetryConfig,
    core::WorkerRegistry,
    policies::PolicyRegistry,
    protocols::{
        chat::ChatCompletionRequest,
21
        classify::ClassifyRequest,
22
23
24
25
26
27
28
29
30
31
32
33
        completion::CompletionRequest,
        embedding::EmbeddingRequest,
        generate::GenerateRequest,
        rerank::RerankRequest,
        responses::{ResponsesGetParams, ResponsesRequest},
    },
    reasoning_parser::ParserFactory as ReasoningParserFactory,
    routers::RouterTrait,
    server::AppContext,
    tokenizer::traits::Tokenizer,
    tool_parser::ParserFactory as ToolParserFactory,
};
34

35
/// gRPC router implementation for SGLang
36
#[derive(Clone)]
37
#[allow(dead_code)]
38
pub struct GrpcRouter {
39
40
    worker_registry: Arc<WorkerRegistry>,
    policy_registry: Arc<PolicyRegistry>,
41
    tokenizer: Arc<dyn Tokenizer>,
42
    reasoning_parser_factory: ReasoningParserFactory,
43
    tool_parser_factory: ToolParserFactory,
44
45
46
    dp_aware: bool,
    api_key: Option<String>,
    retry_config: RetryConfig,
47
48
    configured_reasoning_parser: Option<String>,
    configured_tool_parser: Option<String>,
49
50
    pipeline: RequestPipeline,
    shared_components: Arc<SharedComponents>,
51
}
52
53

impl GrpcRouter {
54
    /// Create a new gRPC router
55
    pub async fn new(ctx: &Arc<AppContext>) -> Result<Self, String> {
56
57
58
59
60
61
62
63
64
65
66
        // Extract necessary components from context
        let tokenizer = ctx
            .tokenizer
            .as_ref()
            .ok_or_else(|| "gRPC router requires tokenizer".to_string())?
            .clone();
        let reasoning_parser_factory = ctx
            .reasoning_parser_factory
            .as_ref()
            .ok_or_else(|| "gRPC router requires reasoning parser factory".to_string())?
            .clone();
67
68
69
70
71
        let tool_parser_factory = ctx
            .tool_parser_factory
            .as_ref()
            .ok_or_else(|| "gRPC router requires tool parser factory".to_string())?
            .clone();
72

73
74
        let worker_registry = ctx.worker_registry.clone();
        let policy_registry = ctx.policy_registry.clone();
Chang Su's avatar
Chang Su committed
75

76
        // Create shared components for pipeline
77
        let shared_components = Arc::new(SharedComponents {
78
79
80
81
82
            tokenizer: tokenizer.clone(),
            tool_parser_factory: tool_parser_factory.clone(),
            reasoning_parser_factory: reasoning_parser_factory.clone(),
        });

83
84
85
86
        // Create pipeline
        let pipeline = RequestPipeline::new_regular(
            worker_registry.clone(),
            policy_registry.clone(),
87
88
89
90
91
92
93
            tokenizer.clone(),
            tool_parser_factory.clone(),
            reasoning_parser_factory.clone(),
            ctx.configured_tool_parser.clone(),
            ctx.configured_reasoning_parser.clone(),
        );

94
        Ok(GrpcRouter {
95
96
            worker_registry,
            policy_registry,
97
98
            tokenizer,
            reasoning_parser_factory,
99
            tool_parser_factory,
100
101
102
            dp_aware: ctx.router_config.dp_aware,
            api_key: ctx.router_config.api_key.clone(),
            retry_config: ctx.router_config.effective_retry_config(),
103
104
            configured_reasoning_parser: ctx.configured_reasoning_parser.clone(),
            configured_tool_parser: ctx.configured_tool_parser.clone(),
105
106
            pipeline,
            shared_components,
107
108
        })
    }
109
110
111
112

    /// Main route_chat implementation
    async fn route_chat_impl(
        &self,
113
        headers: Option<&HeaderMap>,
114
115
116
117
118
119
120
121
        body: &ChatCompletionRequest,
        model_id: Option<&str>,
    ) -> Response {
        debug!(
            "Processing chat completion request for model: {:?}",
            model_id
        );

122
123
124
        // Use pipeline for ALL requests (streaming and non-streaming)
        self.pipeline
            .execute_chat(
125
                Arc::new(body.clone()),
126
127
128
129
130
                headers.cloned(),
                model_id.map(|s| s.to_string()),
                self.shared_components.clone(),
            )
            .await
131
132
    }

133
134
135
    /// Main route_generate implementation
    async fn route_generate_impl(
        &self,
136
        headers: Option<&HeaderMap>,
137
138
139
140
141
        body: &GenerateRequest,
        model_id: Option<&str>,
    ) -> Response {
        debug!("Processing generate request for model: {:?}", model_id);

142
143
144
        // Use pipeline for ALL requests (streaming and non-streaming)
        self.pipeline
            .execute_generate(
145
                Arc::new(body.clone()),
146
147
148
                headers.cloned(),
                model_id.map(|s| s.to_string()),
                self.shared_components.clone(),
149
            )
150
            .await
151
    }
152
153
154
155
}

impl std::fmt::Debug for GrpcRouter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156
        let stats = self.worker_registry.stats();
157
        f.debug_struct("GrpcRouter")
158
            .field("workers_count", &stats.total_workers)
159
160
            .field("dp_aware", &self.dp_aware)
            .finish()
161
162
163
164
165
166
167
168
169
170
    }
}

#[async_trait]
impl RouterTrait for GrpcRouter {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    async fn health_generate(&self, _req: Request<Body>) -> Response {
171
172
173
174
175
176
        // TODO: Implement actual generation test for gRPC
        (
            StatusCode::NOT_IMPLEMENTED,
            "Health generate not yet implemented for gRPC",
        )
            .into_response()
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    }

    async fn get_server_info(&self, _req: Request<Body>) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn get_models(&self, _req: Request<Body>) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn get_model_info(&self, _req: Request<Body>) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn route_generate(
        &self,
193
194
195
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
        model_id: Option<&str>,
196
    ) -> Response {
197
        self.route_generate_impl(headers, body, model_id).await
198
199
200
201
    }

    async fn route_chat(
        &self,
202
        headers: Option<&HeaderMap>,
203
        body: &ChatCompletionRequest,
204
        model_id: Option<&str>,
205
    ) -> Response {
206
        self.route_chat_impl(headers, body, model_id).await
207
208
209
210
211
    }

    async fn route_completion(
        &self,
        _headers: Option<&HeaderMap>,
212
        _body: &CompletionRequest,
213
        _model_id: Option<&str>,
214
215
216
217
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

218
219
220
    async fn route_responses(
        &self,
        _headers: Option<&HeaderMap>,
221
        _body: &ResponsesRequest,
222
        _model_id: Option<&str>,
223
224
225
226
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

227
228
229
230
    async fn get_response(
        &self,
        _headers: Option<&HeaderMap>,
        _response_id: &str,
231
        _params: &ResponsesGetParams,
232
    ) -> Response {
233
234
235
236
237
238
239
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn cancel_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

240
241
242
243
244
245
246
247
248
    async fn route_classify(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &ClassifyRequest,
        _model_id: Option<&str>,
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

249
250
251
    async fn route_embeddings(
        &self,
        _headers: Option<&HeaderMap>,
252
        _body: &EmbeddingRequest,
253
254
        _model_id: Option<&str>,
    ) -> Response {
255
256
257
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

258
259
260
    async fn route_rerank(
        &self,
        _headers: Option<&HeaderMap>,
261
        _body: &RerankRequest,
262
        _model_id: Option<&str>,
263
    ) -> Response {
264
265
266
267
268
269
270
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    fn router_type(&self) -> &'static str {
        "grpc"
    }
}