feat: use message LCP cache affinity keys

2026-06-29 15:47:57 +08:00
parent 229ed6a669
commit 24eb68cc09
6 changed files with 179 additions and 37 deletions
@@ -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,
@@ -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
 	}