request_processing.rs 21.8 KB
Newer Older
1
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
2
use serde_json::{from_str, to_string, to_value, to_vec};
3
4
use std::time::Instant;

5
use sglang_router_rs::core::{BasicWorker, BasicWorkerBuilder, Worker, WorkerType};
6
7
8
use sglang_router_rs::protocols::spec::{
    ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateParameters, GenerateRequest,
    SamplingParams, StringOrArray, UserMessageContent,
9
};
10
11
12
use sglang_router_rs::routers::http::pd_types::{
    generate_room_id, get_hostname, RequestWithBootstrap,
};
13
14

fn create_test_worker() -> BasicWorker {
15
16
    BasicWorkerBuilder::new("http://test-server:8000")
        .worker_type(WorkerType::Prefill {
17
            bootstrap_port: Some(5678),
18
19
        })
        .build()
20
}
21

22
23
24
25
// Helper function to get bootstrap info from worker
fn get_bootstrap_info(worker: &BasicWorker) -> (String, Option<u16>) {
    let hostname = get_hostname(worker.url());
    let bootstrap_port = match worker.worker_type() {
26
        WorkerType::Prefill { bootstrap_port } => bootstrap_port,
27
28
29
30
31
        _ => None,
    };
    (hostname, bootstrap_port)
}

32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
/// Create a default GenerateRequest for benchmarks with minimal fields set
fn default_generate_request() -> GenerateRequest {
    GenerateRequest {
        text: None,
        prompt: None,
        input_ids: None,
        stream: false,
        parameters: None,
        sampling_params: None,
        return_logprob: false,
        // SGLang Extensions
        lora_path: None,
        session_params: None,
        return_hidden_states: false,
        rid: None,
    }
}

/// Create a default ChatCompletionRequest for benchmarks with minimal fields set
51
#[allow(deprecated)]
52
53
fn default_chat_completion_request() -> ChatCompletionRequest {
    ChatCompletionRequest {
54
        // Required fields in OpenAI order
55
        messages: vec![],
56
57
58
59
        model: String::new(),

        // Use default for all other fields
        ..Default::default()
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    }
}

/// Create a default CompletionRequest for benchmarks with minimal fields set
fn default_completion_request() -> CompletionRequest {
    CompletionRequest {
        model: String::new(),
        prompt: StringOrArray::String(String::new()),
        suffix: None,
        max_tokens: None,
        temperature: None,
        top_p: None,
        n: None,
        stream: false,
        stream_options: None,
        logprobs: None,
        echo: false,
        stop: None,
        presence_penalty: None,
        frequency_penalty: None,
        best_of: None,
        logit_bias: None,
        user: None,
        seed: None,
        // SGLang Extensions
        top_k: None,
        min_p: None,
        min_tokens: None,
        repetition_penalty: None,
        regex: None,
        ebnf: None,
        json_schema: None,
        stop_token_ids: None,
        no_stop_trim: false,
        ignore_eos: false,
        skip_special_tokens: true,
        // SGLang Extensions
        lora_path: None,
        session_params: None,
        return_hidden_states: false,
100
        sampling_seed: None,
101
102
103
104
        other: serde_json::Map::new(),
    }
}

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// Sample request data for benchmarks
fn create_sample_generate_request() -> GenerateRequest {
    GenerateRequest {
        text: Some("Write a story about artificial intelligence".to_string()),
        parameters: Some(GenerateParameters {
            max_new_tokens: Some(100),
            temperature: Some(0.8),
            top_p: Some(0.9),
            top_k: Some(50),
            repetition_penalty: Some(1.0),
            ..Default::default()
        }),
        sampling_params: Some(SamplingParams {
            temperature: Some(0.8),
            top_p: Some(0.9),
            top_k: Some(50),
            frequency_penalty: Some(0.0),
            presence_penalty: Some(0.0),
            repetition_penalty: Some(1.0),
            ..Default::default()
        }),
126
        ..default_generate_request()
127
128
129
    }
}

