Files
easyai-ai-gateway/apps/api/internal/runner/admission_test.go
T
wangbo c28bf74230 feat(worker): 实现集群限流与自适应负载
保留平台模型 RPM、TPM 和并发策略语义,增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载,避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标,并扩展本地 acceptance runner,覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证:go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过;48/48 成功,无越限、重复提交、重复计费、重复回调或终态资源泄漏。
2026-08-03 00:13:46 +08:00

237 lines
8.1 KiB
Go

package runner
import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdmissionInputFromSnapshotRefreshesWorkerCapacity(t *testing.T) {
admission := store.TaskAdmission{
TaskID: "task-1",
PlatformID: "platform-1",
PlatformModelID: "model-1",
UserGroupID: "group-1",
QueueKey: "queue-1",
Mode: "async",
Priority: 7,
Scopes: []store.AdmissionScope{
{ScopeType: "platform_model", ScopeKey: "model-1", ConcurrentLimit: 10},
{ScopeType: "worker_capacity", ScopeKey: "global", ConcurrentLimit: 48},
},
ReselectRequestedAt: time.Time{},
}
input := admissionInputFromSnapshot(admission, 24)
if input.TaskID != admission.TaskID || input.PlatformModelID != admission.PlatformModelID || len(input.Scopes) != 2 {
t.Fatalf("unexpected admission input: %+v", input)
}
if input.Scopes[1].ConcurrentLimit != 24 {
t.Fatalf("worker capacity=%v, want 24", input.Scopes[1].ConcurrentLimit)
}
if input.Scopes[0].ConcurrentLimit != 10 {
t.Fatalf("business scope changed: %+v", input.Scopes[0])
}
if admission.Scopes[1].ConcurrentLimit != 48 {
t.Fatalf("snapshot was mutated: %+v", admission.Scopes[1])
}
}
func TestDistributedAdmissionModelTypeBoundary(t *testing.T) {
for _, modelType := range []string{
"image_generate", "image_edit", "image_analysis", "image_vectorize",
"video_generate", "video_enhance", "image_to_video", "text_to_video",
"video_edit", "video_reference", "video_first_last_frame", "video_understanding", "omni_video", "omni",
"audio_generate", "audio_understanding", "text_to_speech", "voice_clone",
} {
if !distributedAdmissionModelType(modelType) {
t.Errorf("%s should use distributed admission", modelType)
}
}
for _, modelType := range []string{"text_generate", "embedding", "rerank", "", "unknown"} {
if distributedAdmissionModelType(modelType) {
t.Errorf("%s should preserve immediate execution/rate-limit behavior", modelType)
}
}
}
func TestAcceptanceAdmissionScopesIsolateRunWithoutChangingLimits(t *testing.T) {
input := []store.AdmissionScope{{
ScopeType: "platform_model",
ScopeKey: "model-1",
ConcurrentLimit: 10,
QueueLimit: 0,
MaxWaitSeconds: 0,
}, {
ScopeType: "worker_capacity",
ScopeKey: "global",
ConcurrentLimit: 48,
}}
got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"}, input)
if got[0].ConcurrentLimit != 10 || got[0].ScopeKey != "acceptance:run-1:model-1" {
t.Fatalf("protocol-emulated acceptance must preserve the limit in an isolated scope, got %+v", got[0])
}
if got[0].QueueLimit != acceptanceQueueLimit || got[0].MaxWaitSeconds != acceptanceQueueMaxWait {
t.Fatalf("acceptance must enable a bounded queue, got %+v", got[0])
}
if got[1].ConcurrentLimit != 48 || got[1].ScopeKey != "acceptance:run-1:global" {
t.Fatalf("acceptance worker capacity changed, got %+v", got[1])
}
if input[0].QueueLimit != 0 || input[0].MaxWaitSeconds != 0 {
t.Fatalf("acceptance queue overlay must not mutate the production scopes, got %+v", input[0])
}
}
func TestAcceptanceCanaryPreservesProductionConcurrency(t *testing.T) {
input := []store.AdmissionScope{{
ScopeType: "platform_model",
ScopeKey: "model-1",
ConcurrentLimit: 10,
}}
got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance_canary"}, input)
if got[0].ConcurrentLimit != 10 {
t.Fatalf("real acceptance canary must preserve production concurrency, got %+v", got[0])
}
}
func TestAcceptanceInfrastructureReservationsIsolateEveryMetric(t *testing.T) {
input := []store.RateLimitReservation{
{ScopeType: "platform_model", ScopeKey: "model-1", Metric: "concurrent", Limit: 10},
{ScopeType: "platform_model", ScopeKey: "model-1", Metric: "rpm", Limit: 600},
{ScopeType: "user_group", ScopeKey: "group-1", Metric: "concurrent", Limit: 20},
}
got := acceptanceInfrastructureReservations(store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"}, input)
if len(got) != 3 || got[0].ScopeKey != "acceptance:run-1:model-1" || got[1].ScopeKey != "acceptance:run-1:model-1" || got[2].ScopeKey != "acceptance:run-1:group-1" {
t.Fatalf("unexpected acceptance reservations: %+v", got)
}
canary := acceptanceInfrastructureReservations(store.GatewayTask{RunMode: "acceptance_canary"}, input)
if len(canary) != len(input) {
t.Fatalf("real canary reservations changed: %+v", canary)
}
}
func TestAcceptanceAdmissionScopesLeaveProductionPolicyUnchanged(t *testing.T) {
input := []store.AdmissionScope{{
ScopeType: "platform_model",
ScopeKey: "model-1",
ConcurrentLimit: 10,
QueueLimit: 3,
MaxWaitSeconds: 7,
}}
got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "production"}, input)
if got[0].QueueLimit != 3 || got[0].MaxWaitSeconds != 7 || got[0].ConcurrentLimit != 10 {
t.Fatalf("production admission policy changed: %+v", got[0])
}
}
func TestAsyncAdmissionBatchCoversCertifiedGlobalCapacity(t *testing.T) {
for _, testCase := range []struct {
globalHardLimit int
want int
}{
{globalHardLimit: 48, want: 48},
{globalHardLimit: 64, want: 64},
{globalHardLimit: 128, want: 64},
{globalHardLimit: 0, want: 64},
} {
if got := asyncAdmissionBatchLimit(testCase.globalHardLimit); got != testCase.want {
t.Fatalf("async admission batch limit for %d = %d, want %d", testCase.globalHardLimit, got, testCase.want)
}
}
}
func TestPinCandidatesToTaskAdmissionPreservesAdmittedCandidate(t *testing.T) {
input := []store.RuntimeModelCandidate{
{PlatformID: "platform-a", PlatformModelID: "model-a"},
{PlatformID: "platform-b", PlatformModelID: "model-b"},
{PlatformID: "platform-c", PlatformModelID: "model-c"},
}
admission := &store.TaskAdmission{
Status: "admitted",
PlatformID: "platform-b",
PlatformModelID: "model-b",
}
got, pinned := pinCandidatesToTaskAdmission(input, admission)
if !pinned {
t.Fatal("expected admitted candidate to be pinned")
}
if got[0].PlatformModelID != "model-b" || got[1].PlatformModelID != "model-a" || got[2].PlatformModelID != "model-c" {
t.Fatalf("unexpected candidate order: %+v", got)
}
if input[0].PlatformModelID != "model-a" {
t.Fatalf("candidate pinning mutated the caller slice: %+v", input)
}
}
func TestPinCandidatesToTaskAdmissionPreservesWaitingCandidate(t *testing.T) {
input := []store.RuntimeModelCandidate{
{PlatformID: "platform-a", PlatformModelID: "model-a"},
{PlatformID: "platform-b", PlatformModelID: "model-b"},
}
admission := &store.TaskAdmission{
Status: "waiting",
PlatformID: "platform-b",
PlatformModelID: "model-b",
}
got, pinned := pinCandidatesToTaskAdmission(input, admission)
if !pinned || got[0].PlatformModelID != "model-b" {
t.Fatalf("waiting candidate was not pinned: pinned=%v candidates=%+v", pinned, got)
}
}
func TestPinCandidatesToTaskAdmissionAllowsRequestedReselection(t *testing.T) {
input := []store.RuntimeModelCandidate{
{PlatformID: "platform-a", PlatformModelID: "model-a"},
{PlatformID: "platform-b", PlatformModelID: "model-b"},
}
admission := &store.TaskAdmission{
Status: "waiting",
PlatformID: "platform-b",
PlatformModelID: "model-b",
ReselectRequestedAt: time.Now(),
}
got, pinned := pinCandidatesToTaskAdmission(input, admission)
if pinned {
t.Fatalf("reselection request unexpectedly pinned the old candidate: %+v", got)
}
if got[0].PlatformModelID != "model-a" {
t.Fatalf("reselection changed the sorted candidate order: %+v", got)
}
}
func TestPinCandidatesToTaskAdmissionIgnoresMissingCandidate(t *testing.T) {
input := []store.RuntimeModelCandidate{
{PlatformID: "platform-a", PlatformModelID: "model-a"},
{PlatformID: "platform-b", PlatformModelID: "model-b"},
}
admission := &store.TaskAdmission{
Status: "admitted",
PlatformID: "platform-c",
PlatformModelID: "model-c",
}
got, pinned := pinCandidatesToTaskAdmission(input, admission)
if pinned {
t.Fatalf("unexpected candidate pin for missing binding: %+v", got)
}
if got[0].PlatformModelID != "model-a" {
t.Fatalf("candidate order changed for missing binding: %+v", got)
}
}