diff --git a/internal/affinity/key.go b/internal/affinity/key.go index 0fc0e40d..aebdc0a3 100644 --- a/internal/affinity/key.go +++ b/internal/affinity/key.go @@ -7,7 +7,7 @@ import ( "gpt-load/internal/protocol" ) -const keyDomain = "gpt-load/affinity/prompt-prefix-hmac/v1" +const keyDomain = "gpt-load/affinity/prompt-prefix-hmac/v2" // Key is an opaque, tenant-scoped affinity cache key. type Key string @@ -26,7 +26,6 @@ func DeriveKey( hasher Hasher, accessKeyID uint, clientProtocol protocol.Protocol, - externalModel string, prefix []byte, ) Key { if hasher == nil || accessKeyID == 0 || !clientProtocol.Valid() || len(prefix) == 0 { @@ -38,7 +37,6 @@ func DeriveKey( binary.BigEndian.PutUint64(encodedID[:], uint64(accessKeyID)) material.Write(encodedID[:]) writeKeyField(&material, []byte(clientProtocol)) - writeKeyField(&material, []byte(externalModel)) writeKeyField(&material, prefix) return Key(hasher.Hash(material.String())) } diff --git a/internal/affinity/key_test.go b/internal/affinity/key_test.go index 4a30222a..88406a8a 100644 --- a/internal/affinity/key_test.go +++ b/internal/affinity/key_test.go @@ -18,7 +18,7 @@ func (testHasher) Hash(value string) string { func TestDeriveKeyScopesStablePrefix(t *testing.T) { t.Parallel() - base := DeriveKey(testHasher{}, 7, protocol.OpenAICompletions, "gpt-4o", []byte("stable-prefix")) + base := DeriveKey(testHasher{}, 7, protocol.OpenAICompletions, []byte("stable-prefix")) if !base.Valid() { t.Fatal("DeriveKey() returned an empty key") } @@ -26,23 +26,21 @@ func TestDeriveKeyScopesStablePrefix(t *testing.T) { name string accessKey uint protocol protocol.Protocol - model string prefix string }{ - {name: "access key", accessKey: 8, protocol: protocol.OpenAICompletions, model: "gpt-4o", prefix: "stable-prefix"}, - {name: "protocol", accessKey: 7, protocol: protocol.Anthropic, model: "gpt-4o", prefix: "stable-prefix"}, - {name: "model", accessKey: 7, protocol: protocol.OpenAICompletions, model: "gpt-4o-mini", prefix: "stable-prefix"}, - {name: "prefix", accessKey: 7, protocol: protocol.OpenAICompletions, model: "gpt-4o", prefix: "other-prefix"}, + {name: "access key", accessKey: 8, protocol: protocol.OpenAICompletions, prefix: "stable-prefix"}, + {name: "protocol", accessKey: 7, protocol: protocol.Anthropic, prefix: "stable-prefix"}, + {name: "prefix", accessKey: 7, protocol: protocol.OpenAICompletions, prefix: "other-prefix"}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - got := DeriveKey(testHasher{}, test.accessKey, test.protocol, test.model, []byte(test.prefix)) + got := DeriveKey(testHasher{}, test.accessKey, test.protocol, []byte(test.prefix)) if !got.Valid() || got == base { t.Fatalf("DeriveKey() = %q, want non-empty key different from base", got) } }) } - if duplicate := DeriveKey(testHasher{}, 7, protocol.OpenAICompletions, "gpt-4o", []byte("stable-prefix")); duplicate != base { + if duplicate := DeriveKey(testHasher{}, 7, protocol.OpenAICompletions, []byte("stable-prefix")); duplicate != base { t.Fatalf("duplicate DeriveKey() = %q, want %q", duplicate, base) } } @@ -64,7 +62,7 @@ func TestDeriveKeyRejectsIncompleteScope(t *testing.T) { } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - if got := DeriveKey(test.hasher, test.accessKey, test.protocol, "model", test.prefix); got.Valid() { + if got := DeriveKey(test.hasher, test.accessKey, test.protocol, test.prefix); got.Valid() { t.Fatalf("DeriveKey() = %q, want empty key", got) } }) @@ -74,8 +72,8 @@ func TestDeriveKeyRejectsIncompleteScope(t *testing.T) { func TestDeriveKeyUsesUnambiguousFieldBoundaries(t *testing.T) { t.Parallel() - left := DeriveKey(testHasher{}, 1, protocol.OpenAICompletions, "ab", []byte("c")) - right := DeriveKey(testHasher{}, 1, protocol.OpenAICompletions, "a", []byte("bc")) + left := DeriveKey(testHasher{}, 1, protocol.OpenAICompletions, []byte("ab")) + right := DeriveKey(testHasher{}, 1, protocol.OpenAICompletions, []byte("a")) if !left.Valid() || !right.Valid() || left == right { t.Fatalf("keys = %q / %q, want distinct non-empty values", left, right) } diff --git a/internal/gateway/affinity.go b/internal/gateway/affinity.go index 3bcb8678..af9ae74c 100644 --- a/internal/gateway/affinity.go +++ b/internal/gateway/affinity.go @@ -17,7 +17,6 @@ func (handler *Handler) resolveRequestAffinity( snapshot *state.ConfigSnapshot, accessKeyID uint, clientProtocol protocol.Protocol, - externalModel string, prefix []byte, allowedCredentialRefs map[uint]state.CredentialRef, ) requestAffinity { @@ -33,7 +32,6 @@ func (handler *Handler) resolveRequestAffinity( handler.encryption, accessKeyID, clientProtocol, - externalModel, prefix, ) if !key.Valid() { diff --git a/internal/gateway/affinity_test.go b/internal/gateway/affinity_test.go index 52ae24fa..62a72bd3 100644 --- a/internal/gateway/affinity_test.go +++ b/internal/gateway/affinity_test.go @@ -12,6 +12,7 @@ import ( "github.com/gin-gonic/gin" "gpt-load/internal/affinity" + "gpt-load/internal/execution" "gpt-load/internal/platform/config" "gpt-load/internal/protocol" "gpt-load/internal/state" @@ -56,6 +57,41 @@ func TestHandlerLearnsAndReusesAutomaticSoftAffinity(t *testing.T) { assertAffinityHits(t, sink.snapshot(), []bool{false, true}) } +func TestHandlerReusesSoftAffinityAcrossModelsWithSameCandidateRange(t *testing.T) { + forwarder := &scriptedForwarder{results: successfulAffinityResults(2)} + handler, manager, _ := newHandlerForTest(t, forwarder, "sk-one", "sk-two") + addAffinityModelRoute(t, manager.Current(), "gpt-4o-mini", 1) + sink := &recordingRequestLogSink{} + handler.requestLogSink = sink + useAffinityRandomValues(handler, 0, affinitySecondCredentialRand) + engine := newAffinityTestEngine(t, handler) + + serveAffinityModelRequest(t, engine, "gpt-4o") + serveAffinityModelRequest(t, engine, "gpt-4o-mini") + + assertAffinityAttemptKeys(t, forwarder.inputs, []string{"sk-one", "sk-one"}) + assertAffinityHits(t, sink.snapshot(), []bool{false, true}) +} + +func TestHandlerRelearnsSoftAffinityWhenModelCandidateRangeChanges(t *testing.T) { + forwarder := &scriptedForwarder{results: successfulAffinityResults(4)} + handler, manager, registry := newHandlerForTest(t, forwarder, "sk-one", "sk-two") + moveSecondAffinityCredentialToGroup(t, manager.Current(), registry) + sink := &recordingRequestLogSink{} + handler.requestLogSink = sink + useAffinityRandomValues(handler, 0, 0, 0, 0) + engine := newAffinityTestEngine(t, handler) + + serveAffinityModelRequest(t, engine, "gpt-4o") + serveAffinityModelRequest(t, engine, "gpt-4o-mini") + serveAffinityModelRequest(t, engine, "gpt-4o-mini") + serveAffinityModelRequest(t, engine, "gpt-4o") + + assertAffinityAttemptKeys(t, forwarder.inputs, []string{"sk-one", "sk-two", "sk-two", "sk-one"}) + assertAffinityHits(t, sink.snapshot(), []bool{false, false, true, false}) + assertAffinityUpstreamModels(t, forwarder.inputs, []string{"gpt-4o", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o"}) +} + func TestHandlerDoesNotLearnAffinityForNonParticipatingGroup(t *testing.T) { forwarder := &scriptedForwarder{results: successfulAffinityResults(2)} handler, manager, _ := newHandlerForTest(t, forwarder, "sk-one", "sk-two") @@ -227,7 +263,6 @@ func TestHandlerIgnoresAffinityAfterCredentialIdentityChanges(t *testing.T) { snapshot, 1, protocol.OpenAICompletions, - "gpt-4o", prefix, map[uint]state.CredentialRef{1: oldRef}, ) @@ -246,7 +281,6 @@ func TestHandlerIgnoresAffinityAfterCredentialIdentityChanges(t *testing.T) { snapshot, 1, protocol.OpenAICompletions, - "gpt-4o", prefix, map[uint]state.CredentialRef{1: oldRef}, ) @@ -259,7 +293,6 @@ func TestHandlerIgnoresAffinityAfterCredentialIdentityChanges(t *testing.T) { snapshot, 1, protocol.OpenAICompletions, - "gpt-4o", prefix, map[uint]state.CredentialRef{1: changedRef}, ) @@ -297,6 +330,75 @@ func useAffinityRandomValues(handler *Handler, values ...int64) { } } +func addAffinityModelRoute( + t *testing.T, + snapshot *state.ConfigSnapshot, + externalModel string, + groupID uint, +) { + t.Helper() + byOperation := snapshot.ExecutionCandidates[protocol.OpenAICompletions] + if byOperation == nil { + t.Fatal("OpenAI Completions execution candidates are missing") + } + byModel := byOperation[execution.OperationChatCompletion] + base := byModel["gpt-4o"] + if len(base) != 1 { + t.Fatalf("gpt-4o routes = %d, want 1", len(base)) + } + route := base[0] + route.GroupID = groupID + route.UpstreamModelID = externalModel + byModel[externalModel] = []state.RouteTarget{route} +} + +func moveSecondAffinityCredentialToGroup( + t *testing.T, + snapshot *state.ConfigSnapshot, + registry *state.CredentialRegistry, +) { + t.Helper() + group := snapshot.Groups[1] + group.ID = 2 + group.Name = "openai-two" + group.Models = []state.ModelConfig{{ID: "gpt-4o-mini"}} + snapshot.Groups[2] = group + catalog := snapshot.GroupCatalog[1] + catalog.ID = 2 + catalog.Name = group.Name + snapshot.GroupCatalog[2] = catalog + addAffinityModelRoute(t, snapshot, "gpt-4o-mini", 2) + + refs := registry.CaptureActiveCredentialRefs([]uint{1}) + if len(refs) != 2 { + t.Fatalf("captured credential refs = %d, want 2", len(refs)) + } + entries := make([]state.CredentialEntry, 0, len(refs)) + for _, ref := range refs { + groupID := uint(1) + if ref.ID == 2 { + groupID = 2 + } + entries = append(entries, state.CredentialEntry{ + ID: ref.ID, GroupID: groupID, Version: ref.Version, + IdentityGeneration: ref.IdentityGeneration, Fingerprint: ref.Fingerprint, + Status: state.CredentialStatusActive, EncryptedValue: ref.EncryptedValue, + }) + } + if err := registry.ReplaceCredentials(entries); err != nil { + t.Fatalf("ReplaceCredentials() error = %v", err) + } +} + +func serveAffinityModelRequest(t *testing.T, engine http.Handler, model string) { + t.Helper() + serveAffinityRequest( + t, + engine, + `{"model":"`+model+`","messages":[{"role":"user","content":"stable conversation"}]}`, + ) +} + func newAffinityTestEngine(t *testing.T, handler *Handler) http.Handler { t.Helper() engine := gin.New() @@ -340,3 +442,14 @@ func assertAffinityHits(t *testing.T, events []telemetry.RequestEvent, want []bo t.Fatalf("affinity hits = %#v, want %#v", got, want) } } + +func assertAffinityUpstreamModels(t *testing.T, inputs []ForwardInput, want []string) { + t.Helper() + got := make([]string, 0, len(inputs)) + for _, input := range inputs { + got = append(got, input.UpstreamModelID) + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("upstream models = %#v, want %#v", got, want) + } +} diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index a408a1f8..a87a4452 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -444,7 +444,6 @@ func (handler *Handler) Handle(ginContext *gin.Context) { snapshot, accessKey.ID, selectedRoute.Protocol, - model, metadata.AffinityPrefix, allowedCredentialRefs, )