pd_router.rs 8.54 KB
Newer Older
1
2
// PD (Prefill-Decode) gRPC Router Implementation

3
use crate::config::types::RetryConfig;
4
use crate::core::{ConnectionMode, WorkerRegistry, WorkerType};
5
use crate::policies::PolicyRegistry;
6
7
8
9
10
11
use crate::protocols::chat::ChatCompletionRequest;
use crate::protocols::completion::CompletionRequest;
use crate::protocols::embedding::EmbeddingRequest;
use crate::protocols::generate::GenerateRequest;
use crate::protocols::rerank::RerankRequest;
use crate::protocols::responses::{ResponsesGetParams, ResponsesRequest};
12
use crate::reasoning_parser::ParserFactory as ReasoningParserFactory;
13
use crate::routers::RouterTrait;
14
use crate::server::AppContext;
15
use crate::tokenizer::traits::Tokenizer;
16
use crate::tool_parser::ParserFactory as ToolParserFactory;
17
18
19
20
use async_trait::async_trait;
use axum::{
    body::Body,
    extract::Request,
21
    http::{HeaderMap, StatusCode},
22
23
    response::{IntoResponse, Response},
};
24
use std::sync::Arc;
25
26

use tracing::debug;
27

28
29
30
use super::context::SharedComponents;
use super::pipeline::RequestPipeline;

31
/// gRPC PD (Prefill-Decode) router implementation for SGLang
32
#[derive(Clone)]
33
34
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcPDRouter {
35
36
    worker_registry: Arc<WorkerRegistry>,
    policy_registry: Arc<PolicyRegistry>,
37
    tokenizer: Arc<dyn Tokenizer>,
38
    reasoning_parser_factory: ReasoningParserFactory,
39
    tool_parser_factory: ToolParserFactory,
40
41
42
    dp_aware: bool,
    api_key: Option<String>,
    retry_config: RetryConfig,
43
44
    configured_reasoning_parser: Option<String>,
    configured_tool_parser: Option<String>,
45
46
    pipeline: RequestPipeline,
    shared_components: Arc<SharedComponents>,
47
}
48
49

impl GrpcPDRouter {
50
    /// Create a new gRPC PD router
51
    pub async fn new(ctx: &Arc<AppContext>) -> Result<Self, String> {
52
53
54
55
        // Get registries from context
        let worker_registry = ctx.worker_registry.clone();
        let policy_registry = ctx.policy_registry.clone();

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 PD router requires tokenizer".to_string())?
            .clone();
        let reasoning_parser_factory = ctx
            .reasoning_parser_factory
            .as_ref()
            .ok_or_else(|| "gRPC PD 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 PD router requires tool parser factory".to_string())?
            .clone();
72

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

80
81
82
83
        // Create PD pipeline
        let pipeline = RequestPipeline::new_pd(
            worker_registry.clone(),
            policy_registry.clone(),
84
85
86
87
88
89
90
            tokenizer.clone(),
            tool_parser_factory.clone(),
            reasoning_parser_factory.clone(),
            ctx.configured_tool_parser.clone(),
            ctx.configured_reasoning_parser.clone(),
        );

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

    /// Main route_generate implementation with PD dual dispatch
    async fn route_generate_impl(
        &self,
110
        headers: Option<&HeaderMap>,
111
112
113
114
115
116
117
118
        body: &GenerateRequest,
        model_id: Option<&str>,
    ) -> Response {
        debug!(
            "Processing generate request for model: {:?} (PD mode)",
            model_id
        );

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

    /// Main route_chat implementation with PD dual dispatch
    async fn route_chat_impl(
        &self,
133
        headers: Option<&HeaderMap>,
134
135
136
137
138
139
140
141
        body: &ChatCompletionRequest,
        model_id: Option<&str>,
    ) -> Response {
        debug!(
            "Processing chat completion request for model: {:?} (PD mode)",
            model_id
        );

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

impl std::fmt::Debug for GrpcPDRouter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156
157
158
159
160
        let prefill_workers = self.worker_registry.get_workers_filtered(
            None,
            Some(WorkerType::Prefill {
                bootstrap_port: None,
            }),
161
            Some(ConnectionMode::Grpc { port: None }),
162
163
164
165
166
            false,
        );
        let decode_workers = self.worker_registry.get_workers_filtered(
            None,
            Some(WorkerType::Decode),
167
            Some(ConnectionMode::Grpc { port: None }),
168
169
            false,
        );
170
        f.debug_struct("GrpcPDRouter")
171
172
            .field("prefill_workers_count", &prefill_workers.len())
            .field("decode_workers_count", &decode_workers.len())
173
174
            .field("dp_aware", &self.dp_aware)
            .finish()
175
176
177
178
179
180
181
182
183
184
    }
}

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

    async fn health_generate(&self, _req: Request<Body>) -> Response {
185
186
187
188
189
190
        // TODO: Implement actual generation test for gRPC PD mode
        (
            StatusCode::NOT_IMPLEMENTED,
            "Health generate not yet implemented for gRPC PD",
        )
            .into_response()
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
    }

    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,
207
208
209
        headers: Option<&HeaderMap>,
        body: &GenerateRequest,
        model_id: Option<&str>,
210
    ) -> Response {
211
        self.route_generate_impl(headers, body, model_id).await
212
213
214
215
    }

    async fn route_chat(
        &self,
216
217
218
        headers: Option<&HeaderMap>,
        body: &ChatCompletionRequest,
        model_id: Option<&str>,
219
    ) -> Response {
220
        self.route_chat_impl(headers, body, model_id).await
221
222
223
224
225
    }

    async fn route_completion(
        &self,
        _headers: Option<&HeaderMap>,
226
        _body: &CompletionRequest,
227
        _model_id: Option<&str>,
228
229
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
230
231
232
233
234
    }

    async fn route_responses(
        &self,
        _headers: Option<&HeaderMap>,
235
        _body: &ResponsesRequest,
236
        _model_id: Option<&str>,
237
238
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
239
240
    }

241
242
243
244
    async fn get_response(
        &self,
        _headers: Option<&HeaderMap>,
        _response_id: &str,
245
        _params: &ResponsesGetParams,
246
    ) -> Response {
247
248
249
250
251
252
253
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

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

254
255
256
    async fn route_embeddings(
        &self,
        _headers: Option<&HeaderMap>,
257
        _body: &EmbeddingRequest,
258
259
        _model_id: Option<&str>,
    ) -> Response {
260
261
262
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

263
264
265
    async fn route_rerank(
        &self,
        _headers: Option<&HeaderMap>,
266
        _body: &RerankRequest,
267
        _model_id: Option<&str>,
268
    ) -> Response {
269
270
271
272
273
274
275
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

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