BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: use message LCP cache affinity keys

Browse Source
This commit is contained in:
wangbo
2026-06-29 15:47:57 +08:00
parent 229ed6a669
commit 24eb68cc09
6 changed files with 179 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/runner/cache_affinity.go
+79 -18
View File
@@ -7,47 +7,78 @@ import (
"strings"
)
type cacheAffinityKeys struct {
Primary string
Lookup []string
Record []string
}
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
return buildCacheAffinityKeys(kind, modelType, body).Primary
}
func buildCacheAffinityKeys(kind string, modelType string, body map[string]any) cacheAffinityKeys {
if len(body) == 0 {
return ""
return cacheAffinityKeys{}
}
for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return "explicit:" + sha256Text(value)
key := "explicit:" + sha256Text(value)
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
}
}
for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return key + ":" + sha256Text(value)
key := key + ":" + sha256Text(value)
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
}
}
if !cacheAffinityPromptHashSupported(kind, modelType) {
return ""
return cacheAffinityKeys{}
}
payload := map[string]any{
basePayload := map[string]any{
"kind": kind,
"modelType": modelType,
}
if messages, ok := body["messages"]; ok {
payload["messages"] = messages
}
if tools, ok := body["tools"]; ok {
payload["tools"] = tools
basePayload["tools"] = tools
}
if instructions, ok := body["instructions"]; ok {
payload["instructions"] = instructions
basePayload["instructions"] = instructions
}
messageHashes := cacheAffinityMessageHashes(body)
if len(messageHashes) > 0 {
baseHash := stableJSONHash(basePayload)
keys := make([]string, 0, len(messageHashes))
for index := range messageHashes {
keyPayload := map[string]any{
"base": baseHash,
"messageHashes": messageHashes[:index+1],
"prefixLength": index + 1,
}
keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
}
if len(keys) == 0 {
return cacheAffinityKeys{}
}
lookup := make([]string, 0, len(keys))
for index := len(keys) - 1; index >= 0; index-- {
lookup = append(lookup, keys[index])
}
return cacheAffinityKeys{
Primary: keys[len(keys)-1],
Lookup: lookup,
Record: keys,
}
}
if input, ok := body["input"]; ok {
payload["input"] = input
basePayload["input"] = input
}
if len(payload) <= 2 {
return ""
if len(basePayload) <= 2 {
return cacheAffinityKeys{}
}
raw, err := json.Marshal(payload)
if err != nil || len(raw) == 0 {
return ""
}
return "prompt:" + sha256Text(string(raw))
key := "prompt:" + stableJSONHash(basePayload)
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
}
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
@@ -63,6 +94,36 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
}
}
func cacheAffinityMessageHashes(body map[string]any) []string {
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
return hashes
}
if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
return hashes
}
messages, ok := body["messages"].([]any)
if !ok || len(messages) == 0 {
return nil
}
hashes := make([]string, 0, len(messages))
for _, message := range messages {
raw, err := json.Marshal(message)
if err != nil || len(raw) == 0 {
continue
}
hashes = append(hashes, sha256Text(string(raw)))
}
return hashes
}
func stableJSONHash(value any) string {
raw, err := json.Marshal(value)
if err != nil || len(raw) == 0 {
return ""
}
return sha256Text(string(raw))
}
func sha256Text(value string) string {
sum := sha256.Sum256([]byte(value))
return hex.EncodeToString(sum[:])
apps/api/internal/runner/cache_affinity_test.go
+36 -5
View File
@@ -1,6 +1,9 @@
package runner
import "testing"
import (
"strings"
"testing"
)
func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
body := map[string]any{
@@ -10,7 +13,7 @@ func TestBuildCacheAffinityKeyUsesExplicitKeyFirst(t *testing.T) {
}
got := buildCacheAffinityKey("chat.completions", "text_generate", body)
if got == "" || got[:9] != "explicit:" {
if !strings.HasPrefix(got, "explicit:") {
t.Fatalf("expected explicit cache affinity key, got %q", got)
}
}
@@ -20,7 +23,7 @@ func TestBuildCacheAffinityKeyUsesSessionContinuity(t *testing.T) {
"session_id": "conversation-1",
})
if got == "" || got[:11] != "session_id:" {
if !strings.HasPrefix(got, "session_id:") {
t.Fatalf("expected session cache affinity key, got %q", got)
}
}
@@ -36,7 +39,35 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
first := buildCacheAffinityKey("chat.completions", "text_generate", body)
second := buildCacheAffinityKey("chat.completions", "text_generate", body)
if first == "" || first != second || first[:7] != "prompt:" {
t.Fatalf("expected stable prompt hash cache affinity key, first=%q second=%q", first, second)
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") {
t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second)
}
}
func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
first := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"messageHashes": []any{"system-hash", "user-hash"},
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
})
next := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"messageHashes": []any{"system-hash", "user-hash", "assistant-hash", "tool-hash"},
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
})
if !strings.HasPrefix(first.Primary, "prompt_lcp:") {
t.Fatalf("expected LCP prompt key, got %+v", first)
}
if len(first.Record) != 2 || len(next.Lookup) != 4 {
t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
}
found := false
for _, key := range next.Lookup {
if key == first.Primary {
found = true
break
}
}
if !found {
t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next)
}
}
apps/api/internal/runner/service.go
+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,