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

3
use crate::config::types::RetryConfig;
4
use crate::core::{
5
    BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, WorkerRegistry, WorkerType,
6
7
};
use crate::metrics::RouterMetrics;
8
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
9
use crate::reasoning_parser::ParserFactory;
10
use crate::routers::RouterTrait;
11
use crate::tokenizer::traits::Tokenizer;
12
use crate::tool_parser::ParserRegistry;
13
14
15
16
17
18
19
use async_trait::async_trait;
use axum::{
    body::Body,
    extract::Request,
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Response},
};
20
use std::sync::Arc;
21
use std::time::Duration;
22
use tracing::info;
23

24
25
26
/// gRPC PD (Prefill-Decode) router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
pub struct GrpcPDRouter {
27
28
29
30
    /// Centralized worker registry
    worker_registry: Arc<WorkerRegistry>,
    /// Centralized policy registry
    policy_registry: Arc<PolicyRegistry>,
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    /// Load balancing policy for prefill
    prefill_policy: Arc<dyn LoadBalancingPolicy>,
    /// Load balancing policy for decode
    decode_policy: Arc<dyn LoadBalancingPolicy>,
    /// Tokenizer for handling text encoding/decoding
    tokenizer: Arc<dyn Tokenizer>,
    /// Reasoning parser factory for structured reasoning outputs
    reasoning_parser_factory: ParserFactory,
    /// Tool parser registry for function/tool calls
    tool_parser_registry: &'static ParserRegistry,
    /// Configuration
    timeout_secs: u64,
    interval_secs: u64,
    dp_aware: bool,
    api_key: Option<String>,
    retry_config: RetryConfig,
    circuit_breaker_config: CircuitBreakerConfig,
}
49
50

impl GrpcPDRouter {
51
52
53
54
55
56
    /// Create a new gRPC PD router
    pub async fn new(
        prefill_urls: Vec<(String, Option<u16>)>,
        decode_urls: Vec<String>,
        prefill_policy: Arc<dyn LoadBalancingPolicy>,
        decode_policy: Arc<dyn LoadBalancingPolicy>,
57
        ctx: &Arc<crate::server::AppContext>,
58
    ) -> Result<Self, String> {
59
60
61
62
        // Get registries from context
        let worker_registry = ctx.worker_registry.clone();
        let policy_registry = ctx.policy_registry.clone();

63
64
65
        // Update metrics
        RouterMetrics::set_active_workers(prefill_urls.len() + decode_urls.len());

66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
        // 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();
        let tool_parser_registry = ctx
            .tool_parser_registry
            .ok_or_else(|| "gRPC PD router requires tool parser registry".to_string())?;

81
        // Convert config CircuitBreakerConfig to core CircuitBreakerConfig
82
        let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
83
84
85
86
87
88
89
        let core_cb_config = CircuitBreakerConfig {
            failure_threshold: circuit_breaker_config.failure_threshold,
            success_threshold: circuit_breaker_config.success_threshold,
            timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
            window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
        };

90
        for (url, bootstrap_port) in &prefill_urls {
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
            let worker = BasicWorkerBuilder::new(url.clone())
                .worker_type(WorkerType::Prefill {
                    bootstrap_port: *bootstrap_port,
                })
                .connection_mode(crate::core::ConnectionMode::Grpc {
                    port: *bootstrap_port,
                })
                .circuit_breaker_config(core_cb_config.clone())
                .health_config(HealthConfig {
                    timeout_secs: ctx.router_config.health_check.timeout_secs,
                    check_interval_secs: ctx.router_config.health_check.check_interval_secs,
                    endpoint: ctx.router_config.health_check.endpoint.clone(),
                    failure_threshold: ctx.router_config.health_check.failure_threshold,
                    success_threshold: ctx.router_config.health_check.success_threshold,
                })
                // No longer passing pre-initialized client - will be created lazily
                .build();
108

109
110
111
112
113
            worker_registry.register(Arc::new(worker));
            info!(
                "Registered gRPC prefill worker at {} (will connect on first use)",
                url
            );
114
115
116
        }

        for url in &decode_urls {
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
            let worker = BasicWorkerBuilder::new(url.clone())
                .worker_type(WorkerType::Decode)
                .connection_mode(crate::core::ConnectionMode::Grpc { port: None })
                .circuit_breaker_config(core_cb_config.clone())
                .health_config(HealthConfig {
                    timeout_secs: ctx.router_config.health_check.timeout_secs,
                    check_interval_secs: ctx.router_config.health_check.check_interval_secs,
                    endpoint: ctx.router_config.health_check.endpoint.clone(),
                    failure_threshold: ctx.router_config.health_check.failure_threshold,
                    success_threshold: ctx.router_config.health_check.success_threshold,
                })
                .build();

            worker_registry.register(Arc::new(worker));
            info!(
                "Registered gRPC decode worker at {} (will connect on first use)",
                url
            );
        }
136

137
138
        if prefill_urls.is_empty() && decode_urls.is_empty() {
            return Err("No gRPC workers configured".to_string());
139
140
141
142
143
144
145
146
147
148
149
        }

        // Initialize policies with workers if needed - filter for gRPC workers only
        let prefill_workers = worker_registry.get_workers_filtered(
            None, // any model
            Some(WorkerType::Prefill {
                bootstrap_port: None,
            }),
            Some(crate::core::ConnectionMode::Grpc { port: None }),
            false, // include unhealthy workers during initialization
        );
150
151
152
153
154
155
156
        if let Some(cache_aware) = prefill_policy
            .as_any()
            .downcast_ref::<crate::policies::CacheAwarePolicy>()
        {
            cache_aware.init_workers(&prefill_workers);
        }

157
158
159
160
161
162
        let decode_workers = worker_registry.get_workers_filtered(
            None, // any model
            Some(WorkerType::Decode),
            Some(crate::core::ConnectionMode::Grpc { port: None }),
            false, // include unhealthy workers during initialization
        );
163
164
165
166
167
168
169
        if let Some(cache_aware) = decode_policy
            .as_any()
            .downcast_ref::<crate::policies::CacheAwarePolicy>()
        {
            cache_aware.init_workers(&decode_workers);
        }

170
        // No need for local health checkers - WorkerRegistry handles health checking
171
172

        Ok(GrpcPDRouter {
173
174
            worker_registry,
            policy_registry,
175
176
177
178
179
            prefill_policy,
            decode_policy,
            tokenizer,
            reasoning_parser_factory,
            tool_parser_registry,
180
181
182
183
184
            timeout_secs: ctx.router_config.worker_startup_timeout_secs,
            interval_secs: ctx.router_config.worker_startup_check_interval_secs,
            dp_aware: ctx.router_config.dp_aware,
            api_key: ctx.router_config.api_key.clone(),
            retry_config: ctx.router_config.effective_retry_config(),
185
186
187
188
189
190
191
            circuit_breaker_config: core_cb_config,
        })
    }
}

