okooooooo/gpt-load
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(affinity): 支持跨模型复用亲和

Browse Source
This commit is contained in:
tbphp committed 2026-08-12 16:58:32 +08:00
1 parent 3d5575c682
commit e6443e8a4d
5 files changed
+126 -20
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/affinity/key.go
+1 -3
View File
@@ -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()))
}
internal/affinity/key_test.go
+9 -11
View File
@@ -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)
}
internal/gateway/affinity.go
-2
View File
@@ -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() {
internal/gateway/affinity_test.go
+116 -3
View File
@@ -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)
}
}
internal/gateway/handler.go
-1
View File
@@ -444,7 +444,6 @@ func (handler *Handler) Handle(ginContext *gin.Context) {
snapshot,
accessKey.ID,
selectedRoute.Protocol,
model,
metadata.AffinityPrefix,
allowedCredentialRefs,
)