130
#[allow(deprecated)]
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
fn create_sample_chat_completion_request() -> ChatCompletionRequest {
    ChatCompletionRequest {
        model: "gpt-3.5-turbo".to_string(),
        messages: vec![
            ChatMessage::System {
                role: "system".to_string(),
                content: "You are a helpful assistant".to_string(),
                name: None,
            },
            ChatMessage::User {
                role: "user".to_string(),
                content: UserMessageContent::Text(
                    "Explain quantum computing in simple terms".to_string(),
                ),
                name: None,
            },
        ],
        max_tokens: Some(150),
        max_completion_tokens: Some(150),
        temperature: Some(0.7),
        top_p: Some(1.0),
        n: Some(1),
        presence_penalty: Some(0.0),
        frequency_penalty: Some(0.0),
        parallel_tool_calls: Some(true),
156
        ..default_chat_completion_request()
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    }
}

fn create_sample_completion_request() -> CompletionRequest {
    CompletionRequest {
        model: "text-davinci-003".to_string(),
        prompt: StringOrArray::String("Complete this sentence: The future of AI is".to_string()),
        max_tokens: Some(50),
        temperature: Some(0.8),
        top_p: Some(1.0),
        n: Some(1),
        presence_penalty: Some(0.0),
        frequency_penalty: Some(0.0),
        best_of: Some(1),
171
        ..default_completion_request()
172
173
174
    }
}

175
#[allow(deprecated)]
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
fn create_large_chat_completion_request() -> ChatCompletionRequest {
    let mut messages = vec![ChatMessage::System {
        role: "system".to_string(),
        content: "You are a helpful assistant with extensive knowledge.".to_string(),
        name: None,
    }];

    // Add many user/assistant pairs to simulate a long conversation
    for i in 0..50 {
        messages.push(ChatMessage::User {
            role: "user".to_string(),
            content: UserMessageContent::Text(format!("Question {}: What do you think about topic number {} which involves complex reasoning about multiple interconnected systems and their relationships?", i, i)),
            name: None,
        });
        messages.push(ChatMessage::Assistant {
            role: "assistant".to_string(),
            content: Some(format!("Answer {}: This is a detailed response about topic {} that covers multiple aspects and provides comprehensive analysis of the interconnected systems you mentioned.", i, i)),
            name: None,
            tool_calls: None,
            function_call: None,
196
            reasoning_content: None,
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
        });
    }

    ChatCompletionRequest {
        model: "gpt-4".to_string(),
        messages,
        max_tokens: Some(1000),
        max_completion_tokens: Some(1000),
        temperature: Some(0.7),
        top_p: Some(0.95),
        n: Some(1),
        presence_penalty: Some(0.1),
        frequency_penalty: Some(0.1),
        top_logprobs: Some(5),
        seed: Some(42),
        parallel_tool_calls: Some(true),
213
        ..default_chat_completion_request()
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
    }
}

// Benchmark JSON serialization
fn bench_json_serialization(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_serialization");

    let generate_req = create_sample_generate_request();
    let chat_req = create_sample_chat_completion_request();
    let completion_req = create_sample_completion_request();
    let large_chat_req = create_large_chat_completion_request();

    group.bench_function("generate_request", |b| {
        b.iter(|| {
            let json = to_string(black_box(&generate_req)).unwrap();
            black_box(json);
        });
    });

    group.bench_function("chat_completion_request", |b| {
        b.iter(|| {
            let json = to_string(black_box(&chat_req)).unwrap();
            black_box(json);
        });
    });

    group.bench_function("completion_request", |b| {
        b.iter(|| {
            let json = to_string(black_box(&completion_req)).unwrap();
            black_box(json);
        });
    });

    group.bench_function("large_chat_completion_request", |b| {
        b.iter(|| {
            let json = to_string(black_box(&large_chat_req)).unwrap();
            black_box(json);
        });
    });

    group.bench_function("generate_request_to_bytes", |b| {
        b.iter(|| {
            let bytes = to_vec(black_box(&generate_req)).unwrap();
            black_box(bytes);
        });
    });

    group.finish();
}

