http_metrics.rs 21.2 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
5
6
use anyhow::Error;
use async_stream::stream;
use dynamo_llm::{
7
8
    http::service::{metrics::Endpoint, service_v2::HttpService},
    model_card::ModelDeploymentCard,
9
10
11
12
13
14
15
    protocols::{
        Annotated,
        openai::chat_completions::{
            NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse,
        },
    },
};
16
use dynamo_runtime::metrics::prometheus_names::frontend_service::METRICS_PREFIX_ENV;
17
18
19
20
21
22
23
use dynamo_runtime::{
    CancellationToken,
    pipeline::{
        AsyncEngine, AsyncEngineContextProvider, ManyOut, ResponseStream, SingleIn, async_trait,
    },
};
use std::{sync::Arc, time::Duration};
24

25
26
27
28
#[path = "common/ports.rs"]
mod ports;
use ports::get_random_port;

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
// Mock engine for testing metrics
struct MockModelEngine {}

#[async_trait]
impl
    AsyncEngine<
        SingleIn<NvCreateChatCompletionRequest>,
        ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
        Error,
    > for MockModelEngine
{
    async fn generate(
        &self,
        request: SingleIn<NvCreateChatCompletionRequest>,
    ) -> Result<ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>, Error> {
        let (request, context) = request.transfer(());
        let ctx = context.context();

        let mut generator = request.response_generator(ctx.id().to_string());

        let stream = stream! {
            // Simulate some processing time
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;

            // Generate 5 response chunks
            for i in 0..5 {
55
                let output = generator.create_choice(i, Some(format!("Mock response {i}")), None, None, None);
56
57
58
59
60
61
62
63
                yield Annotated::from_data(output);
            }
        };

        Ok(ResponseStream::new(Box::pin(stream), ctx))
    }
}

64
#[tokio::test]
65
66
67
68
69
70
71
72
async fn test_metrics_prefix_default() {
    // Test default prefix when no env var is set
    temp_env::async_with_vars([(METRICS_PREFIX_ENV, None::<&str>)], async {
        let port = get_random_port().await;
        let service = HttpService::builder().port(port).build().unwrap();
        let token = CancellationToken::new();
        let handle = service.spawn(token.clone()).await;
        wait_for_metrics_ready(port).await;
73

74
75
76
77
78
79
80
        // Populate labeled metrics
        let state = service.state_clone();
        {
            let _guard = state.metrics_clone().create_inflight_guard(
                "test-model",
                Endpoint::ChatCompletions,
                false,
81
                "",
82
83
            );
        }
84

85
86
87
88
89
90
        let body = reqwest::get(format!("http://localhost:{}/metrics", port))
            .await
            .unwrap()
            .text()
            .await
            .unwrap();
91
92

        // Assert metrics that are actually present in the default configuration
93
        assert!(body.contains("dynamo_frontend_requests_total"));
94
        assert!(body.contains("dynamo_frontend_inflight_requests"));
95
        assert!(body.contains("dynamo_frontend_request_duration_seconds"));
96
        assert!(body.contains("dynamo_frontend_disconnected_clients"));
97

98
99
100
101
102
        token.cancel();
        let _ = handle.await;
    })
    .await;
}
103

104
105
106
107
108
109
110
111
112
#[tokio::test]
async fn test_metrics_prefix_custom() {
    // Test custom prefix override via environment variable
    temp_env::async_with_vars([(METRICS_PREFIX_ENV, Some("custom_prefix"))], async {
        let port = get_random_port().await;
        let service = HttpService::builder().port(port).build().unwrap();
        let token = CancellationToken::new();
        let handle = service.spawn(token.clone()).await;
        wait_for_metrics_ready(port).await;
113

114
115
116
117
118
119
120
        // Populate labeled metrics
        let state = service.state_clone();
        {
            let _guard = state.metrics_clone().create_inflight_guard(
                "test-model",
                Endpoint::ChatCompletions,
                true,
121
                "",
122
123
124
125
126
127
128
129
130
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
            );
        }

        let body = reqwest::get(format!("http://localhost:{}/metrics", port))
            .await
            .unwrap()
            .text()
            .await
            .unwrap();
        assert!(body.contains("custom_prefix_requests_total"));
        assert!(!body.contains("dynamo_frontend_requests_total"));

        token.cancel();
        let _ = handle.await;
    })
    .await;
}

