Repository: vllm-project/aibrix
File: pkg/cache/cache_impl.go (Go)

Code before:
func (c *Store) DoneRequestCount(ctx *types.RoutingContext, requestID string, modelName string, traceTerm int64) {
	if ctx != nil {
		c.donePodStats(ctx)
	}

	meta, ok := c.metaModels.Load(modelName)
	if ok {
		atomic.AddInt32(&meta.pendingRequests, -1)
	}

	// DoneRequest only works for current term, no need to retry.
	if enableGPUOptimizerTracing {
		c.getRequestTrace(modelName).DoneRequest(requestID, traceTerm)
	}
}

// DoneRequestTrace completes request tracing
// Parameters:
//
//	ctx: Routing context
//	requestID: Unique request identifier
//	modelName: Model handling the request
//	inputTokens: Input tokens count
//	outputTokens: Output tokens count
//	traceTerm: Trace term identifier
func (c *Store) DoneRequestTrace(ctx *types.RoutingContext, requestID string, modelName string, inputTokens, outputTokens, traceTerm int64) {
	if ctx != nil {
		c.donePodStats(ctx)
	}

	meta, ok := c.metaModels.Load(modelName)
	if ok {
		atomic.AddInt32(&meta.pendingRequests, -1)
	}

	if enableGPUOptimizerTracing {
		var traceKey string
		for {
			trace := c.getRequestTrace(modelName)
			if traceKey, ok = trace.DoneRequestTrace(requestID, inputTokens, outputTokens, traceKey, traceTerm); ok {
				break
			}
			// In case DoneRequest return false, it has been recycled and we want to retry.
		}
		klog.V(5).Infof("inputTokens: %v, outputTokens: %v, trace key: %s", inputTokens, outputTokens, traceKey)
	}
}

// AddSubscriber registers new metric subscriber
// Parameters:
//

Review comment (suggestion):
Right now, done request trace is called in gateway_resp_body. I identified two scenario where a request can complete before executing resp_body. You can use this PR to address this gap.

1) If error code is returned in resp_headers.
2) If the client cancels the request and error is return from here `if err := srv.Send(resp); err != nil {`

Code after:
func (c *Store) DoneRequestCount(ctx *types.RoutingContext, requestID string, modelName string, traceTerm int64) {
	if ctx != nil {
		c.donePodStats(ctx, requestID)
	}

	meta, ok := c.metaModels.Load(modelName)
	if ok {
		atomic.AddInt32(&meta.pendingRequests, -1)
	}

	// DoneRequest only works for current term, no need to retry.
	if enableGPUOptimizerTracing {
		c.getRequestTrace(modelName).DoneRequest(requestID, traceTerm)
	}
}

// DoneRequestTrace completes request tracing
// Parameters:
//
//	ctx: Routing context
//	requestID: Unique request identifier
//	modelName: Model handling the request
//	inputTokens: Input tokens count
//	outputTokens: Output tokens count
//	traceTerm: Trace term identifier
func (c *Store) DoneRequestTrace(ctx *types.RoutingContext, requestID string, modelName string, inputTokens, outputTokens, traceTerm int64) {
	if ctx != nil {
		c.donePodStats(ctx, requestID)
	}

	meta, ok := c.metaModels.Load(modelName)
	if ok {
		atomic.AddInt32(&meta.pendingRequests, -1)
	}

	if enableGPUOptimizerTracing {
		var traceKey string
		for {
			trace := c.getRequestTrace(modelName)
			if traceKey, ok = trace.DoneRequestTrace(requestID, inputTokens, outputTokens, traceKey, traceTerm); ok {
				break
			}
			// In case DoneRequest return false, it has been recycled and we want to retry.
		}
		klog.V(5).Infof("inputTokens: %v, outputTokens: %v, trace key: %s", inputTokens, outputTokens, traceKey)
	}
}

// AddSubscriber registers new metric subscriber
// Parameters:
//