// Benchmark JSON deserialization
fn bench_json_deserialization(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_deserialization");

    let generate_json = to_string(&create_sample_generate_request()).unwrap();
    let chat_json = to_string(&create_sample_chat_completion_request()).unwrap();
    let completion_json = to_string(&create_sample_completion_request()).unwrap();
    let large_chat_json = to_string(&create_large_chat_completion_request()).unwrap();

    group.bench_function("generate_request", |b| {
        b.iter(|| {
            let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
            black_box(req);
        });
    });

    group.bench_function("chat_completion_request", |b| {
        b.iter(|| {
            let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap();
            black_box(req);
        });
    });

    group.bench_function("completion_request", |b| {
        b.iter(|| {
            let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap();
            black_box(req);
        });
    });

    group.bench_function("large_chat_completion_request", |b| {
        b.iter(|| {
            let req: ChatCompletionRequest = from_str(black_box(&large_chat_json)).unwrap();
            black_box(req);
        });
    });

    group.finish();
}

304
305
306
// Benchmark bootstrap injection (replaces request adaptation)
fn bench_bootstrap_injection(c: &mut Criterion) {
    let mut group = c.benchmark_group("bootstrap_injection");
307
308
309
310
311

    let generate_req = create_sample_generate_request();
    let chat_req = create_sample_chat_completion_request();
    let completion_req = create_sample_completion_request();
    let large_chat_req = create_large_chat_completion_request();
312
    let worker = create_test_worker();
313
    let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
314

315
    group.bench_function("generate_bootstrap_injection", |b| {
316
        b.iter(|| {
317
318
319
320
321
322
323
            let request_with_bootstrap = RequestWithBootstrap {
                original: &generate_req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let json = to_value(black_box(&request_with_bootstrap)).unwrap();
324
            black_box(json);
325
326
327
        });
    });

328
    group.bench_function("chat_completion_bootstrap_injection", |b| {
329
        b.iter(|| {
330
331
332
333
334
335
336
            let request_with_bootstrap = RequestWithBootstrap {
                original: &chat_req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let json = to_value(black_box(&request_with_bootstrap)).unwrap();
337
            black_box(json);
338
339
340
        });
    });

341
    group.bench_function("completion_bootstrap_injection", |b| {
342
        b.iter(|| {
343
344
345
346
347
348
349
            let request_with_bootstrap = RequestWithBootstrap {
                original: &completion_req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let json = to_value(black_box(&request_with_bootstrap)).unwrap();
350
            black_box(json);
351
352
353
        });
    });

354
    group.bench_function("large_chat_completion_bootstrap_injection", |b| {
355
        b.iter(|| {
356
357
358
359
360
361
362
            let request_with_bootstrap = RequestWithBootstrap {
                original: &large_chat_req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let json = to_value(black_box(&request_with_bootstrap)).unwrap();
363
            black_box(json);
364
365
366
367
368
369
        });
    });

    group.finish();
}

370
371
372
// Benchmark direct JSON routing (replaces regular routing)
fn bench_direct_json_routing(c: &mut Criterion) {
    let mut group = c.benchmark_group("direct_json_routing");
373
374
375
376
377
378
379

    let generate_req = create_sample_generate_request();
    let chat_req = create_sample_chat_completion_request();
    let completion_req = create_sample_completion_request();

    group.bench_function("generate_to_json", |b| {
        b.iter(|| {
380
381
382
383
384
385
386
387
            let json = to_value(black_box(&generate_req)).unwrap();
            black_box(json);
        });
    });

    group.bench_function("generate_to_json_string", |b| {
        b.iter(|| {
            let json = to_string(black_box(&generate_req)).unwrap();
388
389
390
391
392
393
            black_box(json);
        });
    });

    group.bench_function("generate_to_bytes", |b| {
        b.iter(|| {
394
            let bytes = to_vec(black_box(&generate_req)).unwrap();
395
396
397
398
399
400
            black_box(bytes);
        });
    });

    group.bench_function("chat_completion_to_json", |b| {
        b.iter(|| {
401
            let json = to_value(black_box(&chat_req)).unwrap();
402
403
404
405
            black_box(json);
        });
    });

406
    group.bench_function("chat_completion_to_json_string", |b| {
407
        b.iter(|| {
408
409
            let json = to_string(black_box(&chat_req)).unwrap();
            black_box(json);
410
411
412
413
414
        });
    });

    group.bench_function("completion_to_json", |b| {
        b.iter(|| {
415
            let json = to_value(black_box(&completion_req)).unwrap();
416
417
418
419
420
421
422
423
424
425
426
427
428
429
            black_box(json);
        });
    });

    group.finish();
}

// Benchmark throughput with different request sizes
fn bench_throughput_by_size(c: &mut Criterion) {
    let mut group = c.benchmark_group("throughput_by_size");

    // Create requests of different sizes
    let small_generate = GenerateRequest {
        text: Some("Hi".to_string()),
430
        ..default_generate_request()
431
432
433
434
    };

    let medium_generate = GenerateRequest {
        text: Some("Write a medium length story about AI".repeat(10)),
435
        ..default_generate_request()
436
437
438
439
    };

    let large_generate = GenerateRequest {
        text: Some("Write a very long and detailed story about artificial intelligence and its impact on society".repeat(100)),
440
        ..default_generate_request()
441
442
    };

443
    let worker = create_test_worker();
444
    let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
445

446
447
448
449
450
451
452
    for (name, req) in [
        ("small", &small_generate),
        ("medium", &medium_generate),
        ("large", &large_generate),
    ] {
        let json = to_string(req).unwrap();
        let size_bytes = json.len();
453
        let hostname_clone = hostname.clone();
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473

        group.throughput(Throughput::Bytes(size_bytes as u64));
        group.bench_with_input(BenchmarkId::new("serialize", name), &req, |b, req| {
            b.iter(|| {
                let json = to_string(black_box(req)).unwrap();
                black_box(json);
            });
        });

        group.bench_with_input(
            BenchmarkId::new("deserialize", name),
            &json,
            |b, json_str| {
                b.iter(|| {
                    let req: GenerateRequest = black_box(from_str(json_str)).unwrap();
                    black_box(req);
                });
            },
        );

474
475
476
        group.bench_with_input(
            BenchmarkId::new("bootstrap_inject", name),
            &req,
477
478
            move |b, req| {
                let hostname = hostname_clone.clone();
479
                b.iter(|| {
480
481
482
483
484
485
486
                    let request_with_bootstrap = RequestWithBootstrap {
                        original: req,
                        bootstrap_host: hostname.clone(),
                        bootstrap_port,
                        bootstrap_room: generate_room_id(),
                    };
                    let json = to_value(&request_with_bootstrap).unwrap();
487
488
489
490
                    black_box(json);
                });
            },
        );
491
492
493
494
495
    }

    group.finish();
}

496
// Benchmark full round-trip: deserialize -> inject bootstrap -> serialize
497
498
499
500
501
502
fn bench_full_round_trip(c: &mut Criterion) {
    let mut group = c.benchmark_group("full_round_trip");

    let generate_json = to_string(&create_sample_generate_request()).unwrap();
    let chat_json = to_string(&create_sample_chat_completion_request()).unwrap();
    let completion_json = to_string(&create_sample_completion_request()).unwrap();
503
    let worker = create_test_worker();
504
    let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
505
506
507
508
509

    group.bench_function("generate_openai_to_pd_pipeline", |b| {
        b.iter(|| {
            // Deserialize OpenAI request
            let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
510
511
512
513
514
515
516
            // Create wrapper with bootstrap fields
            let request_with_bootstrap = RequestWithBootstrap {
                original: &req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
517
            // Serialize final request
518
            let pd_json = to_string(&request_with_bootstrap).unwrap();
519
520
521
522
523
524
525
            black_box(pd_json);
        });
    });

    group.bench_function("chat_completion_openai_to_pd_pipeline", |b| {
        b.iter(|| {
            let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap();
526
527
528
529
530
531
532
            let request_with_bootstrap = RequestWithBootstrap {
                original: &req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let pd_json = to_string(&request_with_bootstrap).unwrap();
533
534
535
536
537
538
539
            black_box(pd_json);
        });
    });

    group.bench_function("completion_openai_to_pd_pipeline", |b| {
        b.iter(|| {
            let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap();
540
541
542
543
544
545
546
            let request_with_bootstrap = RequestWithBootstrap {
                original: &req,
                bootstrap_host: hostname.clone(),
                bootstrap_port,
                bootstrap_room: generate_room_id(),
            };
            let pd_json = to_string(&request_with_bootstrap).unwrap();
547
548
549
550
            black_box(pd_json);
        });
    });

551
    group.bench_function("generate_direct_json_pipeline", |b| {
552
553
554
        b.iter(|| {
            // Deserialize OpenAI request
            let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap();
555
556
557
558
            // Convert to JSON for direct routing (no bootstrap injection)
            let routing_json = to_value(&req).unwrap();
            let json_string = to_string(&routing_json).unwrap();
            black_box(json_string);
559
560
561
562
563
564
565
566
567
568
569
570
571
572
        });
    });

    group.finish();
}

fn benchmark_summary(c: &mut Criterion) {
    let group = c.benchmark_group("benchmark_summary");

    println!("\nSGLang Router Performance Benchmark Suite");
    println!("=============================================");

    // Quick performance overview
    let generate_req = create_sample_generate_request();
573
    let worker = create_test_worker();
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596

    println!("\nQuick Performance Overview:");

    // Measure serialization
    let start = Instant::now();
    for _ in 0..1000 {
        let _ = black_box(to_string(&generate_req).unwrap());
    }
    let serialize_time = start.elapsed().as_nanos() / 1000;
    println!("  * Serialization (avg):     {:>8} ns/req", serialize_time);

    // Measure deserialization
    let json = to_string(&generate_req).unwrap();
    let start = Instant::now();
    for _ in 0..1000 {
        let _: GenerateRequest = black_box(from_str(&json).unwrap());
    }
    let deserialize_time = start.elapsed().as_nanos() / 1000;
    println!(
        "  * Deserialization (avg):   {:>8} ns/req",
        deserialize_time
    );

597
    // Measure bootstrap injection (replaces adaptation)
598
    let (hostname, bootstrap_port) = get_bootstrap_info(&worker);
599
600
    let start = Instant::now();
    for _ in 0..1000 {
601
602
603
604
605
606
607
        let request_with_bootstrap = RequestWithBootstrap {
            original: &generate_req,
            bootstrap_host: hostname.clone(),
            bootstrap_port,
            bootstrap_room: generate_room_id(),
        };
        let _ = black_box(to_value(&request_with_bootstrap).unwrap());
608
    }
609
610
    let inject_time = start.elapsed().as_nanos() / 1000;
    println!("  * Bootstrap Injection (avg): {:>6} ns/req", inject_time);
611
612

    // Calculate ratios
613
    let total_pipeline = serialize_time + deserialize_time + inject_time;
614
615
616
617
618
619
    println!("  * Total Pipeline (avg):    {:>8} ns/req", total_pipeline);

    println!("\nPerformance Insights:");
    if deserialize_time > serialize_time * 2 {
        println!("  • Deserialization is significantly faster than serialization");
    }
620
    if inject_time < serialize_time / 10 {
621
        println!(
622
623
            "  • Bootstrap injection overhead is negligible ({:.1}% of serialization)",
            (inject_time as f64 / serialize_time as f64) * 100.0
624
625
        );
    }
626
627
    if total_pipeline < 100_000 {
        println!("  • Total pipeline latency is excellent (< 100μs)");
628
629
    }

630
631
632
633
634
635
    println!("\nSimplification Benefits:");
    println!("  • Eliminated complex type conversion layer");
    println!("  • Reduced memory allocations");
    println!("  • Automatic field preservation (no manual mapping)");
    println!("  • Direct JSON manipulation improves performance");

636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
    println!("\nRecommendations:");
    if serialize_time > deserialize_time {
        println!("  • Focus optimization efforts on serialization rather than deserialization");
    }
    println!("  • PD mode overhead is minimal - safe to use for latency-sensitive workloads");
    println!("  • Consider batching small requests to improve overall throughput");

    println!("\n{}", "=".repeat(50));

    group.finish();
}

criterion_group!(
    benches,
    benchmark_summary,
    bench_json_serialization,
    bench_json_deserialization,
653
654
    bench_bootstrap_injection,
    bench_direct_json_routing,
655
656
657
658
    bench_throughput_by_size,
    bench_full_round_trip
);
criterion_main!(benches);