Skip to content

Commit 4bf87df

Browse files
committed
fix: remove mid-stream TPOT predictions from predicted-latency producer
The mid-stream prediction path in processTokenForLatencyPrediction was observability-only and low quality: computed from the scheduling-time metrics snapshot, averaged together with the scheduling-time TPOT prediction, and the only remaining caller of buildPredictionRequest using the scraped RunningRequestsSize gauge (train/serve inconsistency). avgPredictedTPOT now reduces to the scheduling-time prediction, which is the meaningful predicted-vs-actual comparison. The decodeTokenSampler becomes dead code with this removal (its only remaining consumer was the removed prediction path) and is deleted, along with the write-only tpotObservations field. The samplingMean and maxDecodeTokenSamplesForPrediction config parameters are kept in Config for compatibility with existing configs (the strict decoder would otherwise reject them) but are ignored, with a deprecation log when set. Fixes #2016 Signed-off-by: Michele Campi <215741962+MicheleCampi@users.noreply.github.com>
1 parent 9a8c999 commit 4bf87df

6 files changed

Lines changed: 16 additions & 296 deletions

File tree

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/decode_token_sampler.go

Lines changed: 0 additions & 111 deletions
This file was deleted.

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/decode_token_sampler_test.go

Lines changed: 0 additions & 96 deletions
This file was deleted.

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin.go

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,10 @@ func (pl *PredictedLatency) readInFlightLoad(endpoint fwksched.Endpoint) inFligh
227227
}
228228

229229
type Config struct {
230-
SamplingMean float64 `json:"samplingMean,omitempty"`
230+
// Deprecated: no longer used. Mid-stream TPOT predictions were removed;
231+
// the value is accepted for config compatibility and ignored.
232+
SamplingMean float64 `json:"samplingMean,omitempty"`
233+
// Deprecated: no longer used. Accepted for config compatibility and ignored.
231234
MaxDecodeTokenSamplesForPrediction int `json:"maxDecodeTokenSamplesForPrediction,omitempty"`
232235
SLOBufferFactor float64 `json:"sloBufferFactor,omitempty"`
233236
ContextTTL time.Duration `json:"contextTTL,omitempty"`
@@ -279,6 +282,9 @@ func PredictedLatencyFactory(name string, rawParameters *json.Decoder, handle pl
279282
if handle == nil {
280283
return nil, errors.New("plugin handle is required")
281284
}
285+
if parameters.SamplingMean != DefaultConfig.SamplingMean || parameters.MaxDecodeTokenSamplesForPrediction != DefaultConfig.MaxDecodeTokenSamplesForPrediction {
286+
log.FromContext(handle.Context()).Info("Deprecated: samplingMean and maxDecodeTokenSamplesForPrediction are ignored; mid-stream TPOT predictions were removed")
287+
}
282288
if err := registerMetrics(handle.Metrics()); err != nil {
283289
return nil, err
284290
}
@@ -294,14 +300,6 @@ func PredictedLatencyFactory(name string, rawParameters *json.Decoder, handle pl
294300
func (c *Config) validate() error {
295301
var errs []error
296302

297-
if c.SamplingMean <= 0 {
298-
errs = append(errs, fmt.Errorf("samplingMean must be > 0, got %f", c.SamplingMean))
299-
}
300-
301-
if c.MaxDecodeTokenSamplesForPrediction < 0 {
302-
errs = append(errs, fmt.Errorf("maxDecodeTokenSamplesForPrediction must be >= 0, got %d", c.MaxDecodeTokenSamplesForPrediction))
303-
}
304-
305303
if c.SLOBufferFactor <= 0 {
306304
errs = append(errs, fmt.Errorf("sloBufferFactor must be > 0, got %f", c.SLOBufferFactor))
307305
}
@@ -383,8 +381,6 @@ type predictedLatencyCtx struct {
383381
predictedTTFT float64
384382
avgTPOT float64
385383
avgPredictedTPOT float64
386-
decodeTokenSampler *decodeTokenSampler
387-
tpotObservations []float64
388384
predictedTPOTObservations []float64
389385

390386
inputTokenCount int

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/plugin_test.go

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -391,16 +391,10 @@ func TestPredictedLatencyFactory(t *testing.T) {
391391
expectErr: false,
392392
},
393393
{
394-
name: "invalid samplingMean <= 0",
395-
pluginName: "bad-sampling-mean",
396-
jsonParams: `{"samplingMean": -1.0}`,
397-
expectErr: true,
398-
},
399-
{
400-
name: "invalid maxSampledTokens < 0",
401-
pluginName: "bad-max-tokens",
402-
jsonParams: `{"maxDecodeTokenSamplesForPrediction": -1}`,
403-
expectErr: true,
394+
name: "deprecated sampling params are accepted and ignored",
395+
pluginName: "deprecated-sampling-params",
396+
jsonParams: `{"samplingMean": -1.0, "maxDecodeTokenSamplesForPrediction": -1}`,
397+
expectErr: false,
404398
},
405399
{
406400
name: "invalid sloBufferFactor <= 0",

pkg/epp/framework/plugins/requestcontrol/dataproducer/predictedlatency/requestcontrol_hooks.go

Lines changed: 5 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -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-
302290
func 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.
317305
func 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

Comments
 (0)