#[tokio::test]
async fn test_metrics_prefix_sanitized() {
    // Test that invalid prefix characters are sanitized
    temp_env::async_with_vars([(METRICS_PREFIX_ENV, Some("nv-llm/http service"))], async {
        let port = get_random_port().await;
        let service = HttpService::builder().port(port).build().unwrap();
        let token = CancellationToken::new();
        let handle = service.spawn(token.clone()).await;
        wait_for_metrics_ready(port).await;

        let state = service.state_clone();
        {
            let _guard = state.metrics_clone().create_inflight_guard(
                "test-model",
                Endpoint::ChatCompletions,
                true,
156
                "",
157
158
159
160
161
162
163
164
165
166
167
            );
        }

        let body = reqwest::get(format!("http://localhost:{}/metrics", port))
            .await
            .unwrap()
            .text()
            .await
            .unwrap();
        assert!(body.contains("nv_llm_http_service_requests_total"));
        assert!(!body.contains("dynamo_frontend_requests_total"));
168

169
170
171
172
        token.cancel();
        let _ = handle.await;
    })
    .await;
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
}

// Poll /metrics until ready or timeout
async fn wait_for_metrics_ready(port: u16) {
    let url = format!("http://localhost:{}/metrics", port);
    let start = tokio::time::Instant::now();
    let timeout = Duration::from_secs(5);
    loop {
        if start.elapsed() > timeout {
            panic!("Timed out waiting for metrics endpoint at {}", url);
        }
        match reqwest::get(&url).await {
            Ok(resp) if resp.status().is_success() => break,
            _ => tokio::time::sleep(Duration::from_millis(50)).await,
        }
    }
}
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211

#[tokio::test]
async fn test_metrics_with_mock_model() {
    // Test metrics collection with a mock model serving requests
    // Ensure we use the default prefix
    temp_env::async_with_vars([(METRICS_PREFIX_ENV, None::<&str>)], async {
        let port = get_random_port().await;
        let service = HttpService::builder()
            .port(port)
            .enable_chat_endpoints(true)
            .build()
            .unwrap();

        let state = service.state_clone();
        let manager = state.manager();

        // Start the HTTP service
        let token = CancellationToken::new();
        let cancel_token = token.clone();
        let task = tokio::spawn(async move { service.run(token.clone()).await });

        // Add mock model engine
212
        let card = ModelDeploymentCard::with_name_only("mockmodel");
213
214
        let mock_engine = Arc::new(MockModelEngine {});
        manager
215
            .add_chat_completions_model("mockmodel", card.mdcsum(), mock_engine)
216
217
218
219
220
221
222
223
            .unwrap();

        // Wait for service to be ready
        wait_for_metrics_ready(port).await;

        let client = reqwest::Client::new();

        // Create a chat completion request
224
225
226
        let message = dynamo_protocols::types::ChatCompletionRequestMessage::User(
            dynamo_protocols::types::ChatCompletionRequestUserMessage {
                content: dynamo_protocols::types::ChatCompletionRequestUserMessageContent::Text(
227
228
229
230
231
232
                    "Hello, mock model!".to_string(),
                ),
                name: None,
            },
        );

233
        let request = dynamo_protocols::types::CreateChatCompletionRequestArgs::default()
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
            .model("mockmodel")
            .messages(vec![message])
            .max_tokens(50u32)
            .stream(true)
            .build()
            .expect("Failed to build request");

        // Make the request to the HTTP service
        let response = client
            .post(format!("http://localhost:{}/v1/chat/completions", port))
            .json(&request)
            .send()
            .await
            .unwrap();

        assert!(
            response.status().is_success(),
            "Request failed: {:?}",
            response
        );

        // Consume the response stream to complete the request
        let _response_body = response.bytes().await.unwrap();

        // Give some time for metrics to be updated
        tokio::time::sleep(Duration::from_millis(100)).await;

        // Fetch and verify metrics
        let metrics_response = client
            .get(format!("http://localhost:{}/metrics", port))
            .send()
            .await
            .unwrap();

        assert!(metrics_response.status().is_success());
        let metrics_body = metrics_response.text().await.unwrap();

        println!("=== METRICS WITH MOCK MODEL ===");
        println!("{}", metrics_body);
        println!("=== END METRICS ===");

        // Assert that key metrics are present with the mockmodel
        assert!(metrics_body.contains("dynamo_frontend_requests_total"));
        assert!(metrics_body.contains("model=\"mockmodel\""));
278
        assert!(metrics_body.contains("dynamo_frontend_inflight_requests"));
279
280
        assert!(metrics_body.contains("dynamo_frontend_request_duration_seconds"));
        assert!(metrics_body.contains("dynamo_frontend_output_sequence_tokens"));
281
        assert!(metrics_body.contains("dynamo_frontend_queued_requests"));
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299

        // Verify specific request counter incremented
        assert!(metrics_body.contains("endpoint=\"chat_completions\""));
        assert!(metrics_body.contains("request_type=\"stream\""));
        assert!(metrics_body.contains("status=\"success\""));

        // Clean up
        cancel_token.cancel();
        task.await.unwrap().unwrap();
    })
    .await;
}

