feat: use message LCP cache affinity keys

This commit is contained in:
2026-06-29 15:47:57 +08:00
parent 229ed6a669
commit 24eb68cc09
6 changed files with 179 additions and 37 deletions
+6 -4
View File
@@ -149,9 +149,10 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
if err != nil {
return Result{}, err
}
cacheAffinityKey := buildCacheAffinityKey(task.Kind, modelType, body)
cacheAffinityKeys := buildCacheAffinityKeys(task.Kind, modelType, body)
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
CacheAffinityKey: cacheAffinityKey,
CacheAffinityKey: cacheAffinityKeys.Primary,
CacheAffinityKeys: cacheAffinityKeys.Lookup,
CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
})
if err != nil {
@@ -324,7 +325,7 @@ candidatesLoop:
break candidatesLoop
}
candidateBody := preprocessing.Body
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy)
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
if err == nil {
attemptNo = nextAttemptNo
billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
@@ -535,7 +536,7 @@ candidatesLoop:
return Result{Task: failed, Output: failed.Result}, lastErr
}
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool, cacheAffinityPolicy map[string]any) (clients.Response, error) {
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -769,6 +770,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
}
if err := s.store.RecordCacheAffinityObservation(context.WithoutCancel(ctx), store.CacheAffinityObservationInput{
CacheAffinityKey: candidate.CacheAffinity.Key,
CacheAffinityKeys: cacheAffinityRecordKeys,
CacheAffinityPolicy: cacheAffinityPolicy,
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,