BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Actions 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
apps/api/internal/httpapi/request_preparation.go
View File

@ -63,12 +63,14 @@ func (s *Server) prepareTaskRequest(ctx context.Context, r *http.Request, user *
return preparedTaskRequest{}, err
}
inputs := make([]store.ConversationMessageInput, 0, len(messages))
messageHashes := make([]any, 0, len(messages))
for _, rawMessage := range messages {
message, _ := rawMessage.(map[string]any)
if message == nil {
message = map[string]any{"content": rawMessage}
}
hash, assetHashes := canonicalConversationMessageHash(message)
messageHashes = append(messageHashes, hash)
inputs = append(inputs, store.ConversationMessageInput{
Hash: hash,
Role: stringFromRequestAny(message["role"]),
@ -83,6 +85,7 @@ func (s *Server) prepareTaskRequest(ctx context.Context, r *http.Request, user *
preparedBody["conversationId"] = conversationKey
preparedBody["conversationRecordId"] = conversationID
preparedBody["messageRefs"] = messageRefsForRequest(refs)
preparedBody["messageHashes"] = messageHashes
preparedBody["newMessageCount"] = newCount
delete(preparedBody, "messages")
result.ConversationID = conversationID

97
apps/api/internal/runner/cache_affinity.go
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[:])

41
apps/api/internal/runner/cache_affinity_test.go
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)
}
}

10
apps/api/internal/runner/service.go
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,

15
apps/api/internal/store/cache_affinity.go
View File

@ -7,6 +7,7 @@ import (
type CacheAffinityObservationInput struct {
CacheAffinityKey string
CacheAffinityKeys []string
CacheAffinityPolicy map[string]any
PlatformID string
PlatformModelID string
@ -19,9 +20,10 @@ type CacheAffinityObservationInput struct {
func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
input.CacheAffinityKeys = normalizedCacheAffinityKeys(input.CacheAffinityKey, input.CacheAffinityKeys)
input.ClientID = strings.TrimSpace(input.ClientID)
input.ModelType = strings.TrimSpace(input.ModelType)
if input.CacheAffinityKey == "" || input.ClientID == "" || input.ModelType == "" {
if len(input.CacheAffinityKeys) == 0 || input.ClientID == "" || input.ModelType == "" {
return nil
}
if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
@ -39,6 +41,15 @@ func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheA
}
alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
for _, key := range input.CacheAffinityKeys {
if err := s.recordCacheAffinityObservationKey(ctx, input, key, hitRatio, alpha); err != nil {
return err
}
}
return nil
}
func (s *Store) recordCacheAffinityObservationKey(ctx context.Context, input CacheAffinityObservationInput, key string, hitRatio float64, alpha float64) error {
_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_cache_affinity_stats (
client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
@ -60,7 +71,7 @@ func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheA
last_observed_at = now(),
updated_at = now()`,
input.ClientID,
input.CacheAffinityKey,
key,
input.PlatformID,
input.PlatformModelID,
input.ModelType,

50
apps/api/internal/store/candidates.go
View File

@ -13,6 +13,7 @@ import (
type ListModelCandidatesOptions struct {
CacheAffinityKey string
CacheAffinityKeys []string
CacheAffinityPolicy map[string]any
}
@ -46,6 +47,7 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
COALESCE(s.waiting_count, 0)::float8,
COALESCE(s.limiter_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8,
COALESCE(ca.cache_affinity_key, '')::text,
COALESCE(ca.request_count, 0)::float8,
COALESCE(ca.input_tokens, 0)::float8,
COALESCE(ca.cached_input_tokens, 0)::float8,
@ -59,10 +61,15 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
LEFT JOIN gateway_cache_affinity_stats ca
ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
AND ca.cache_affinity_key = NULLIF($4::text, '')
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
LEFT JOIN LATERAL (
SELECT ca.cache_affinity_key, ca.request_count, ca.input_tokens, ca.cached_input_tokens, ca.ema_hit_ratio, ca.last_hit_ratio, ca.last_observed_at
FROM unnest($4::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank)
JOIN gateway_cache_affinity_stats ca ON ca.cache_affinity_key = affinity_keys.cache_affinity_key
WHERE ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
ORDER BY affinity_keys.affinity_rank ASC, ca.cached_input_tokens DESC, ca.ema_hit_ratio DESC, ca.last_observed_at DESC
LIMIT 1
) ca ON TRUE
LEFT JOIN (
SELECT scope_key, SUM(lease_value) AS active
FROM gateway_concurrency_leases
@ -175,7 +182,7 @@ WHERE p.status = 'enabled'
COALESCE(s.running_count, 0) ASC,
COALESCE(s.waiting_count, 0) ASC,
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKey, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
if err != nil {
return nil, err
}
@ -211,6 +218,7 @@ WHERE p.status = 'enabled'
var stateWaitingCount float64
var stateLimiterRatio float64
var lastAssignedUnix float64
var cacheAffinityKey string
var cacheRequestCount float64
var cacheInputTokens float64
var cacheCachedInputTokens float64
@ -269,6 +277,7 @@ WHERE p.status = 'enabled'
&stateWaitingCount,
&stateLimiterRatio,
&lastAssignedUnix,
&cacheAffinityKey,
&cacheRequestCount,
&cacheInputTokens,
&cacheCachedInputTokens,
@ -302,6 +311,7 @@ WHERE p.status = 'enabled'
item.WaitingCount = maxFloat(queuedWaiting, stateWaitingCount)
item.LastAssignedUnix = lastAssignedUnix
applyRuntimeCandidateCacheAffinity(&item, listOptions, runtimeCandidateCacheAffinityInput{
Key: cacheAffinityKey,
RequestCount: int(cacheRequestCount),
InputTokens: int(cacheInputTokens),
CachedInputTokens: int(cacheCachedInputTokens),
@ -421,6 +431,7 @@ type runtimeCandidateLoadInput struct {
}
type runtimeCandidateCacheAffinityInput struct {
Key string
RequestCount int
InputTokens int
CachedInputTokens int
@ -435,15 +446,38 @@ func normalizeListModelCandidatesOptions(modelType string, options ...ListModelC
}
out := options[0]
out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey)
if out.CacheAffinityKey == "" || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
out.CacheAffinityKeys = normalizedCacheAffinityKeys(out.CacheAffinityKey, out.CacheAffinityKeys)
if len(out.CacheAffinityKeys) > 0 {
out.CacheAffinityKey = out.CacheAffinityKeys[0]
}
if len(out.CacheAffinityKeys) == 0 || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
out.CacheAffinityKey = ""
out.CacheAffinityKeys = nil
}
return out
}
func normalizedCacheAffinityKeys(primary string, keys []string) []string {
seen := map[string]bool{}
out := make([]string, 0, len(keys)+1)
for _, key := range append([]string{primary}, keys...) {
key = strings.TrimSpace(key)
if key == "" || seen[key] {
continue
}
seen[key] = true
out = append(out, key)
}
return out
}
func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, options ListModelCandidatesOptions, input runtimeCandidateCacheAffinityInput) {
key := strings.TrimSpace(input.Key)
if key == "" {
key = options.CacheAffinityKey
}
affinity := RuntimeCandidateCacheAffinity{
Key: options.CacheAffinityKey,
Key: key,
RequestCount: input.RequestCount,
InputTokens: input.InputTokens,
CachedInputTokens: input.CachedInputTokens,
@ -455,7 +489,7 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
hasSampledAffinity := input.RequestCount >= minSamples && affinity.EMAHitRatio > 0
if options.CacheAffinityKey == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
if key == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
candidate.CacheAffinity = affinity
return
}