// Integration tests that require distributed runtime with etcd
#[cfg(feature = "integration")]
mod integration_tests {
    use super::*;
    use dynamo_llm::{
300
        discovery::ModelWatcher, engines::make_echo_engine, entrypoint::EngineConfig,
301
        local_model::LocalModelBuilder, namespace::NamespaceFilter,
302
303
    };
    use dynamo_runtime::DistributedRuntime;
304
    use dynamo_runtime::discovery::DiscoveryQuery;
305
    use std::sync::Arc;
306
307
308
309
310

    #[tokio::test]
    #[ignore = "Requires etcd and distributed runtime"]
    async fn test_metrics_with_mdc_registration() {
        // Integration test for metrics collection with full MDC registration (like real model servers)
311
        let port = get_random_port().await;
312

313
314
315
316
317
        // Create distributed runtime (required for MDC registration)
        let runtime = dynamo_runtime::Runtime::from_settings().unwrap();
        let distributed_runtime = DistributedRuntime::from_settings(runtime.clone())
            .await
            .unwrap();
318

319
320
321
322
323
324
        // Create LocalModel with realistic configuration for testing
        let mut local_model = LocalModelBuilder::default()
            .model_name(Some("test-mdc-model".to_string()))
            .build()
            .await
            .unwrap();
325

326
        // Create EngineConfig with EchoEngine
327
        let engine_config = EngineConfig::InProcessText {
328
329
330
            engine: make_echo_engine(),
            model: Box::new(local_model.clone()),
        };
331

332
333
334
335
336
        let service = HttpService::builder()
            .port(port)
            .enable_chat_endpoints(true)
            .build()
            .unwrap();
337

338
        // Set up model watcher to discover models via discovery interface (like production)
339
340
341
342
343
        // This is crucial for the polling task to find model entries

        let model_watcher = ModelWatcher::new(
            distributed_runtime.clone(),
            service.state().manager_clone(),
344
            dynamo_llm::entrypoint::RouterConfig::default(),
345
            0, // migration_limit
346
            None,
347
            service.state().metrics_clone(),
348
        );
349
350
351
352
353
354
355
356
357
        // Start watching for model registrations via discovery interface
        let discovery = distributed_runtime.discovery();
        let discovery_stream = discovery
            .list_and_watch(
                DiscoveryQuery::AllModels,
                Some(distributed_runtime.primary_token()),
            )
            .await
            .unwrap();
358

359
        // Spawn watcher task to discover models
360
        let _watcher_task = tokio::spawn(async move {
361
362
363
            model_watcher
                .watch(discovery_stream, NamespaceFilter::Global)
                .await;
364
        });
365

366
367
        let EngineConfig::InProcessText { engine, model, .. } = engine_config else {
            panic!("Expected InProcessText config");
368
        };
369

370
        let card = local_model.card().clone();
371
372
373
        let engine = Arc::new(dynamo_llm::engines::StreamingEngineAdapter::new(engine));
        let manager = service.model_manager();
        manager
374
            .add_chat_completions_model(model.service_name(), card.mdcsum(), engine.clone())
375
            .unwrap();
376

377
378
379
380
381
382
        // Now do the proper MDC registration via LocalModel::attach()
        // Create a component and endpoint for proper registration
        let namespace = distributed_runtime.namespace("test-namespace").unwrap();
        let test_component = namespace.component("test-mdc-component").unwrap();
        let test_endpoint = test_component.endpoint("test-mdc-endpoint");

383
        // This will store the MDC in key-value store for discovery
384
385
386
387
388
        local_model
            .attach(
                &test_endpoint,
                dynamo_llm::model_type::ModelType::Chat,
                dynamo_llm::model_type::ModelInput::Text,
389
                None,
390
391
392
            )
            .await
            .unwrap();
393

394
395
        // Manually save the model card and update metrics
        // This simulates what the ModelWatcher polling task would do in production
396
        let _ = manager.save_model_card("test-mdc-key", card.clone());
397
398
399
400
401
402
403
404
405
406
407
408
409

        if let Err(e) = service
            .state()
            .metrics_clone()
            .update_metrics_from_mdc(&card)
        {
            tracing::debug!(
                model = %card.display_name,
                error = %e,
                "Failed to update MDC metrics in test"
            );
        }

410
411
412
413
414
        // Start the HTTP service
        let token = CancellationToken::new();
        let cancel_token = token.clone();
        let service_for_task = service.clone();
        let task = tokio::spawn(async move { service_for_task.run(token.clone()).await });
415

416
417
        // Wait for service to be ready
        wait_for_metrics_ready(port).await;
418

419
420
        // Give a bit more time for background metrics collection
        tokio::time::sleep(Duration::from_secs(5)).await;
421

422
        let client = reqwest::Client::new();
423

424
        // Create a chat completion request
425
426
427
        let message = dynamo_protocols::types::ChatCompletionRequestMessage::User(
            dynamo_protocols::types::ChatCompletionRequestUserMessage {
                content: dynamo_protocols::types::ChatCompletionRequestUserMessageContent::Text(
428
429
430
431
432
                    "Hello, MDC model!".to_string(),
                ),
                name: None,
            },
        );
433

434
        let request = dynamo_protocols::types::CreateChatCompletionRequestArgs::default()
435
436
437
438
439
440
            .model(model.service_name())
            .messages(vec![message])
            .max_tokens(50u32)
            .stream(true)
            .build()
            .expect("Failed to build request");
441

442
443
444
445
446
447
448
        // Make the request to the HTTP service
        let response = client
            .post(format!("http://localhost:{}/v1/chat/completions", port))
            .json(&request)
            .send()
            .await
            .unwrap();
449

450
451
452
453
454
        assert!(
            response.status().is_success(),
            "Request failed: {:?}",
            response
        );
455

456
457
        // Consume the response stream to complete the request
        let _response_body = response.bytes().await.unwrap();
458

459
460
        // Wait for the fast polling interval (50ms) for MDC metrics
        tokio::time::sleep(Duration::from_millis(50)).await;
461

462
463
464
465
466
467
        // Fetch and verify metrics
        let metrics_response = client
            .get(format!("http://localhost:{}/metrics", port))
            .send()
            .await
            .unwrap();
468

469
470
        assert!(metrics_response.status().is_success());
        let metrics_body = metrics_response.text().await.unwrap();
471

472
473
474
        println!("=== METRICS WITH FULL MDC REGISTRATION ===");
        println!("{}", metrics_body);
        println!("=== END METRICS ===");
475

476
477
478
479
        // Assert basic metrics are present (using service_name from the model)
        let model_name = model.service_name();
        assert!(metrics_body.contains("dynamo_frontend_requests_total"));
        assert!(metrics_body.contains(&format!("model=\"{}\"", model_name)));
480
        assert!(metrics_body.contains("dynamo_frontend_inflight_requests"));
481
482
        assert!(metrics_body.contains("dynamo_frontend_request_duration_seconds"));
        assert!(metrics_body.contains("dynamo_frontend_output_sequence_tokens"));
483
        assert!(metrics_body.contains("dynamo_frontend_queued_requests"));
484

485
486
487
488
489
490
491
        // Assert MDC-based model configuration metrics are present
        // These MUST be present for the test to pass
        assert!(
            metrics_body.contains("dynamo_frontend_model_context_length"),
            "MDC metrics not found! Metrics body: {}",
            metrics_body
        );
492

493
494
        assert!(metrics_body.contains("dynamo_frontend_model_kv_cache_block_size"));
        assert!(metrics_body.contains("dynamo_frontend_model_migration_limit"));
495

496
497
498
499
500
        // Note: The following metrics are not present in this test because they require
        // actual inference engines (vllm/sglang/trtllm *.py) with real runtime configurations:
        // - dynamo_frontend_model_total_kv_blocks (requires actual KV cache from real engines)
        // - dynamo_frontend_model_max_num_seqs (requires actual batching config from real engines)
        // - dynamo_frontend_model_max_num_batched_tokens (requires actual batching config from real engines)
501

502
503
504
505
        // Verify specific request counter incremented
        assert!(metrics_body.contains("endpoint=\"chat_completions\""));
        assert!(metrics_body.contains("request_type=\"stream\""));
        assert!(metrics_body.contains("status=\"success\""));
506

507
508
        // etcd lease will ensure we everything is deleted from etcd

509
        // Now test the complete lifecycle: remove the model from etcd
510
        // We don't need to cleanup model manager because it's local to this test
511

512
513
        /*
        // Clean up
514
515
516
517
518
519
        // Remove the model using the cleaner ModelWatcher approach
        if let Some(etcd_client) = distributed_runtime.etcd_client() {
            // Use ModelWatcher to find and remove the model (following ModelWatcher::handle_delete pattern)
            let watcher = ModelWatcher::new(
                distributed_runtime.clone(),
                service.state().manager_clone(),
520
                dynamo_llm::entrypoint::RouterConfig::default(),
521
                0, // migration_limit
522
                None,
523
                service.state().metrics_clone(),
524
            );
525

526
            // Get all model entries for our test model
527
            let model_entries = watcher.entries_for_model("test-mdc-model").await.unwrap();
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559

            if !model_entries.is_empty() {
                // For each model entry, we need to find its etcd key and remove it
                // This follows the same pattern as ModelWatcher::handle_delete
                for model_entry in model_entries {
                    // Find the etcd key for this specific model entry
                    // etcd keys follow pattern: "models/{UUID}"
                    // Example: "models/11dff335-316d-4c9f-8229-88ad8e8dac5e"
                    let kvs = etcd_client.kv_get_prefix("models").await.unwrap();

                    // Find the key by matching ModelEntry JSON structure:
                    // {
                    //   "name": "test-mdc-model",
                    //   "endpoint": { "namespace": "...", "component": "...", "name": "..." },
                    //   "model_type": "Chat",
                    //   "runtime_config": { ... },
                    //   "model_input": "Text"
                    // }
                    let key = kvs
                        .iter()
                        .find(|kv| {
                            serde_json::from_slice::<ModelEntry>(kv.value())
                                .map(|entry| {
                                    entry.name == model_entry.name
                                        && entry.endpoint_id == model_entry.endpoint_id
                                })
                                .unwrap_or(false)
                        })
                        .map(|kv| kv.key_str().unwrap().to_string());

                    if let Some(key) = key {
                        // Remove from ModelManager first (this returns the ModelEntry)
560
                        if let Some(_removed_card) = manager.remove_model_card(&key) {
561
562
                            // Remove entire model (following ModelWatcher::handle_delete pattern)
                            manager.remove_model(&model_entry.name);
563
564
565
566
567

                            // Then delete from etcd
                            etcd_client.kv_delete(key.as_str(), None).await.unwrap();
                        }
                    }
568
569
                }
            }
570
        }
571
        */
572

573
574
        cancel_token.cancel();
        task.await.unwrap().unwrap();
575
576
    }
}