@@ -11,6 +11,7 @@ import (
1111 "github.com/onsi/gomega"
1212 "github.com/openai/openai-go"
1313 "github.com/openai/openai-go/option"
14+ "github.com/openai/openai-go/packages/param"
1415)
1516
1617func newOpenAIClient () * openai.Client {
@@ -24,6 +25,30 @@ func extractInferenceHeaders(httpResp *http.Response) (string, string, string) {
2425 httpResp .Header .Get ("x-inference-port" )
2526}
2627
28+ func generateAndCheckLoad (count int ) {
29+ for range count {
30+ prefillPods , decodePods := getModelServerPods (podSelector , prefillSelector , decodeSelector )
31+ gomega .Expect (prefillPods ).Should (gomega .BeEmpty ())
32+ gomega .Expect (decodePods ).Should (gomega .HaveLen (1 ))
33+
34+ nsHdr , podHdr , _ := runCompletion (simplePrompt , simModelName )
35+ gomega .Expect (nsHdr ).Should (gomega .Equal (nsName ))
36+ gomega .Expect (podHdr ).Should (gomega .Equal (decodePods [0 ]))
37+
38+ nsHdr , podHdr , _ = runChatCompletion (simplePrompt , simModelName )
39+ gomega .Expect (nsHdr ).Should (gomega .Equal (nsName ))
40+ gomega .Expect (podHdr ).Should (gomega .Equal (decodePods [0 ]))
41+
42+ nsHdr , podHdr , _ = runEmbeddings (singleEmbedding , simModelName )
43+ gomega .Expect (nsHdr ).Should (gomega .Equal (nsName ))
44+ gomega .Expect (podHdr ).Should (gomega .Equal (decodePods [0 ]))
45+
46+ nsHdr , podHdr , _ = runEmbeddings (doubleEmbedding , simModelName )
47+ gomega .Expect (nsHdr ).Should (gomega .Equal (nsName ))
48+ gomega .Expect (podHdr ).Should (gomega .Equal (decodePods [0 ]))
49+ }
50+ }
51+
2752// doPost sends a POST request with a JSON body to the given path, asserts HTTP 200,
2853// and returns the x-inference-namespace, x-inference-pod headers and the response body.
2954func doPost (path , body string , extraHeaders map [string ]string ) (string , string , []byte ) {
@@ -46,6 +71,27 @@ func doPost(path, body string, extraHeaders map[string]string) (string, string,
4671 return resp .Header .Get ("x-inference-namespace" ), resp .Header .Get ("x-inference-pod" ), respBody
4772}
4873
74+ // doPostWithError sends a POST request with a JSON body to the given path
75+ // and returns the status code and the response body.
76+ func doPostWithError (path , body string , extraHeaders map [string ]string ) (int , []byte ) {
77+ req , err := http .NewRequest (http .MethodPost , fmt .Sprintf ("http://localhost:%s%s" , port , path ), strings .NewReader (body ))
78+ gomega .Expect (err ).ShouldNot (gomega .HaveOccurred ())
79+ req .Header .Set ("Content-Type" , "application/json" )
80+ for k , v := range extraHeaders {
81+ req .Header .Set (k , v )
82+ }
83+ resp , err := http .DefaultClient .Do (req )
84+ gomega .Expect (err ).ShouldNot (gomega .HaveOccurred ())
85+ defer func () {
86+ gomega .Expect (resp .Body .Close ()).ToNot (gomega .HaveOccurred ())
87+ }()
88+
89+ respBody , err := io .ReadAll (resp .Body )
90+ gomega .Expect (err ).ShouldNot (gomega .HaveOccurred ())
91+
92+ return resp .StatusCode , respBody
93+ }
94+
4995func runCompletion (prompt string , theModel openai.CompletionNewParamsModel ) (string , string , string ) {
5096 var httpResp * http.Response
5197
@@ -115,6 +161,27 @@ func runChatCompletion(prompt, modelName string) (string, string, string) {
115161 return extractInferenceHeaders (httpResp )
116162}
117163
164+ func runEmbeddings (embeddings []string , modelName string ) (string , string , string ) {
165+ var httpResp * http.Response
166+
167+ input := openai.EmbeddingNewParamsInputUnion {}
168+ if len (embeddings ) == 1 {
169+ input .OfString = param .NewOpt (embeddings [0 ])
170+ } else {
171+ input .OfArrayOfStrings = embeddings
172+ }
173+
174+ params := openai.EmbeddingNewParams {
175+ Input : input ,
176+ Model : modelName ,
177+ }
178+ resp , err := newOpenAIClient ().Embeddings .New (testConfig .Context , params , option .WithResponseInto (& httpResp ))
179+ gomega .Expect (err ).ShouldNot (gomega .HaveOccurred ())
180+ gomega .Expect (resp .Data ).Should (gomega .HaveLen (len (embeddings )))
181+
182+ return extractInferenceHeaders (httpResp )
183+ }
184+
118185// runRawChatCompletion POSTs the given JSON body to /v1/chat/completions and returns
119186// the x-inference-namespace and x-inference-pod response headers.
120187func runRawChatCompletion (body string ) (string , string ) {
@@ -217,3 +284,86 @@ func cacheThresholdHeaders(force bool) map[string]string {
217284 }
218285 return nil
219286}
287+
288+ func verifyMetrics (infPoolName string , numTargetPorts int ) {
289+
290+ generateAndCheckLoad (10 )
291+
292+ // Send a few errors
293+ for range 10 {
294+ doPostWithError ("/v1/chat/completions" , "an invalid body" , nil )
295+ }
296+
297+ metricsURL := fmt .Sprintf ("http://localhost:%s/metrics" , metricsPort )
298+
299+ if k8sContext != "" {
300+ // Use port-forward to access the EPP pod's metrics endpoint.
301+ startEPPMetricsPortForward ()
302+ }
303+
304+ theMetrics := getMetrics (metricsURL )
305+ gomega .Expect (theMetrics ).ShouldNot (gomega .BeEmpty ())
306+ metricsAsString := strings .Join (theMetrics , "\n " )
307+
308+ _ , decodePods := getModelServerPods (podSelector , prefillSelector , decodeSelector )
309+
310+ // Define the metrics we expect to see
311+ preset := []string { //nolint:prealloc
312+ "inference_objective_request_total" ,
313+ "inference_objective_request_error_total" ,
314+ "inference_objective_request_duration_seconds" ,
315+ "inference_objective_normalized_time_per_output_token_seconds" ,
316+ "inference_objective_request_sizes" ,
317+ "inference_objective_response_sizes" ,
318+ "inference_objective_input_tokens" ,
319+ "inference_objective_output_tokens" ,
320+ "inference_pool_average_kv_cache_utilization" ,
321+ "inference_pool_average_queue_size" ,
322+ "inference_pool_per_pod_queue_size" ,
323+ "inference_objective_running_requests" ,
324+ "inference_pool_ready_pods" ,
325+ "inference_extension_info" ,
326+
327+ // llm_d metrics
328+ "llm_d_epp_request_total" ,
329+ "llm_d_epp_request_error_total" ,
330+ "llm_d_epp_request_duration_seconds" ,
331+ "llm_d_epp_normalized_time_per_output_token_seconds" ,
332+ "llm_d_epp_request_sizes" ,
333+ "llm_d_epp_response_sizes" ,
334+ "llm_d_epp_input_tokens" ,
335+ "llm_d_epp_output_tokens" ,
336+ "llm_d_epp_average_kv_cache_utilization" ,
337+ "llm_d_epp_average_queue_size" ,
338+ "llm_d_epp_per_endpoint_queue_size" ,
339+ "llm_d_epp_running_requests" ,
340+ "llm_d_epp_ready_endpoints" ,
341+ "llm_d_epp_info" ,
342+ }
343+ expectedMetrics := make ([]string , 0 , len (preset )+ len (decodePods )* numTargetPorts * 2 )
344+ expectedMetrics = append (expectedMetrics , preset ... )
345+
346+ for _ , modelServerPodName := range decodePods {
347+ for rank := range numTargetPorts {
348+ metricQueueSize := fmt .Sprintf (
349+ "inference_pool_per_pod_queue_size{model_server_pod=\" %s-rank-%d\" ,name=\" %s\" }" ,
350+ modelServerPodName ,
351+ rank ,
352+ infPoolName )
353+ expectedMetrics = append (expectedMetrics , metricQueueSize )
354+
355+ metricQueueSizeNew := fmt .Sprintf (
356+ "llm_d_epp_per_endpoint_queue_size{model_server_endpoint=\" %s-rank-%d\" ,name=\" %s\" }" ,
357+ modelServerPodName ,
358+ rank ,
359+ infPoolName ,
360+ )
361+ expectedMetrics = append (expectedMetrics , metricQueueSizeNew )
362+ }
363+ }
364+
365+ // Check if all expected metrics are present in the metrics output.
366+ for _ , metric := range expectedMetrics {
367+ gomega .Expect (metricsAsString ).Should (gomega .ContainSubstring (metric ))
368+ }
369+ }
0 commit comments