impl std::fmt::Debug for GrpcPDRouter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        let prefill_workers = self.worker_registry.get_workers_filtered(
            None,
            Some(WorkerType::Prefill {
                bootstrap_port: None,
            }),
            Some(crate::core::ConnectionMode::Grpc { port: None }),
            false,
        );
        let decode_workers = self.worker_registry.get_workers_filtered(
            None,
            Some(WorkerType::Decode),
            Some(crate::core::ConnectionMode::Grpc { port: None }),
            false,
        );
206
        f.debug_struct("GrpcPDRouter")
207
208
            .field("prefill_workers_count", &prefill_workers.len())
            .field("decode_workers_count", &decode_workers.len())
209
210
211
212
            .field("timeout_secs", &self.timeout_secs)
            .field("interval_secs", &self.interval_secs)
            .field("dp_aware", &self.dp_aware)
            .finish()
213
214
215
216
217
218
219
220
221
222
    }
}

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

    async fn health_generate(&self, _req: Request<Body>) -> Response {
223
224
225
226
227
228
        // TODO: Implement actual generation test for gRPC PD mode
        (
            StatusCode::NOT_IMPLEMENTED,
            "Health generate not yet implemented for gRPC PD",
        )
            .into_response()
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    }

    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,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::GenerateRequest,
247
        _model_id: Option<&str>,
248
249
250
251
252
253
254
255
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn route_chat(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::ChatCompletionRequest,
256
        _model_id: Option<&str>,
257
258
259
260
261
262
263
264
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

    async fn route_completion(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::CompletionRequest,
265
        _model_id: Option<&str>,
266
267
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
268
269
270
271
272
273
    }

    async fn route_responses(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::ResponsesRequest,
274
        _model_id: Option<&str>,
275
276
    ) -> Response {
        (StatusCode::NOT_IMPLEMENTED).into_response()
277
278
    }

279
280
281
282
283
284
    async fn get_response(
        &self,
        _headers: Option<&HeaderMap>,
        _response_id: &str,
        _params: &crate::protocols::spec::ResponsesGetParams,
    ) -> Response {
285
286
287
288
289
290
291
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

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

292
293
294
295
296
297
    async fn route_embeddings(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::EmbeddingRequest,
        _model_id: Option<&str>,
    ) -> Response {
298
299
300
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

301
302
303
304
    async fn route_rerank(
        &self,
        _headers: Option<&HeaderMap>,
        _body: &crate::protocols::spec::RerankRequest,
305
        _model_id: Option<&str>,
306
    ) -> Response {
307
308
309
310
311
312
313
        (StatusCode::NOT_IMPLEMENTED).into_response()
    }

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