feat: add runner failover rules and cache affinity

2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions
@@ -0,0 +1,73 @@
+package store
+
+import (
+	"context"
+	"strings"
+)
+
+type CacheAffinityObservationInput struct {
+	CacheAffinityKey       string
+	CacheAffinityPolicy    map[string]any
+	PlatformID             string
+	PlatformModelID        string
+	ClientID               string
+	ModelType              string
+	InputTokens            int
+	CachedInputTokens      int
+	CachedInputTokensKnown bool
+}
+
+func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
+	input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
+	input.ClientID = strings.TrimSpace(input.ClientID)
+	input.ModelType = strings.TrimSpace(input.ModelType)
+	if input.CacheAffinityKey == "" || input.ClientID == "" || input.ModelType == "" {
+		return nil
+	}
+	if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
+		return nil
+	}
+	minInputTokens := cacheAffinityMinInputTokens(input.CacheAffinityPolicy)
+	if input.InputTokens < minInputTokens || input.InputTokens <= 0 {
+		return nil
+	}
+	if input.CachedInputTokens < 0 {
+		input.CachedInputTokens = 0
+	}
+	if input.CachedInputTokens > input.InputTokens {
+		input.CachedInputTokens = input.InputTokens
+	}
+	alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
+	hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
+	_, err := s.pool.Exec(ctx, `
+	INSERT INTO gateway_cache_affinity_stats (
+	  client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
+	  request_count, input_tokens, cached_input_tokens, ema_hit_ratio, last_hit_ratio, last_observed_at
+	)
+	VALUES (
+	  $1, $2, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, $5,
+	  1, $6, $7, $8, $8, now()
+	)
+	ON CONFLICT (client_id, cache_affinity_key) DO UPDATE
+	SET platform_id = COALESCE(EXCLUDED.platform_id, gateway_cache_affinity_stats.platform_id),
+	    platform_model_id = COALESCE(EXCLUDED.platform_model_id, gateway_cache_affinity_stats.platform_model_id),
+	    model_type = EXCLUDED.model_type,
+	    request_count = gateway_cache_affinity_stats.request_count + 1,
+	    input_tokens = gateway_cache_affinity_stats.input_tokens + EXCLUDED.input_tokens,
+	    cached_input_tokens = gateway_cache_affinity_stats.cached_input_tokens + EXCLUDED.cached_input_tokens,
+	    ema_hit_ratio = ($9::float8 * EXCLUDED.last_hit_ratio) + ((1 - $9::float8) * gateway_cache_affinity_stats.ema_hit_ratio),
+	    last_hit_ratio = EXCLUDED.last_hit_ratio,
+	    last_observed_at = now(),
+	    updated_at = now()`,
+		input.ClientID,
+		input.CacheAffinityKey,
+		input.PlatformID,
+		input.PlatformModelID,
+		input.ModelType,
+		input.InputTokens,
+		input.CachedInputTokens,
+		hitRatio,
+		alpha,
+	)
+	return err
+}