@@ -149,15 +149,15 @@ func (pl *PredictedLatency) ResponseBody(ctx context.Context, request *fwksched.
149149
150150 if predictedLatencyCtx .ttft == 0 {
151151 if pl .config .StreamingMode && ! response .EndOfStream {
152- processFirstTokenForLatencyPrediction (ctx , pl .latencypredictor , pl .config .StreamingMode , pl .config .EndpointRoleLabel , predictedLatencyCtx , now , pl . config . SamplingMean , pl . config . MaxDecodeTokenSamplesForPrediction )
152+ processFirstTokenForLatencyPrediction (ctx , pl .latencypredictor , pl .config .StreamingMode , pl .config .EndpointRoleLabel , predictedLatencyCtx , now )
153153 }
154154 } else {
155- processTokenForLatencyPrediction (ctx , pl . typedName . Name , pl . typedName . Type , pl . latencypredictor , pl . config . EndpointRoleLabel , predictedLatencyCtx , targetMetadata , now , pl . config . SamplingMean , pl . config . MaxDecodeTokenSamplesForPrediction )
155+ processTokenForLatencyPrediction (ctx , predictedLatencyCtx , now )
156156 }
157157
158158 if response .EndOfStream {
159159 if ! pl .config .StreamingMode {
160- processFirstTokenForLatencyPrediction (ctx , pl .latencypredictor , pl .config .StreamingMode , pl .config .EndpointRoleLabel , predictedLatencyCtx , now , pl . config . SamplingMean , pl . config . MaxDecodeTokenSamplesForPrediction )
160+ processFirstTokenForLatencyPrediction (ctx , pl .latencypredictor , pl .config .StreamingMode , pl .config .EndpointRoleLabel , predictedLatencyCtx , now )
161161 }
162162
163163 if predictedLatencyCtx .ttft > 0 {
@@ -249,12 +249,9 @@ func processFirstTokenForLatencyPrediction(
249249 endpointRoleLabel string ,
250250 predictedLatencyCtx * predictedLatencyCtx ,
251251 now time.Time ,
252- samplingMean float64 ,
253- maxDecodeTokenSamplesForPrediction int ,
254252) {
255253 logger := log .FromContext (ctx )
256254
257- initializeSampler (ctx , predictedLatencyCtx , samplingMean , maxDecodeTokenSamplesForPrediction )
258255 predictedLatencyCtx .ttft = float64 (now .Sub (predictedLatencyCtx .requestReceivedTimestamp ).Milliseconds ())
259256 predictedLatencyCtx .generatedTokenCount = 1
260257
@@ -290,15 +287,6 @@ func processFirstTokenForLatencyPrediction(
290287 refreshLastSeenMetrics (ctx , predictedLatencyCtx )
291288}
292289
293- func initializeSampler (ctx context.Context , predictedLatencyCtx * predictedLatencyCtx , samplingMean float64 , maxDecodeTokenSamplesForPrediction int ) {
294- if predictedLatencyCtx .decodeTokenSampler == nil {
295- logger := log .FromContext (ctx )
296- requestID := predictedLatencyCtx .schedulingRequest .Headers [reqcommon .RequestIDHeaderKey ]
297- predictedLatencyCtx .decodeTokenSampler = newDecodeTokenSampler (requestID , samplingMean , maxDecodeTokenSamplesForPrediction )
298- logger .V (logutil .DEBUG ).Info ("Initialized token sampler for first token" , "request_id" , requestID , "next_prediction_token" , predictedLatencyCtx .decodeTokenSampler .getNextSampleToken ())
299- }
300- }
301-
302290func predictFirstTPOT (ctx context.Context , predictedLatencyCtx * predictedLatencyCtx ) {
303291 logger := log .FromContext (ctx )
304292 targetName := predictedLatencyCtx .targetMetadata .NamespacedName .Name
@@ -313,69 +301,20 @@ func predictFirstTPOT(ctx context.Context, predictedLatencyCtx *predictedLatency
313301 }
314302}
315303
316- // processTokenForLatencyPrediction records actual inter-token latency, sampled predictions, and advances timestamp.
304+ // processTokenForLatencyPrediction records the actual TPOT for the token and advances the timestamp.
317305func processTokenForLatencyPrediction (
318306 ctx context.Context ,
319- pluginName , pluginType string ,
320- predictor latencypredictor.PredictorInterface ,
321- endpointRoleLabel string ,
322307 predictedLatencyCtx * predictedLatencyCtx ,
323- targetEndpointMetadata * fwkdl.EndpointMetadata ,
324308 now time.Time ,
325- samplingMean float64 ,
326- maxDecodeTokenSamplesForPrediction int ,
327309) {
328310 logger := log .FromContext (ctx )
329311
330- if predictedLatencyCtx .decodeTokenSampler == nil {
331- requestID := predictedLatencyCtx .schedulingRequest .Headers [reqcommon .RequestIDHeaderKey ]
332- predictedLatencyCtx .decodeTokenSampler = newDecodeTokenSampler (requestID , samplingMean , maxDecodeTokenSamplesForPrediction )
333- logger .V (logutil .DEBUG ).Info ("Initialized token sampler for subsequent tokens" , "request_id" , requestID , "next_prediction_token" , predictedLatencyCtx .decodeTokenSampler .getNextSampleToken ())
334- }
335-
336312 latencyMs := float64 (now .Sub (predictedLatencyCtx .lastTokenTimestamp ).Milliseconds ())
337313 predictedLatencyCtx .generatedTokenCount ++
338314
339- if predictedLatencyCtx .generatedTokenCount == 2 || predictedLatencyCtx .decodeTokenSampler .shouldPredict (predictedLatencyCtx .generatedTokenCount ) {
340- predictedLatencyCtx .tpotObservations = append (predictedLatencyCtx .tpotObservations , latencyMs )
341- }
342315 if predictedLatencyCtx .generatedTokenCount == 2 {
343316 logger .V (logutil .DEBUG ).Info ("First inter-token latency observed" ,
344- "actual_tpot_ms" , latencyMs ,
345- "predicted_tpot_ms" , predictedLatencyCtx .avgPredictedTPOT )
346- }
347-
348- m , err := getLatestMetricsForProfile (predictedLatencyCtx , "" )
349- if err != nil {
350- logger .V (logutil .DEBUG ).Info ("Skipping TPOT prediction due to missing metrics or schedulingResult" , "error" , err )
351- return
352- }
353-
354- if predictedLatencyCtx .decodeTokenSampler .shouldPredict (predictedLatencyCtx .generatedTokenCount ) {
355- in := buildPredictionRequest (
356- endpointRoleLabel ,
357- targetEndpointMetadata ,
358- m ,
359- predictedLatencyCtx .inputTokenCount ,
360- predictedLatencyCtx .generatedTokenCount ,
361- 0 ,
362- 0 ,
363- 0 ,
364- )
365- start := time .Now ()
366- p , err := predictor .Predict (ctx , in )
367- dur := time .Since (start )
368- if err != nil || p == nil {
369- logger .V (logutil .DEBUG ).Error (err , "TPOT predict failed" , "duration_ms" , dur .Milliseconds ())
370- predictedLatencyCtx .predictedTPOTObservations = append (predictedLatencyCtx .predictedTPOTObservations , 0 )
371- predictedLatencyCtx .avgPredictedTPOT = calculateRunningAverage (predictedLatencyCtx .avgPredictedTPOT , 0 , len (predictedLatencyCtx .predictedTPOTObservations ))
372- } else {
373- logger .V (logutil .DEBUG ).Info ("TPOT predict succeeded" , "value_ms" , p .TPOT , "duration_ms" , dur .Milliseconds ())
374- predictedLatencyCtx .predictedTPOTObservations = append (predictedLatencyCtx .predictedTPOTObservations , p .TPOT )
375- predictedLatencyCtx .avgPredictedTPOT = calculateRunningAverage (predictedLatencyCtx .avgPredictedTPOT , p .TPOT , len (predictedLatencyCtx .predictedTPOTObservations ))
376- }
377- recordRequestTPOTPredictionDuration (ctx , pluginName , pluginType , predictedLatencyCtx .schedulingRequest .TargetModel , predictedLatencyCtx .incomingModelName , dur .Seconds ())
378- predictedLatencyCtx .decodeTokenSampler .recordPrediction (predictedLatencyCtx .generatedTokenCount )
317+ "actual_tpot_ms" , latencyMs )
379318 }
380319
381320 predictedLatencyCtx .lastTokenTimestamp = now
0 commit comments