From c28bf74230f63d46a7458df5df0d04ff38db38d1 Mon Sep 17 00:00:00 2001 From: wangbo Date: Mon, 3 Aug 2026 00:13:46 +0800 Subject: [PATCH] =?UTF-8?q?feat(worker):=20=E5=AE=9E=E7=8E=B0=E9=9B=86?= =?UTF-8?q?=E7=BE=A4=E9=99=90=E6=B5=81=E4=B8=8E=E8=87=AA=E9=80=82=E5=BA=94?= =?UTF-8?q?=E8=B4=9F=E8=BD=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 保留平台模型 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 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。 --- .dockerignore | 1 + Dockerfile | 4 +- apps/api/cmd/acceptance-load/main.go | 219 +++++-- apps/api/cmd/acceptance-load/main_test.go | 91 ++- apps/api/docs/swagger.json | 126 ++++ apps/api/docs/swagger.yaml | 81 +++ .../api/internal/acceptanceemulator/server.go | 38 +- .../acceptanceemulator/server_test.go | 38 ++ .../internal/acceptanceworkload/profiles.go | 13 +- .../acceptanceworkload/profiles_test.go | 7 +- apps/api/internal/config/config.go | 7 + apps/api/internal/config/config_test.go | 5 + .../internal/httpapi/acceptance_handlers.go | 16 + apps/api/internal/httpapi/server.go | 1 + apps/api/internal/runner/admission.go | 51 +- apps/api/internal/runner/admission_test.go | 44 +- apps/api/internal/runner/queue_worker.go | 112 +++- apps/api/internal/runner/service.go | 50 +- apps/api/internal/runner/worker_load.go | 28 + apps/api/internal/securityevents/metrics.go | 185 ++++-- .../internal/securityevents/metrics_test.go | 22 + apps/api/internal/store/acceptance.go | 2 +- apps/api/internal/store/admission_queue.go | 39 +- .../store/admission_queue_integration_test.go | 45 ++ .../internal/store/async_worker_capacity.go | 35 +- apps/api/internal/store/candidates.go | 2 +- apps/api/internal/store/rate_limit_status.go | 2 +- apps/api/internal/store/rate_limits.go | 22 +- .../store/rate_limits_integration_test.go | 153 ++++- apps/api/internal/store/tasks_runtime.go | 10 +- apps/api/internal/store/worker_registry.go | 158 ++++- .../internal/store/worker_registry_test.go | 12 + apps/api/internal/workerload/controller.go | 381 ++++++++++++ .../internal/workerload/controller_test.go | 95 +++ apps/api/internal/workerload/sampler.go | 131 ++++ apps/api/internal/workerload/sampler_test.go | 34 + .../migrations/0101_worker_adaptive_load.sql | 44 ++ apps/web/src/App.tsx | 20 +- apps/web/src/api.ts | 5 + apps/web/src/app-state.ts | 2 + apps/web/src/navigation.ts | 2 +- apps/web/src/pages/AdminPage.tsx | 1 + .../web/src/pages/admin/RealtimeLoadPanel.tsx | 93 ++- .../local-acceptance/dependencies.lock | 2 + deploy/kubernetes/local-acceptance/k3d.yaml | 4 + packages/contracts/src/index.ts | 37 ++ scripts/acceptance/local-cluster.sh | 33 +- scripts/acceptance/local-control-plane.sh | 5 +- scripts/acceptance/provider-burst.sh | 515 ++++++++++++++++ scripts/acceptance/report.mjs | 4 +- scripts/acceptance/run-local-acceptance.sh | 366 ++++++++++- scripts/cluster/run-production-acceptance.sh | 579 +++++++++++++++++- 52 files changed, 3700 insertions(+), 272 deletions(-) create mode 100644 apps/api/internal/runner/worker_load.go create mode 100644 apps/api/internal/workerload/controller.go create mode 100644 apps/api/internal/workerload/controller_test.go create mode 100644 apps/api/internal/workerload/sampler.go create mode 100644 apps/api/internal/workerload/sampler_test.go create mode 100644 apps/api/migrations/0101_worker_adaptive_load.sql create mode 100644 scripts/acceptance/provider-burst.sh diff --git a/.dockerignore b/.dockerignore index de00df7..f15509b 100644 --- a/.dockerignore +++ b/.dockerignore @@ -8,6 +8,7 @@ .env *.log +.local-secrets node_modules **/node_modules diff --git a/Dockerfile b/Dockerfile index 27021c6..f645f44 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,8 +6,8 @@ ARG NODE_BUILD_IMAGE=node:${NODE_VERSION}-alpine ARG WEB_RUNTIME_IMAGE=nginx:1.27-alpine FROM --platform=$BUILDPLATFORM ${GO_BUILD_IMAGE} AS api-builder -ARG TARGETOS=linux -ARG TARGETARCH=amd64 +ARG TARGETOS +ARG TARGETARCH ARG GOPROXY=https://goproxy.cn,direct ENV GOPROXY=$GOPROXY diff --git a/apps/api/cmd/acceptance-load/main.go b/apps/api/cmd/acceptance-load/main.go index ae38f54..bb85519 100644 --- a/apps/api/cmd/acceptance-load/main.go +++ b/apps/api/cmd/acceptance-load/main.go @@ -38,8 +38,10 @@ const ( ) const ( - geminiRequestPrefix = `{"contents":[{"role":"user","parts":[{"text":"将参考图编辑为蓝色赛博朋克风格，保留主体结构"},{"inlineData":{"mimeType":"image/png","data":"` - geminiRequestSuffix = `"}}]}],"generationConfig":{"responseModalities":["IMAGE"]}}` + geminiRequestPrefix = `{"contents":[{"role":"user","parts":[{"text":"将多张参考图融合编辑为蓝色赛博朋克风格，保留主体结构"}` + geminiRequestImageStart = `,{"inlineData":{"mimeType":"image/png","data":"` + geminiRequestImageEnd = `"}}` + geminiRequestSuffix = `]}],"generationConfig":{"responseModalities":["IMAGE"]}}` ) var reportURLPattern = regexp.MustCompile(`(?i)(?:https?|postgres(?:ql)?):\/\/[^\s"'<>]+`) @@ -64,6 +66,8 @@ type options struct { shardIndex int shardCount int executionID string + requestCount int + externalClient *http.Client } type report struct { @@ -91,6 +95,7 @@ type phaseReport struct { P99MS float64 `json:"p99Ms"` DecodedOutputBytes int64 `json:"decodedOutputBytes,omitempty"` OutputSHA256 string `json:"outputSha256,omitempty"` + InputImagesPerRequest int `json:"inputImagesPerRequest,omitempty"` UniqueTaskIDs int `json:"uniqueTaskIds,omitempty"` UniqueImageCombos int `json:"uniqueImageCombinations,omitempty"` ForcedConversionTasks int `json:"forcedConversionTasks,omitempty"` @@ -148,7 +153,7 @@ func main() { ctx, cancel := context.WithTimeout(signalContext, opts.timeout) defer cancel() result := report{ - SchemaVersion: "acceptance-load-report/v1", + SchemaVersion: "acceptance-load-report/v2", RunID: opts.runID, Profile: opts.profile, StartedAt: time.Now().UTC(), SecretSafe: true, } phaseResults, runErr := run(ctx, opts) @@ -205,6 +210,7 @@ func parseOptions() (options, error) { var shardIndex int var shardCount int var executionID string + var requestCount int flag.StringVar(&profile, "profile", env("AI_GATEWAY_ACCEPTANCE_PROFILE", "simulated-all"), "simulated-all, Gemini profile, video profile, or real-canary") flag.StringVar(&reportPath, "report", env("AI_GATEWAY_ACCEPTANCE_REPORT", ""), "secret-safe JSON report path") flag.DurationVar(&timeout, "timeout", 45*time.Minute, "overall timeout") @@ -214,6 +220,7 @@ func parseOptions() (options, error) { flag.IntVar(&shardIndex, "shard-index", 0, "zero-based distributed load shard index") flag.IntVar(&shardCount, "shard-count", 1, "distributed load shard count") flag.StringVar(&executionID, "execution-id", "", "unique workload execution identifier") + flag.IntVar(&requestCount, "requests", 0, "override the fixed profile request count") flag.Parse() opts := options{ gateways: splitCSV(os.Getenv("AI_GATEWAY_ACCEPTANCE_GATEWAYS")), @@ -235,6 +242,7 @@ func parseOptions() (options, error) { shardIndex: shardIndex, shardCount: shardCount, executionID: strings.ToLower(strings.TrimSpace(executionID)), + requestCount: requestCount, } if opts.executionID == "" { opts.executionID = opts.profile @@ -245,8 +253,8 @@ func parseOptions() (options, error) { if opts.shardCount < 1 || opts.shardCount > 32 || opts.shardIndex < 0 || opts.shardIndex >= opts.shardCount { return options{}, fmt.Errorf("invalid load shard %d/%d", opts.shardIndex, opts.shardCount) } - if len(opts.gateways) < 1 || len(opts.gateways) > 2 || (opts.shardCount == 1 && len(opts.gateways) != 2) { - return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain two API base URLs, or one URL for a distributed shard") + if len(opts.gateways) < 1 || len(opts.gateways) > 2 { + return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_GATEWAYS must contain one or two API base URLs") } for index := range opts.gateways { opts.gateways[index] = strings.TrimRight(opts.gateways[index], "/") @@ -267,8 +275,11 @@ func parseOptions() (options, error) { if opts.timeout <= 0 { return options{}, errors.New("timeout must be positive") } + if opts.requestCount < 0 || opts.requestCount > 10000 { + return options{}, errors.New("requests must be between 0 and 10000") + } switch opts.profile { - case "simulated-smoke", "simulated-all", "gemini-baseline", "gemini-large", "gemini-peak", "video-throughput", "video-recovery", + case "simulated-smoke", "simulated-all", "gemini-baseline", "gemini-multi-image", "gemini-large", "gemini-peak", "video-throughput", "video-recovery", "mixed-soak", "mixed-overload": if opts.emulatorURL == "" { return options{}, errors.New("AI_GATEWAY_ACCEPTANCE_EMULATOR_URL is required for simulated profiles") @@ -322,28 +333,44 @@ func run(ctx context.Context, opts options) ([]phaseResult, error) { TLSClientConfig: tlsConfig, }, } + opts.externalClient = &http.Client{ + Timeout: opts.timeout, + Transport: &http.Transport{ + MaxIdleConns: 512, + MaxIdleConnsPerHost: 256, + ForceAttemptHTTP2: true, + }, + } results := make([]phaseResult, 0, 5) runPhase := func(name string) error { var result phaseResult if profile, ok := acceptanceworkload.GeminiProfileByName(name); ok { + requests := profile.Requests + if opts.requestCount > 0 { + requests = opts.requestCount + } result = runGemini( - ctx, client, opts, name, opts.shardRequestCount(profile.Requests), - profile.InputBytes, profile.OutputBytes, false, opts.shardIndex, opts.shardCount, + ctx, client, opts, name, opts.shardRequestCount(requests), + profile.InputImages, profile.InputBytes, profile.OutputBytes, false, opts.shardIndex, opts.shardCount, ) } else { switch name { case "smoke-gemini": - result = runGemini(ctx, client, opts, name, opts.shardRequestCount(4), 256<<10, 256<<10, false, opts.shardIndex, opts.shardCount) + result = runGemini(ctx, client, opts, name, opts.shardRequestCount(4), 1, 256<<10, 256<<10, false, opts.shardIndex, opts.shardCount) case "smoke-video": result = runVideo(ctx, client, opts, name, opts.shardRequestCount(10), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount) case "video-throughput": - result = runVideo(ctx, client, opts, name, opts.shardRequestCount(1200), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount) + requests := 1200 + if opts.requestCount > 0 { + requests = opts.requestCount + } + result = runVideo(ctx, client, opts, name, opts.shardRequestCount(requests), false, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount) case "video-recovery": result = runVideo(ctx, client, opts, name, opts.shardRequestCount(96), true, opts.emulatorFixtureURLs(), false, opts.shardIndex, opts.shardCount) case "mixed-soak", "mixed-overload": result = runMixed(ctx, client, opts, name == "mixed-overload") case "real-gemini-canary": - result = runGemini(ctx, client, opts, name, 1, 256<<10, 0, true, 0, 1) + result = runGemini(ctx, client, opts, name, 1, 1, 256<<10, 0, true, 0, 1) case "real-video-canary": result = runVideo(ctx, client, opts, name, 1, false, opts.realImageURLs, true, 0, 1) default: @@ -375,6 +402,7 @@ func runGemini( opts options, name string, requestCount int, + inputImages int, inputBytes int, expectedOutputBytes int, realUpstream bool, @@ -412,8 +440,14 @@ func runGemini( defer func() { <-slots }() requestStarted := time.Now() endpoint := opts.gateways[logicalIndex%len(opts.gateways)] + "/v1beta/models/" + url.PathEscape(opts.geminiModel) + ":generateContent" - input := paddedPNGVariant(inputBytes, geminiInputVariant(opts.runID+"/"+opts.executionID, logicalIndex)) - req, err := newGeminiRequest(ctx, endpoint, input) + var req *http.Request + var err error + if name == acceptanceworkload.GeminiMultiImage.Name { + req, err = newGeminiFileDataRequest(ctx, endpoint, geminiFixtureURLs(opts, inputImages, logicalIndex)) + } else { + inputs := geminiInputs(inputImages, inputBytes, opts.runID+"/"+opts.executionID, logicalIndex) + req, err = newGeminiRequest(ctx, endpoint, inputs) + } if err == nil { opts.setHeaders(req, logicalIndex, realUpstream) req.Header.Set("Content-Type", "application/json") @@ -462,6 +496,7 @@ func runGemini( result := phaseReport{ Name: name, Requests: requestCount, Completed: completed, Failed: failures, DurationMS: elapsed.Milliseconds(), DecodedOutputBytes: decodedBytes, OutputSHA256: expectedHash, + InputImagesPerRequest: inputImages, } if firstErr == nil && name == "gemini-baseline" && elapsed > 8*time.Minute { firstErr = fmt.Errorf("Gemini baseline exceeded 8 minutes: %s", elapsed) @@ -472,7 +507,7 @@ func runGemini( return phaseResult{report: result, latencies: latencies, err: firstErr} } -func streamGeminiRequestBody(input []byte) io.ReadCloser { +func streamGeminiRequestBody(inputs [][]byte) io.ReadCloser { reader, writer := io.Pipe() go func() { closeWithError := func(err error) { @@ -482,14 +517,24 @@ func streamGeminiRequestBody(input []byte) io.ReadCloser { closeWithError(err) return } - encoder := base64.NewEncoder(base64.StdEncoding, writer) - if _, err := encoder.Write(input); err != nil { - closeWithError(err) - return - } - if err := encoder.Close(); err != nil { - closeWithError(err) - return + for _, input := range inputs { + if _, err := io.WriteString(writer, geminiRequestImageStart); err != nil { + closeWithError(err) + return + } + encoder := base64.NewEncoder(base64.StdEncoding, writer) + if _, err := encoder.Write(input); err != nil { + closeWithError(err) + return + } + if err := encoder.Close(); err != nil { + closeWithError(err) + return + } + if _, err := io.WriteString(writer, geminiRequestImageEnd); err != nil { + closeWithError(err) + return + } } if _, err := io.WriteString(writer, geminiRequestSuffix); err != nil { closeWithError(err) @@ -500,9 +545,9 @@ func streamGeminiRequestBody(input []byte) io.ReadCloser { return reader } -func newGeminiRequest(ctx context.Context, endpoint string, input []byte) (*http.Request, error) { +func newGeminiRequest(ctx context.Context, endpoint string, inputs [][]byte) (*http.Request, error) { getBody := func() (io.ReadCloser, error) { - return streamGeminiRequestBody(input), nil + return streamGeminiRequestBody(inputs), nil } body, _ := getBody() request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, body) @@ -511,10 +556,57 @@ func newGeminiRequest(ctx context.Context, endpoint string, input []byte) (*http return nil, err } request.GetBody = getBody - request.ContentLength = int64(len(geminiRequestPrefix) + base64.StdEncoding.EncodedLen(len(input)) + len(geminiRequestSuffix)) + contentLength := len(geminiRequestPrefix) + len(geminiRequestSuffix) + for _, input := range inputs { + contentLength += len(geminiRequestImageStart) + base64.StdEncoding.EncodedLen(len(input)) + len(geminiRequestImageEnd) + } + request.ContentLength = int64(contentLength) return request, nil } +func newGeminiFileDataRequest(ctx context.Context, endpoint string, imageURLs []string) (*http.Request, error) { + parts := make([]any, 0, len(imageURLs)+1) + parts = append(parts, map[string]any{"text": "将多张参考图融合编辑为蓝色赛博朋克风格，保留主体结构"}) + for _, imageURL := range imageURLs { + parts = append(parts, map[string]any{"fileData": map[string]any{ + "mimeType": geminiFixtureMIMEType(imageURL), + "fileUri": imageURL, + }}) + } + body, err := json.Marshal(map[string]any{ + "contents": []any{map[string]any{"role": "user", "parts": parts}}, + "generationConfig": map[string]any{"responseModalities": []string{"IMAGE"}}, + }) + if err != nil { + return nil, err + } + return http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) +} + +func geminiFixtureMIMEType(imageURL string) string { + lower := strings.ToLower(imageURL) + switch { + case strings.Contains(lower, ".webp"): + return "image/webp" + case strings.Contains(lower, ".jpg"), strings.Contains(lower, ".jpeg"): + return "image/jpeg" + default: + return "image/png" + } +} + +func geminiFixtureURLs(opts options, count int, logicalIndex int) []string { + fixtures := opts.emulatorFixtureURLs() + if count > len(fixtures) { + count = len(fixtures) + } + urls := make([]string, 0, count) + for imageIndex := 0; imageIndex < count; imageIndex++ { + urls = append(urls, fixtures[(logicalIndex*count+imageIndex)%len(fixtures)]) + } + return urls +} + func runVideo( ctx context.Context, client *http.Client, @@ -629,7 +721,7 @@ func runVideo( } wg.Wait() submissionDuration := time.Since(submitStartedAt) - if firstErr == nil && name == "video-throughput" && submissionDuration > 10*time.Second { + if firstErr == nil && name == "video-throughput" && requestCount == 1200 && submissionDuration > 10*time.Second { firstErr = fmt.Errorf("1200 video submissions exceeded 10 seconds: %s", submissionDuration) } uniqueTaskIDs := map[string]struct{}{} @@ -672,18 +764,25 @@ func runVideo( }() } wg.Wait() - comboSet := map[string]struct{}{} - for _, combo := range combinations { - comboSet[strings.Join(combo, "\x00")] = struct{}{} - } - if firstErr == nil && !realUpstream && len(comboSet) != 128 { - firstErr = fmt.Errorf("video image combinations=%d, want 128", len(comboSet)) + uniqueImageCombinations := requestedVideoCombinationCount( + combinations, + requestCount, + requestOffset, + requestStride, + ) + expectedCombinations := min(requestCount, 128) + if firstErr == nil && !realUpstream && uniqueImageCombinations < expectedCombinations { + firstErr = fmt.Errorf( + "video image combinations=%d, want at least %d", + uniqueImageCombinations, + expectedCombinations, + ) } elapsed := time.Since(startedAt) result := phaseReport{ Name: name, Requests: requestCount, Completed: completed, Failed: failures, DurationMS: elapsed.Milliseconds(), SubmissionDurationMS: submissionDuration.Milliseconds(), - UniqueTaskIDs: len(uniqueTaskIDs), UniqueImageCombos: len(comboSet), + UniqueTaskIDs: len(uniqueTaskIDs), UniqueImageCombos: uniqueImageCombinations, ForcedConversionTasks: forcedConversions, } if firstErr == nil && completed != requestCount { @@ -755,7 +854,7 @@ func runMixed( defer func() { <-slots }() var result phaseResult if isImage { - result = runGemini(runCtx, client, opts, "mixed-image", 1, 256<<10, 256<<10, false, requestIndex, 1) + result = runGemini(runCtx, client, opts, "mixed-image", 1, 1, 256<<10, 256<<10, false, requestIndex, 1) } else { result = runVideo( runCtx, @@ -931,7 +1030,14 @@ func validateVideoAsset(ctx context.Context, client *http.Client, opts options, if gatewayRequest { opts.setHeaders(request, index, false) } - response, err := client.Do(request) + mediaClient := client + if !gatewayRequest && opts.gatewayTLSName != "" { + mediaClient = opts.externalClient + if mediaClient == nil { + mediaClient = &http.Client{Timeout: client.Timeout} + } + } + response, err := mediaClient.Do(request) if err != nil { return fmt.Errorf("download final video: %w", err) } @@ -1056,6 +1162,31 @@ func videoCombinations(images []string, count int) [][]string { return out } +func requestedVideoCombinationCount( + combinations [][]string, + requestCount int, + requestOffset int, + requestStride int, +) int { + if len(combinations) == 0 || requestCount <= 0 { + return 0 + } + if requestStride < 1 { + requestStride = 1 + } + seen := map[string]struct{}{} + for index := 0; index < requestCount; index++ { + logicalIndex := requestOffset + index*requestStride + imageCount := acceptanceworkload.VideoImageCount(logicalIndex) + combo := combinations[logicalIndex%len(combinations)] + if len(combo) > imageCount { + combo = combo[:imageCount] + } + seen[strings.Join(combo, "\x00")] = struct{}{} + } + return len(seen) +} + func streamGeminiImageHash(reader io.Reader) (int64, string, error) { buffered := bufio.NewReaderSize(reader, 64<<10) if err := scanUntil(buffered, []byte(`"inlineData"`), 256<<20); err != nil { @@ -1176,6 +1307,24 @@ func paddedPNG(size int) []byte { return paddedPNGVariant(size, 0) } +func geminiInputs(count int, totalBytes int, executionID string, logicalIndex int) [][]byte { + if count < 1 { + count = 1 + } + inputs := make([][]byte, 0, count) + baseSize := totalBytes / count + remainder := totalBytes % count + for imageIndex := 0; imageIndex < count; imageIndex++ { + size := baseSize + if imageIndex < remainder { + size++ + } + variant := geminiInputVariant(executionID, logicalIndex*count+imageIndex) + inputs = append(inputs, paddedPNGVariant(size, variant)) + } + return inputs +} + func geminiInputVariant(runID string, logicalIndex int) int { digest := sha256.Sum256([]byte(fmt.Sprintf("%s:%d", strings.TrimSpace(runID), logicalIndex))) return int(binary.BigEndian.Uint64(digest[:8]) & uint64(^uint(0)>>1)) diff --git a/apps/api/cmd/acceptance-load/main_test.go b/apps/api/cmd/acceptance-load/main_test.go index bd9916f..13edc5c 100644 --- a/apps/api/cmd/acceptance-load/main_test.go +++ b/apps/api/cmd/acceptance-load/main_test.go @@ -17,6 +17,12 @@ import ( "testing" ) +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return f(request) +} + func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) { payload := paddedPNG(4 << 20) response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`, @@ -44,8 +50,8 @@ func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) { } } -func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) { - input := paddedPNGVariant(2<<20, 37) +func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) { + inputs := geminiInputs(3, 2<<20, "multi-image-test", 37) var body struct { Contents []struct { Parts []struct { @@ -59,24 +65,26 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) { ResponseModalities []string `json:"responseModalities"` } `json:"generationConfig"` } - stream := streamGeminiRequestBody(input) + stream := streamGeminiRequestBody(inputs) defer stream.Close() if err := json.NewDecoder(stream).Decode(&body); err != nil { t.Fatalf("decode streamed Gemini request: %v", err) } - if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 2 || body.Contents[0].Parts[1].InlineData == nil { + if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 { t.Fatalf("unexpected Gemini body structure: %+v", body) } - inlineData := body.Contents[0].Parts[1].InlineData - if inlineData.MIMEType != "image/png" { - t.Fatalf("mime type=%q", inlineData.MIMEType) - } - decoded, err := base64.StdEncoding.DecodeString(inlineData.Data) - if err != nil { - t.Fatalf("decode streamed input: %v", err) - } - if !bytes.Equal(decoded, input) { - t.Fatal("streamed input differs from source") + for index, input := range inputs { + inlineData := body.Contents[0].Parts[index+1].InlineData + if inlineData == nil || inlineData.MIMEType != "image/png" { + t.Fatalf("image %d inline data=%+v", index, inlineData) + } + decoded, err := base64.StdEncoding.DecodeString(inlineData.Data) + if err != nil { + t.Fatalf("decode streamed input %d: %v", index, err) + } + if !bytes.Equal(decoded, input) { + t.Fatalf("streamed input %d differs from source", index) + } } if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" { t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities) @@ -84,8 +92,8 @@ func TestStreamGeminiRequestBodyPreservesInput(t *testing.T) { } func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) { - input := paddedPNGVariant(256<<10, 91) - request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", input) + inputs := geminiInputs(3, 768<<10, "replay-test", 91) + request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs) if err != nil { t.Fatalf("new Gemini request: %v", err) } @@ -114,6 +122,37 @@ func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) { } } +func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) { + request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{ + "https://fixtures.example/one.png", + "https://fixtures.example/two.png", + "https://fixtures.example/three.png", + }) + if err != nil { + t.Fatalf("new multi-image request: %v", err) + } + var body struct { + Contents []struct { + Parts []struct { + FileData *struct { + FileURI string `json:"fileUri"` + } `json:"fileData"` + } `json:"parts"` + } `json:"contents"` + } + if err := json.NewDecoder(request.Body).Decode(&body); err != nil { + t.Fatalf("decode multi-image request: %v", err) + } + if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 { + t.Fatalf("unexpected multi-image body: %+v", body) + } + for index := 1; index < 4; index++ { + if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" { + t.Fatalf("missing fileData at part %d", index) + } + } +} + func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) { images := make([]string, 16) for index := range images { @@ -133,6 +172,9 @@ func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) { if len(seen) != 128 { t.Fatalf("unique combinations=%d", len(seen)) } + if got := requestedVideoCombinationCount(combinations, 144, 0, 1); got != 131 { + t.Fatalf("requested combinations=%d, want 131", got) + } } func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) { @@ -178,7 +220,7 @@ func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) { gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1", runToken: "token-1", geminiModel: "gemini-image-test", } - result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 256<<10, 256<<10, false, 0, 1) + result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1) if result.err != nil || result.report.Completed != 8 { t.Fatalf("result=%+v err=%v", result.report, result.err) } @@ -254,6 +296,21 @@ func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) { } } +func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "video/mp4") + _, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'}) + })) + defer server.Close() + poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("gateway-only transport must not be used for external media") + })} + opts := options{gatewayTLSName: "gateway.easyai.local"} + if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil { + t.Fatalf("validate external video with independent transport: %v", err) + } +} + func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" { diff --git a/apps/api/docs/swagger.json b/apps/api/docs/swagger.json index 99853fd..411729c 100644 --- a/apps/api/docs/swagger.json +++ b/apps/api/docs/swagger.json @@ -2575,6 +2575,30 @@ } } }, + "/api/admin/runtime/workers": { + "get": { + "security": [ + { + "BearerAuth": [] + } + ], + "produces": [ + "application/json" + ], + "tags": [ + "acceptance" + ], + "summary": "获取集群 Worker 实时负载与领取容量", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/store.WorkerClusterRuntime" + } + } + } + } + }, "/api/admin/system/acceptance/capacity-profiles": { "get": { "security": [ @@ -16347,6 +16371,108 @@ "$ref": "#/definitions/store.GatewayWalletAccount" } } + }, + "store.WorkerClusterRuntime": { + "type": "object", + "properties": { + "capturedAt": { + "type": "string" + }, + "queue": { + "$ref": "#/definitions/store.WorkerQueueRuntime" + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/store.WorkerInstanceRuntime" + } + } + } + }, + "store.WorkerInstanceRuntime": { + "type": "object", + "properties": { + "activeLeases": { + "type": "integer" + }, + "allocatedCapacity": { + "type": "integer" + }, + "capacityLimit": { + "type": "integer" + }, + "drainingAt": { + "type": "string" + }, + "finalizingTasks": { + "type": "integer" + }, + "hardCapacityLimit": { + "type": "integer" + }, + "heartbeatAt": { + "type": "string" + }, + "heavyCapacity": { + "type": "integer" + }, + "instanceId": { + "type": "string" + }, + "loadSampledAt": { + "type": "string" + }, + "podName": { + "type": "string" + }, + "podUid": { + "type": "string" + }, + "preparingTasks": { + "type": "integer" + }, + "pressureReason": { + "type": "string" + }, + "pressureState": { + "type": "string" + }, + "reportedActiveTasks": { + "type": "integer" + }, + "revision": { + "type": "string" + }, + "runningTasks": { + "type": "integer" + }, + "safeCapacity": { + "type": "integer" + }, + "site": { + "type": "string" + }, + "status": { + "type": "string" + }, + "waitingUpstreamTasks": { + "type": "integer" + } + } + }, + "store.WorkerQueueRuntime": { + "type": "object", + "properties": { + "oldestWaitSeconds": { + "type": "number" + }, + "queued": { + "type": "integer" + }, + "running": { + "type": "integer" + } + } } }, "securityDefinitions": { diff --git a/apps/api/docs/swagger.yaml b/apps/api/docs/swagger.yaml index fdca05a..826f817 100644 --- a/apps/api/docs/swagger.yaml +++ b/apps/api/docs/swagger.yaml @@ -4196,6 +4196,73 @@ definitions: primaryAccount: $ref: '#/definitions/store.GatewayWalletAccount' type: object + store.WorkerClusterRuntime: + properties: + capturedAt: + type: string + queue: + $ref: '#/definitions/store.WorkerQueueRuntime' + workers: + items: + $ref: '#/definitions/store.WorkerInstanceRuntime' + type: array + type: object + store.WorkerInstanceRuntime: + properties: + activeLeases: + type: integer + allocatedCapacity: + type: integer + capacityLimit: + type: integer + drainingAt: + type: string + finalizingTasks: + type: integer + hardCapacityLimit: + type: integer + heartbeatAt: + type: string + heavyCapacity: + type: integer + instanceId: + type: string + loadSampledAt: + type: string + podName: + type: string + podUid: + type: string + preparingTasks: + type: integer + pressureReason: + type: string + pressureState: + type: string + reportedActiveTasks: + type: integer + revision: + type: string + runningTasks: + type: integer + safeCapacity: + type: integer + site: + type: string + status: + type: string + waitingUpstreamTasks: + type: integer + type: object + store.WorkerQueueRuntime: + properties: + oldestWaitSeconds: + type: number + queued: + type: integer + running: + type: integer + type: object info: contact: {} description: |- @@ -5846,6 +5913,20 @@ paths: summary: 更新 Runner 策略 tags: - runtime + /api/admin/runtime/workers: + get: + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/store.WorkerClusterRuntime' + security: + - BearerAuth: [] + summary: 获取集群 Worker 实时负载与领取容量 + tags: + - acceptance /api/admin/system/acceptance/capacity-profiles: get: produces: diff --git a/apps/api/internal/acceptanceemulator/server.go b/apps/api/internal/acceptanceemulator/server.go index fd578d0..1f2e89a 100644 --- a/apps/api/internal/acceptanceemulator/server.go +++ b/apps/api/internal/acceptanceemulator/server.go @@ -59,6 +59,7 @@ type Server struct { type Report struct { GeminiRequests int64 `json:"geminiRequests"` GeminiInvalid int64 `json:"geminiInvalid"` + GeminiInputImages int64 `json:"geminiInputImages"` GeminiInputBytes int64 `json:"geminiInputBytes"` GeminiOutputBytes int64 `json:"geminiOutputBytes"` VideoSubmissions int64 `json:"videoSubmissions"` @@ -163,7 +164,7 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { writeProtocolError(w, http.StatusBadRequest, err.Error()) return } - inputBytes, err := validateGeminiImageRequest(body) + inputBytes, inputImages, err := validateGeminiImageRequest(body) if err != nil { s.recordGeminiInvalid() writeProtocolError(w, http.StatusBadRequest, err.Error()) @@ -180,6 +181,7 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { s.geminiIdempotency[idempotencyKey] = struct{}{} } s.report.GeminiRequests++ + s.report.GeminiInputImages += int64(inputImages) s.report.GeminiInputBytes += int64(inputBytes) s.report.GeminiOutputBytes += int64(outputBytes) s.mu.Unlock() @@ -533,7 +535,7 @@ func decodeImageDataURL(raw string) ([]byte, error) { return payload, nil } -func validateGeminiImageRequest(body map[string]any) (int, error) { +func validateGeminiImageRequest(body map[string]any) (int, int, error) { generationConfig, _ := body["generationConfig"].(map[string]any) modalities, _ := generationConfig["responseModalities"].([]any) hasImageModality := false @@ -543,9 +545,10 @@ func validateGeminiImageRequest(body map[string]any) (int, error) { } } if !hasImageModality { - return 0, errors.New("generationConfig.responseModalities must contain IMAGE") + return 0, 0, errors.New("generationConfig.responseModalities must contain IMAGE") } total := 0 + images := 0 contents, _ := body["contents"].([]any) for _, rawContent := range contents { content, _ := rawContent.(map[string]any) @@ -558,22 +561,41 @@ func validateGeminiImageRequest(body map[string]any) (int, error) { } encoded := strings.TrimSpace(stringValue(inline["data"])) if encoded == "" { + fileData, _ := part["fileData"].(map[string]any) + if fileData == nil { + fileData, _ = part["file_data"].(map[string]any) + } + fileURI := strings.TrimSpace(stringValue(fileData["fileUri"])) + if fileURI == "" { + fileURI = strings.TrimSpace(stringValue(fileData["file_uri"])) + } + if fileURI == "" { + fileURI = strings.TrimSpace(stringValue(fileData["uri"])) + } + if fileURI == "" { + continue + } + if !strings.HasPrefix(fileURI, "http://") && !strings.HasPrefix(fileURI, "https://") { + return 0, 0, errors.New("Gemini fileData URI must use HTTP or HTTPS") + } + images++ continue } decoded, err := base64.StdEncoding.DecodeString(encoded) if err != nil { - return 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err) + return 0, 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err) } if len(decoded) == 0 { - return 0, errors.New("Gemini inlineData is empty") + return 0, 0, errors.New("Gemini inlineData is empty") } total += len(decoded) + images++ } } - if total == 0 { - return 0, errors.New("Gemini request has no inlineData image") + if images == 0 { + return 0, 0, errors.New("Gemini request has no inlineData or fileData image") } - return total, nil + return total, images, nil } func (s *Server) recordGeminiInvalid() { diff --git a/apps/api/internal/acceptanceemulator/server_test.go b/apps/api/internal/acceptanceemulator/server_test.go index f5cd4bb..6477365 100644 --- a/apps/api/internal/acceptanceemulator/server_test.go +++ b/apps/api/internal/acceptanceemulator/server_test.go @@ -155,6 +155,44 @@ func TestForcedConversionRejectsImagesOutsideOfficialSeedanceRange(t *testing.T) } } +func TestGeminiProtocolAcceptsMultipleFileDataImages(t *testing.T) { + server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }}) + httpServer := httptest.NewServer(server.Handler()) + defer httpServer.Close() + parts := []any{map[string]any{"text": "combine references"}} + for index := 0; index < 3; index++ { + parts = append(parts, map[string]any{"fileData": map[string]any{ + "mimeType": "image/png", + "fileUri": fmt.Sprintf("https://fixtures.example/%d.png", index), + }}) + } + body, _ := json.Marshal(map[string]any{ + "contents": []any{map[string]any{"parts": parts}}, + "generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}}, + }) + response, err := postIdempotent(httpServer.URL+"/v1beta/models/gemini-test:generateContent", body, "gemini-multi") + if err != nil { + t.Fatalf("Gemini multi-image request: %v", err) + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + payload, _ := io.ReadAll(response.Body) + t.Fatalf("Gemini multi-image status=%d body=%s", response.StatusCode, payload) + } + reportResponse, err := http.Get(httpServer.URL + "/report") + if err != nil { + t.Fatalf("get report: %v", err) + } + defer reportResponse.Body.Close() + var report Report + if err := json.NewDecoder(reportResponse.Body).Decode(&report); err != nil { + t.Fatalf("decode report: %v", err) + } + if report.GeminiRequests != 1 || report.GeminiInputImages != 3 || report.GeminiInvalid != 0 { + t.Fatalf("unexpected multi-image report: %+v", report) + } +} + func TestImageFixturesFullyDecode(t *testing.T) { for name, fixture := range buildFixtures() { if !strings.HasPrefix(fixture.ContentType, "image/") { diff --git a/apps/api/internal/acceptanceworkload/profiles.go b/apps/api/internal/acceptanceworkload/profiles.go index 519f2b5..15cf70d 100644 --- a/apps/api/internal/acceptanceworkload/profiles.go +++ b/apps/api/internal/acceptanceworkload/profiles.go @@ -5,6 +5,7 @@ import "time" type GeminiProfile struct { Name string Requests int + InputImages int InputBytes int OutputBytes int DelayMin time.Duration @@ -13,15 +14,19 @@ type GeminiProfile struct { var ( GeminiBaseline = GeminiProfile{ - Name: "gemini-baseline", Requests: 1000, InputBytes: 256 << 10, OutputBytes: 256 << 10, + Name: "gemini-baseline", Requests: 1000, InputImages: 1, InputBytes: 256 << 10, OutputBytes: 256 << 10, DelayMin: 4 * time.Second, DelayMax: 4 * time.Second, } + GeminiMultiImage = GeminiProfile{ + Name: "gemini-multi-image", Requests: 192, InputImages: 3, InputBytes: 768 << 10, OutputBytes: 0, + DelayMin: 8 * time.Second, DelayMax: 15 * time.Second, + } GeminiLarge = GeminiProfile{ - Name: "gemini-large", Requests: 128, InputBytes: 2 << 20, OutputBytes: 4 << 20, + Name: "gemini-large", Requests: 128, InputImages: 1, InputBytes: 2 << 20, OutputBytes: 4 << 20, DelayMin: 8 * time.Second, DelayMax: 15 * time.Second, } GeminiPeak = GeminiProfile{ - Name: "gemini-peak", Requests: 32, InputBytes: 8 << 20, OutputBytes: 8 << 20, + Name: "gemini-peak", Requests: 32, InputImages: 1, InputBytes: 8 << 20, OutputBytes: 8 << 20, DelayMin: 15 * time.Second, DelayMax: 30 * time.Second, } ) @@ -30,6 +35,8 @@ func GeminiProfileByName(name string) (GeminiProfile, bool) { switch name { case GeminiBaseline.Name: return GeminiBaseline, true + case GeminiMultiImage.Name: + return GeminiMultiImage, true case GeminiLarge.Name: return GeminiLarge, true case GeminiPeak.Name: diff --git a/apps/api/internal/acceptanceworkload/profiles_test.go b/apps/api/internal/acceptanceworkload/profiles_test.go index 434898c..1d7283b 100644 --- a/apps/api/internal/acceptanceworkload/profiles_test.go +++ b/apps/api/internal/acceptanceworkload/profiles_test.go @@ -3,15 +3,18 @@ package acceptanceworkload import "testing" func TestGeminiProfilesAndVideoDistribution(t *testing.T) { - for _, profile := range []GeminiProfile{GeminiBaseline, GeminiLarge, GeminiPeak} { + for _, profile := range []GeminiProfile{GeminiBaseline, GeminiMultiImage, GeminiLarge, GeminiPeak} { resolved, ok := GeminiProfileByName(profile.Name) if !ok || resolved != profile { t.Fatalf("profile %q resolved to %+v, ok=%v", profile.Name, resolved, ok) } - if profile.Requests <= 0 || profile.InputBytes <= 0 || profile.OutputBytes <= 0 || + if profile.Requests <= 0 || profile.InputImages <= 0 || profile.InputBytes < profile.InputImages || profile.OutputBytes < 0 || profile.DelayMin <= 0 || profile.DelayMax < profile.DelayMin { t.Fatalf("invalid Gemini profile: %+v", profile) } + if profile.Name != GeminiMultiImage.Name && profile.OutputBytes == 0 { + t.Fatalf("fixed-output Gemini profile has no output size: %+v", profile) + } } counts := map[int]int{} for index := 0; index < 100; index++ { diff --git a/apps/api/internal/config/config.go b/apps/api/internal/config/config.go index 88b0a15..d3d5be4 100644 --- a/apps/api/internal/config/config.go +++ b/apps/api/internal/config/config.go @@ -79,6 +79,7 @@ type Config struct { AsyncWorkerHardLimit int AsyncWorkerInstanceHardLimit int AsyncWorkerRefreshIntervalSeconds int + AsyncWorkerLoadMode string AsyncAdmissionMicrobatchSize int AsyncAdmissionDispatcherEnabled bool AsyncAdmissionDispatcherConfigured bool @@ -189,6 +190,7 @@ func Load() Config { ), AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32), AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5), + AsyncWorkerLoadMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_WORKER_LOAD_MODE", "adaptive"))), AsyncAdmissionMicrobatchSize: envIntValidated("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", 8), AsyncAdmissionDispatcherEnabled: envValue("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED") == "true", AsyncAdmissionDispatcherConfigured: envValue( @@ -299,6 +301,11 @@ func (c Config) Validate() error { if c.AsyncWorkerRefreshIntervalSeconds < 1 { return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive") } + switch strings.ToLower(strings.TrimSpace(c.AsyncWorkerLoadMode)) { + case "", "adaptive", "legacy": + default: + return errors.New("AI_GATEWAY_WORKER_LOAD_MODE must be adaptive or legacy") + } if c.AsyncAdmissionMicrobatchSize < 1 || c.AsyncAdmissionMicrobatchSize > 32 { return errors.New("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE must be between 1 and 32") } diff --git a/apps/api/internal/config/config_test.go b/apps/api/internal/config/config_test.go index 48fbd58..bb43ddc 100644 --- a/apps/api/internal/config/config_test.go +++ b/apps/api/internal/config/config_test.go @@ -95,6 +95,11 @@ func TestValidateAsyncWorkerSettings(t *testing.T) { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") { t.Fatalf("Validate() error = %v, want invalid refresh interval", err) } + cfg.AsyncWorkerRefreshIntervalSeconds = 5 + cfg.AsyncWorkerLoadMode = "invalid" + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOAD_MODE") { + t.Fatalf("Validate() error = %v, want invalid worker load mode", err) + } t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "not-an-integer") loaded := Load() if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") { diff --git a/apps/api/internal/httpapi/acceptance_handlers.go b/apps/api/internal/httpapi/acceptance_handlers.go index d1c842f..649eb40 100644 --- a/apps/api/internal/httpapi/acceptance_handlers.go +++ b/apps/api/internal/httpapi/acceptance_handlers.go @@ -264,6 +264,22 @@ func (s *Server) listCapacityProfiles(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, profiles) } +// getWorkerClusterRuntime godoc +// @Summary 获取集群 Worker 实时负载与领取容量 +// @Tags acceptance +// @Produce json +// @Security BearerAuth +// @Success 200 {object} store.WorkerClusterRuntime +// @Router /api/admin/runtime/workers [get] +func (s *Server) getWorkerClusterRuntime(w http.ResponseWriter, r *http.Request) { + runtime, err := s.acceptanceStore().GetWorkerClusterRuntime(r.Context()) + if err != nil { + writeError(w, http.StatusInternalServerError, "get worker cluster runtime failed") + return + } + writeJSON(w, http.StatusOK, runtime) +} + // createAcceptanceRun godoc // @Summary 创建生产同构验收 Run // @Tags acceptance diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index 45579b8..6419ea4 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -264,6 +264,7 @@ func NewServerWithStores( mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/promote", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.promoteAcceptanceRun))) mux.Handle("POST /api/admin/system/acceptance/runs/{runID}/abort", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.abortAcceptanceRun))) mux.Handle("GET /api/admin/system/acceptance/capacity-profiles", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listCapacityProfiles))) + mux.Handle("GET /api/admin/runtime/workers", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getWorkerClusterRuntime))) mux.Handle("GET /api/admin/system/identity/configuration", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityConfiguration))) mux.Handle("POST /api/admin/system/identity/pairings", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.startIdentityPairing))) mux.Handle("GET /api/admin/system/identity/pairings/{pairingID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getIdentityPairing))) diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index 40bdefb..2291b0c 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -205,23 +205,12 @@ func (s *Service) taskAdmissionScopes( } func acceptanceAdmissionScopes(task store.GatewayTask, scopes []store.AdmissionScope) []store.AdmissionScope { - if task.RunMode != "acceptance" && task.RunMode != "acceptance_canary" { + if task.RunMode != "acceptance" { return scopes } out := append([]store.AdmissionScope(nil), scopes...) for index := range out { - // Protocol-emulated acceptance measures Gateway and Worker capacity, so - // the isolated Run is bounded by the worker_capacity scope instead of a - // production supplier quota. The real acceptance_canary path deliberately - // retains the production platform-model concurrency limit. - if task.RunMode == "acceptance" && out[index].ScopeType == "platform_model" { - out[index].ConcurrentLimit = 0 - } - if out[index].ConcurrentLimit <= 0 { - if task.RunMode != "acceptance" || out[index].ScopeType != "platform_model" { - continue - } - } + out[index].ScopeKey = acceptanceScopeKey(task, out[index].ScopeKey) out[index].QueueLimit = acceptanceQueueLimit out[index].MaxWaitSeconds = acceptanceQueueMaxWait } @@ -235,16 +224,21 @@ func acceptanceInfrastructureReservations( if task.RunMode != "acceptance" { return reservations } - out := make([]store.RateLimitReservation, 0, len(reservations)) - for _, reservation := range reservations { - if reservation.ScopeType == "platform_model" && reservation.Metric == "concurrent" { - continue - } - out = append(out, reservation) + out := append([]store.RateLimitReservation(nil), reservations...) + for index := range out { + out[index].ScopeKey = acceptanceScopeKey(task, out[index].ScopeKey) } return out } +func acceptanceScopeKey(task store.GatewayTask, scopeKey string) string { + runID := strings.TrimSpace(task.AcceptanceRunID) + if runID == "" { + runID = "unbound" + } + return "acceptance:" + runID + ":" + scopeKey +} + func (s *Service) loadAsyncTaskAdmission(ctx context.Context, task store.GatewayTask) (*store.TaskAdmission, error) { if !task.AsyncMode { return nil, nil @@ -265,6 +259,7 @@ func pinCandidatesToTaskAdmission( ) ([]store.RuntimeModelCandidate, bool) { if admission == nil || (admission.Status != "waiting" && admission.Status != "admitted") || + !admission.ReselectRequestedAt.IsZero() || len(candidates) < 2 { return candidates, false } @@ -785,6 +780,24 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st return false, outcome.Err } if !outcome.Result.Admitted { + platformModelID := outcome.Result.Admission.PlatformModelID + if platformModelID == "" { + for _, input := range inputs { + if input.TaskID == outcome.TaskID { + platformModelID = input.PlatformModelID + break + } + } + } + if platformModelID != "" { + marked, markErr := s.store.RequestWaitingTaskAdmissionReselect(ctx, platformModelID) + if markErr != nil { + return false, markErr + } + if marked > 0 { + s.observeTaskAdmission("candidate_reselect_requested") + } + } return true, nil } } diff --git a/apps/api/internal/runner/admission_test.go b/apps/api/internal/runner/admission_test.go index 357edd2..fbdecf5 100644 --- a/apps/api/internal/runner/admission_test.go +++ b/apps/api/internal/runner/admission_test.go @@ -57,7 +57,7 @@ func TestDistributedAdmissionModelTypeBoundary(t *testing.T) { } } -func TestAcceptanceAdmissionScopesUseWorkerCapacityInsteadOfSupplierConcurrency(t *testing.T) { +func TestAcceptanceAdmissionScopesIsolateRunWithoutChangingLimits(t *testing.T) { input := []store.AdmissionScope{{ ScopeType: "platform_model", ScopeKey: "model-1", @@ -70,15 +70,15 @@ func TestAcceptanceAdmissionScopesUseWorkerCapacityInsteadOfSupplierConcurrency( ConcurrentLimit: 48, }} - got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance"}, input) + got := acceptanceAdmissionScopes(store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"}, input) - if got[0].ConcurrentLimit != 0 { - t.Fatalf("protocol-emulated acceptance must defer to worker capacity, got %+v", got[0]) + 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 { + 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 { @@ -100,15 +100,15 @@ func TestAcceptanceCanaryPreservesProductionConcurrency(t *testing.T) { } } -func TestAcceptanceInfrastructureReservationsOnlyRemoveSupplierConcurrency(t *testing.T) { +func TestAcceptanceInfrastructureReservationsIsolateEveryMetric(t *testing.T) { input := []store.RateLimitReservation{ - {ScopeType: "platform_model", Metric: "concurrent", Limit: 10}, - {ScopeType: "platform_model", Metric: "rpm", Limit: 600}, - {ScopeType: "user_group", Metric: "concurrent", Limit: 20}, + {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"}, input) - if len(got) != 2 || got[0].Metric != "rpm" || got[1].ScopeType != "user_group" { + 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) @@ -192,6 +192,28 @@ func TestPinCandidatesToTaskAdmissionPreservesWaitingCandidate(t *testing.T) { } } +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"}, diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 82ff0fd..400fff0 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -10,8 +10,10 @@ import ( "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload" "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" "github.com/riverqueue/river" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivermigrate" @@ -45,6 +47,10 @@ type asyncTaskWorker struct { service *Service } +type workerLoadSampler interface { + Sample(databaseConnections, databaseMax int32) workerload.ResourceSample +} + func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs]) error { task, err := w.service.store.GetTask(ctx, job.Args.TaskID) if err != nil { @@ -53,6 +59,12 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" { return nil } + loadLease, admitted := w.service.tryStartWorkerTask() + if !admitted { + return river.JobSnooze(workerLoadRetryDelay(task.ID)) + } + defer loadLease.Release() + ctx = context.WithValue(ctx, workerLoadLeaseContextKey{}, loadLease) executionToken := uuid.NewString() result, runErr := w.service.executeWithToken(ctx, task, authUserFromTask(task), nil, executionToken) if runErr == nil { @@ -323,21 +335,36 @@ func (s *Service) makeAsyncExecutionClient(capacity int) (asyncExecutionClient, func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorkerCapacitySnapshot, error) { if s.asyncCapacityLoader != nil { - return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit) + snapshot, err := s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit) + if err == nil && s.workerLoad != nil { + s.workerLoad.SetClaimLimit(snapshot.Capacity) + } + return snapshot, err } snapshot, err := s.coordinationStore.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit) if err != nil { return store.AsyncWorkerCapacitySnapshot{}, err } + loadSnapshot := s.sampleWorkerLoad() allocation, err := s.coordinationStore.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{ - InstanceID: s.workerInstanceID, - PodUID: strings.TrimSpace(os.Getenv("POD_UID")), - PodName: strings.TrimSpace(os.Getenv("POD_NAME")), - Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), - Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), - DesiredCapacity: snapshot.Capacity, - CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit, - HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, + InstanceID: s.workerInstanceID, + PodUID: strings.TrimSpace(os.Getenv("POD_UID")), + PodName: strings.TrimSpace(os.Getenv("POD_NAME")), + Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), + Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), + DesiredCapacity: snapshot.Capacity, + CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit, + LoadMode: loadSnapshot.Mode, + SafeCapacity: loadSnapshot.SafeCapacity, + HeavyCapacity: loadSnapshot.HeavyLimit, + ActiveTasks: loadSnapshot.ActiveTasks, + PreparingTasks: loadSnapshot.PreparingTasks, + WaitingUpstreamTasks: loadSnapshot.WaitingUpstreamTasks, + FinalizingTasks: loadSnapshot.FinalizingTasks, + PressureState: string(loadSnapshot.PressureState), + PressureReason: loadSnapshot.PressureReason, + LoadSampledAt: loadSnapshot.SampledAt, + HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, }) if err != nil { return store.AsyncWorkerCapacitySnapshot{}, err @@ -346,9 +373,60 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke snapshot.GlobalCapacity = allocation.GlobalAllocated snapshot.ActiveInstances = allocation.ActiveInstances snapshot.InstanceID = allocation.InstanceID + loadSnapshot = s.workerLoad.SetClaimLimit(allocation.Allocated) + snapshot.LoadMode = loadSnapshot.Mode + snapshot.LocalSafeCapacity = loadSnapshot.SafeCapacity + snapshot.LocalHeavyCapacity = loadSnapshot.HeavyLimit + snapshot.LocalActiveTasks = loadSnapshot.ActiveTasks + snapshot.LocalPreparingTasks = loadSnapshot.PreparingTasks + snapshot.LocalWaitingTasks = loadSnapshot.WaitingUpstreamTasks + snapshot.LocalFinalizingTasks = loadSnapshot.FinalizingTasks + snapshot.LocalPressureState = string(loadSnapshot.PressureState) + snapshot.LocalPressureReason = loadSnapshot.PressureReason return snapshot, nil } +func (s *Service) sampleWorkerLoad() workerload.Snapshot { + if s.workerLoad == nil { + return workerload.Snapshot{Mode: workerload.ModeLegacy, SafeCapacity: s.cfg.AsyncWorkerInstanceHardLimit, HeavyLimit: s.cfg.AsyncWorkerInstanceHardLimit} + } + if s.workerLoadSampler == nil { + return s.workerLoad.Snapshot() + } + var connections int32 + var maximum int32 + seen := make(map[*pgxpool.Pool]struct{}, 3) + for _, database := range []*store.Store{s.store, s.coordinationStore, s.riverStore} { + if database == nil || database.Pool() == nil { + continue + } + pool := database.Pool() + if _, ok := seen[pool]; ok { + continue + } + seen[pool] = struct{}{} + statistics := pool.Stat() + connections += statistics.AcquiredConns() + maximum += statistics.MaxConns() + } + return s.workerLoad.Observe(s.workerLoadSampler.Sample(connections, maximum)) +} + +func (s *Service) tryStartWorkerTask() (*workerload.Lease, bool) { + if s.workerLoad == nil { + return nil, true + } + return s.workerLoad.TryStart() +} + +func workerLoadRetryDelay(taskID string) time.Duration { + var value uint32 + for _, character := range []byte(taskID) { + value = value*33 + uint32(character) + } + return 250*time.Millisecond + time.Duration(value%501)*time.Millisecond +} + func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) { ticker := time.NewTicker(time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * time.Second) defer ticker.Stop() @@ -427,6 +505,12 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) { "modelDesired", snapshot.ModelDesired, "groupDesired", snapshot.GroupDesired, "capped", snapshot.Capped, + "loadMode", snapshot.LoadMode, + "localSafeCapacity", snapshot.LocalSafeCapacity, + "localHeavyCapacity", snapshot.LocalHeavyCapacity, + "localActiveTasks", snapshot.LocalActiveTasks, + "pressureState", snapshot.LocalPressureState, + "pressureReason", snapshot.LocalPressureReason, ) if oldClient != nil { go s.drainAsyncWorkerClient(oldClient) @@ -493,6 +577,16 @@ func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacityS if ok { distributedObserver.SetDistributedWorkerCapacity(snapshot.ActiveInstances, snapshot.GlobalCapacity, snapshot.Capacity) } + loadObserver, ok := s.billingMetrics.(interface { + SetWorkerLoad(activeLimit, heavyLimit, active, preparing, waiting, finalizing int, pressure string) + }) + if ok { + loadObserver.SetWorkerLoad( + snapshot.LocalSafeCapacity, snapshot.LocalHeavyCapacity, snapshot.LocalActiveTasks, + snapshot.LocalPreparingTasks, snapshot.LocalWaitingTasks, snapshot.LocalFinalizingTasks, + snapshot.LocalPressureState, + ) + } } func (s *Service) observeAsyncWorkerResize(outcome string) { diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 4eaa998..1b3af7d 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -19,6 +19,7 @@ import ( "github.com/easyai/easyai-ai-gateway/apps/api/internal/config" scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/riverqueue/river" @@ -39,6 +40,8 @@ type Service struct { riverDrainingClients map[asyncExecutionClient]struct{} riverWorkerCapacity int workerInstanceID string + workerLoad *workerload.Controller + workerLoadSampler workerLoadSampler asyncCapacityLoader func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) admissionWakeMu sync.Mutex admissionWake chan struct{} @@ -129,6 +132,9 @@ func NewWithStores( if cfg.AsyncWorkerRefreshIntervalSeconds == 0 { cfg.AsyncWorkerRefreshIntervalSeconds = 5 } + if strings.TrimSpace(cfg.AsyncWorkerLoadMode) == "" { + cfg.AsyncWorkerLoadMode = workerload.ModeAdaptive + } if cfg.MediaMaterializationConcurrency == 0 { cfg.MediaMaterializationConcurrency = 8 } @@ -184,8 +190,13 @@ func NewWithStores( "universal": clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor}, "simulation": clients.SimulationClient{}, }, - httpClients: httpClients, - workerInstanceID: asyncWorkerID(), + httpClients: httpClients, + workerInstanceID: asyncWorkerID(), + workerLoad: workerload.New(workerload.Config{ + Mode: cfg.AsyncWorkerLoadMode, HardLimit: cfg.AsyncWorkerInstanceHardLimit, + InitialActive: 4, InitialHeavy: 1, HealthySamples: 3, + }), + workerLoadSampler: workerload.NewSystemSampler(), admissionWake: make(chan struct{}, 4096), asyncAdmissionWake: make(chan struct{}, 1), admissionTaskWaiters: map[string]*admissionTaskWaiter{}, @@ -1266,7 +1277,7 @@ func (s *Service) runCandidate( } simulated := isSimulation(task, candidate) baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing)) - reservations := s.rateLimitReservations(ctx, user, candidate, body) + reservations := acceptanceInfrastructureReservations(task, s.rateLimitReservations(ctx, user, candidate, body)) if admittedPlatformModelID == candidate.PlatformModelID && len(admittedLeases) > 0 { filtered := make([]store.RateLimitReservation, 0, len(reservations)) for _, reservation := range reservations { @@ -1278,6 +1289,10 @@ func (s *Service) runCandidate( } limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, "", reservations) if err != nil { + var limitErr *store.RateLimitExceededError + if errors.As(err, &limitErr) { + s.observeProviderQuotaWait(limitErr.Metric) + } retryable := store.RateLimitRetryable(err) clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable} return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)} @@ -1438,6 +1453,9 @@ func (s *Service) runCandidate( ); err != nil { return clients.Response{}, fmt.Errorf("restore upstream submission status: %w", err) } + if err := enterWorkerWaiting(ctx); err != nil { + return clients.Response{}, err + } } setSubmissionStatus := func(status string) error { if submissionStatus == "response_received" && status != "response_received" { @@ -1484,7 +1502,10 @@ func (s *Service) runCandidate( if err := s.persistCompatibilitySubmission(context.WithoutCancel(ctx), task, candidate, remoteTaskID, checkpoint, submissionWire); err != nil { return err } - return setSubmissionStatus("response_received") + if err := setSubmissionStatus("response_received"); err != nil { + return err + } + return enterWorkerWaiting(ctx) }, OnRemoteTaskPolled: func(remoteTaskID string, payload map[string]any) error { if strings.TrimSpace(remoteTaskID) == "" { @@ -1496,18 +1517,21 @@ func (s *Service) runCandidate( } task.RemoteTaskID = remoteTaskID task.RemoteTaskPayload = checkpoint - return setSubmissionStatus("response_received") + if err := setSubmissionStatus("response_received"); err != nil { + return err + } + return enterWorkerWaiting(ctx) }, OnUpstreamSubmissionStarted: func() error { if err := setSubmissionStatus("submitting"); err != nil { return err } markUpstreamSubmissionStarted(ctx) - return nil + return enterWorkerWaiting(ctx) }, OnUpstreamResponseReceived: func() error { submissionStatus = "response_received" - return nil + return enterWorkerFinalizing(ctx) }, OnUpstreamWireResponse: func(wire *clients.WireResponse) error { submissionWire = wire @@ -1521,6 +1545,9 @@ func (s *Service) runCandidate( UpstreamPreviousResponseID: responseExecution.UpstreamPreviousResponseID, PreviousResponseTurns: responseExecution.PreviousTurns, }) + if phaseErr := enterWorkerFinalizing(runCtx); err == nil && phaseErr != nil { + err = phaseErr + } if leaseErr := stopLeaseRenewal(); leaseErr != nil { err = &clients.ClientError{ Code: "concurrency_lease_lost", @@ -1717,6 +1744,15 @@ func (s *Service) runCandidate( return response, nil } +func (s *Service) observeProviderQuotaWait(metric string) { + observer, ok := s.billingMetrics.(interface { + ObserveProviderQuotaWait(string) + }) + if ok { + observer.ObserveProviderQuotaWait(metric) + } +} + func minimalRemoteTaskCheckpoint(provider string, specType string, payload map[string]any) map[string]any { const maxBytes = 8192 provider = strings.ToLower(strings.TrimSpace(provider)) diff --git a/apps/api/internal/runner/worker_load.go b/apps/api/internal/runner/worker_load.go new file mode 100644 index 0000000..1fd035a --- /dev/null +++ b/apps/api/internal/runner/worker_load.go @@ -0,0 +1,28 @@ +package runner + +import ( + "context" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload" +) + +type workerLoadLeaseContextKey struct{} + +func workerLoadLease(ctx context.Context) *workerload.Lease { + lease, _ := ctx.Value(workerLoadLeaseContextKey{}).(*workerload.Lease) + return lease +} + +func enterWorkerWaiting(ctx context.Context) error { + if lease := workerLoadLease(ctx); lease != nil { + return lease.EnterWaiting() + } + return nil +} + +func enterWorkerFinalizing(ctx context.Context) error { + if lease := workerLoadLease(ctx); lease != nil { + return lease.EnterFinalizing(ctx) + } + return nil +} diff --git a/apps/api/internal/securityevents/metrics.go b/apps/api/internal/securityevents/metrics.go index 3fafe7d..0c4bc5b 100644 --- a/apps/api/internal/securityevents/metrics.go +++ b/apps/api/internal/securityevents/metrics.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "net/http" + "strings" "sync/atomic" "time" @@ -22,53 +23,64 @@ type DynamicMetricsSnapshotProvider interface { } type Metrics struct { - accepted atomic.Uint64 - rejected atomic.Uint64 - duplicate atomic.Uint64 - sessionsDeleted atomic.Uint64 - watermarkRejected atomic.Uint64 - verificationAccepted atomic.Uint64 - heartbeatAccepted atomic.Uint64 - heartbeatFailed atomic.Uint64 - introspectionActive atomic.Uint64 - introspectionInactive atomic.Uint64 - introspectionFailed atomic.Uint64 - jwksSSFFailed atomic.Uint64 - jwksOIDCFailed atomic.Uint64 - processingCount atomic.Uint64 - processingNanos atomic.Uint64 - processingBuckets [6]atomic.Uint64 - billingSettlementCompleted atomic.Uint64 - billingSettlementRetry atomic.Uint64 - billingManualReview atomic.Uint64 - billingEstimateFailed atomic.Uint64 - billingIdempotentReplay atomic.Uint64 - billingPricingUnavailable atomic.Uint64 - asyncWorkerCapacity atomic.Int64 - asyncWorkerDesiredCapacity atomic.Int64 - asyncWorkerHardLimit atomic.Int64 - asyncWorkerCapacityCapped atomic.Int64 - asyncWorkerActiveInstances atomic.Int64 - asyncWorkerGlobalCapacity atomic.Int64 - asyncWorkerAllocated atomic.Int64 - asyncWorkerResizeSuccess atomic.Uint64 - asyncWorkerRefreshFailed atomic.Uint64 - asyncWorkerCreateFailed atomic.Uint64 - asyncWorkerStartFailed atomic.Uint64 - leaseRenewalSuccess atomic.Uint64 - leaseRenewalFailure atomic.Uint64 - leaseRenewalLost atomic.Uint64 - taskEventDuplicate atomic.Uint64 - taskEventUnknownType atomic.Uint64 - taskEventBudgetExceeded atomic.Uint64 - taskAdmissionAdmitted atomic.Uint64 - taskAdmissionQueueFull atomic.Uint64 - taskAdmissionTimeout atomic.Uint64 - taskAdmissionCancelled atomic.Uint64 - taskAdmissionExpired atomic.Uint64 - taskAdmissionMigrated atomic.Uint64 - taskAdmissionWaitBuckets [11]atomic.Uint64 - taskAdmissionWaitMicros atomic.Uint64 + accepted atomic.Uint64 + rejected atomic.Uint64 + duplicate atomic.Uint64 + sessionsDeleted atomic.Uint64 + watermarkRejected atomic.Uint64 + verificationAccepted atomic.Uint64 + heartbeatAccepted atomic.Uint64 + heartbeatFailed atomic.Uint64 + introspectionActive atomic.Uint64 + introspectionInactive atomic.Uint64 + introspectionFailed atomic.Uint64 + jwksSSFFailed atomic.Uint64 + jwksOIDCFailed atomic.Uint64 + processingCount atomic.Uint64 + processingNanos atomic.Uint64 + processingBuckets [6]atomic.Uint64 + billingSettlementCompleted atomic.Uint64 + billingSettlementRetry atomic.Uint64 + billingManualReview atomic.Uint64 + billingEstimateFailed atomic.Uint64 + billingIdempotentReplay atomic.Uint64 + billingPricingUnavailable atomic.Uint64 + asyncWorkerCapacity atomic.Int64 + asyncWorkerDesiredCapacity atomic.Int64 + asyncWorkerHardLimit atomic.Int64 + asyncWorkerCapacityCapped atomic.Int64 + asyncWorkerActiveInstances atomic.Int64 + asyncWorkerGlobalCapacity atomic.Int64 + asyncWorkerAllocated atomic.Int64 + workerSafeCapacity atomic.Int64 + workerHeavyCapacity atomic.Int64 + workerActiveTasks atomic.Int64 + workerPreparingTasks atomic.Int64 + workerWaitingTasks atomic.Int64 + workerFinalizingTasks atomic.Int64 + workerPressureState atomic.Int64 + providerQuotaWaitRPM atomic.Uint64 + providerQuotaWaitTPM atomic.Uint64 + providerQuotaWaitConcurrent atomic.Uint64 + providerQuotaWaitOther atomic.Uint64 + asyncWorkerResizeSuccess atomic.Uint64 + asyncWorkerRefreshFailed atomic.Uint64 + asyncWorkerCreateFailed atomic.Uint64 + asyncWorkerStartFailed atomic.Uint64 + leaseRenewalSuccess atomic.Uint64 + leaseRenewalFailure atomic.Uint64 + leaseRenewalLost atomic.Uint64 + taskEventDuplicate atomic.Uint64 + taskEventUnknownType atomic.Uint64 + taskEventBudgetExceeded atomic.Uint64 + taskAdmissionAdmitted atomic.Uint64 + taskAdmissionQueueFull atomic.Uint64 + taskAdmissionTimeout atomic.Uint64 + taskAdmissionCancelled atomic.Uint64 + taskAdmissionExpired atomic.Uint64 + taskAdmissionMigrated atomic.Uint64 + taskAdmissionWaitBuckets [11]atomic.Uint64 + taskAdmissionWaitMicros atomic.Uint64 } var processingDurationBounds = [...]time.Duration{ @@ -178,6 +190,38 @@ func (m *Metrics) SetDistributedWorkerCapacity(activeInstances, globalCapacity, m.asyncWorkerAllocated.Store(int64(allocatedCapacity)) } +func (m *Metrics) SetWorkerLoad(activeLimit, heavyLimit, active, preparing, waiting, finalizing int, pressure string) { + m.workerSafeCapacity.Store(int64(activeLimit)) + m.workerHeavyCapacity.Store(int64(heavyLimit)) + m.workerActiveTasks.Store(int64(active)) + m.workerPreparingTasks.Store(int64(preparing)) + m.workerWaitingTasks.Store(int64(waiting)) + m.workerFinalizingTasks.Store(int64(finalizing)) + state := int64(-1) + switch pressure { + case "normal": + state = 0 + case "busy": + state = 1 + case "critical": + state = 2 + } + m.workerPressureState.Store(state) +} + +func (m *Metrics) ObserveProviderQuotaWait(metric string) { + switch metric { + case "rpm": + m.providerQuotaWaitRPM.Add(1) + case "tpm", "tpm_total", "tpm_input", "tpm_output": + m.providerQuotaWaitTPM.Add(1) + case "concurrent": + m.providerQuotaWaitConcurrent.Add(1) + default: + m.providerQuotaWaitOther.Add(1) + } +} + func (m *Metrics) ObserveTaskAdmission(event string) { switch event { case "admitted": @@ -297,6 +341,17 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str }); ok { riverPostgresPool = poolProvider.RiverPostgresPoolMetrics() } + modelRateLimits := []store.ModelRateLimitStatus{} + if rateLimitProvider, ok := provider.(interface { + ListModelRateLimitStatuses(context.Context) ([]store.ModelRateLimitStatus, error) + }); ok { + var err error + modelRateLimits, err = rateLimitProvider.ListModelRateLimitStatuses(r.Context()) + if err != nil { + http.Error(w, "metrics unavailable", http.StatusServiceUnavailable) + return + } + } w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8") outcomeCounters(w, "easyai_gateway_ssf_receipts_total", "Received SETs by bounded outcome.", []outcomeValue{ {"accepted", m.accepted.Load()}, {"rejected", m.rejected.Load()}, {"duplicate", m.duplicate.Load()}, @@ -368,6 +423,20 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str plainGauge(w, "easyai_gateway_worker_active_instances", "Active distributed worker instances with a fresh heartbeat.", m.asyncWorkerActiveInstances.Load()) plainGauge(w, "easyai_gateway_worker_global_capacity", "Global asynchronous execution capacity before instance allocation.", m.asyncWorkerGlobalCapacity.Load()) plainGauge(w, "easyai_gateway_worker_allocated_capacity", "Asynchronous execution capacity allocated to this worker instance.", m.asyncWorkerAllocated.Load()) + plainGauge(w, "easyai_gateway_worker_safe_capacity", "Locally resource-safe asynchronous task capacity.", m.workerSafeCapacity.Load()) + plainGauge(w, "easyai_gateway_worker_heavy_capacity", "Locally resource-safe preparing and finalizing capacity.", m.workerHeavyCapacity.Load()) + plainGauge(w, "easyai_gateway_worker_active_tasks", "Tasks currently owned by this Worker process.", m.workerActiveTasks.Load()) + plainGauge(w, "easyai_gateway_worker_preparing_tasks", "Tasks in the local preparing phase.", m.workerPreparingTasks.Load()) + plainGauge(w, "easyai_gateway_worker_waiting_upstream_tasks", "Tasks waiting for an upstream result.", m.workerWaitingTasks.Load()) + plainGauge(w, "easyai_gateway_worker_finalizing_tasks", "Tasks in the local finalizing phase.", m.workerFinalizingTasks.Load()) + plainGauge(w, "easyai_gateway_worker_pressure_state", "Local Worker pressure state: -1 unknown, 0 normal, 1 busy, 2 critical.", m.workerPressureState.Load()) + outcomeCounters(w, "easyai_gateway_provider_quota_waits_total", "Tasks delayed before an upstream call by a cluster-wide provider quota.", []outcomeValue{ + {"rpm", m.providerQuotaWaitRPM.Load()}, + {"tpm", m.providerQuotaWaitTPM.Load()}, + {"concurrent", m.providerQuotaWaitConcurrent.Load()}, + {"other", m.providerQuotaWaitOther.Load()}, + }) + platformModelRateLimitUtilizationGauges(w, modelRateLimits) plainGauge(w, "easyai_gateway_postgres_pool_max_connections", "Maximum PostgreSQL connections in this process pool.", int64(postgresPool.MaxConnections)) plainGauge(w, "easyai_gateway_postgres_pool_total_connections", "Current PostgreSQL connections in this process pool.", int64(postgresPool.TotalConnections)) plainGauge(w, "easyai_gateway_postgres_pool_acquired_connections", "Currently acquired PostgreSQL connections in this process pool.", int64(postgresPool.AcquiredConnections)) @@ -463,6 +532,28 @@ func plainFloatGauge(w http.ResponseWriter, name, help string, value float64) { fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s gauge\n%s %.6f\n", name, help, name, name, value) } +func platformModelRateLimitUtilizationGauges(w http.ResponseWriter, statuses []store.ModelRateLimitStatus) { + const name = "easyai_gateway_platform_model_rate_limit_utilization" + fmt.Fprintf(w, "# HELP %s Cluster-wide platform model quota utilization.\n# TYPE %s gauge\n", name, name) + for _, status := range statuses { + platformModelID := escapePrometheusLabel(status.PlatformModelID) + for _, metric := range []struct { + name string + ratio float64 + }{ + {name: "rpm", ratio: status.RPM.Ratio}, + {name: "tpm", ratio: status.TPM.Ratio}, + {name: "concurrent", ratio: status.Concurrent.Ratio}, + } { + fmt.Fprintf(w, "%s{platform_model_id=\"%s\",metric=\"%s\"} %.6f\n", name, platformModelID, metric.name, metric.ratio) + } + } +} + +func escapePrometheusLabel(value string) string { + return strings.NewReplacer("\\", "\\\\", "\n", "\\n", "\"", "\\\"").Replace(value) +} + func taskAdmissionWaitHistogram(w http.ResponseWriter, metrics *Metrics) { const name = "easyai_gateway_task_admission_wait_seconds" fmt.Fprintf(w, "# HELP %s Time spent waiting for persistent task admission.\n# TYPE %s histogram\n", name, name) diff --git a/apps/api/internal/securityevents/metrics_test.go b/apps/api/internal/securityevents/metrics_test.go index fac43a1..8d5a324 100644 --- a/apps/api/internal/securityevents/metrics_test.go +++ b/apps/api/internal/securityevents/metrics_test.go @@ -25,6 +25,15 @@ func (m metricsSnapshot) PostgresPoolMetrics() store.PostgresPoolMetricsSnapshot return m.pool } +func (m metricsSnapshot) ListModelRateLimitStatuses(context.Context) ([]store.ModelRateLimitStatus, error) { + return []store.ModelRateLimitStatus{{ + PlatformModelID: "platform-model-1", + RPM: store.RateLimitMetricStatus{Ratio: .25}, + TPM: store.RateLimitMetricStatus{Ratio: .5}, + Concurrent: store.RateLimitMetricStatus{Ratio: .75}, + }}, nil +} + func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { metrics := &Metrics{} metrics.ObserveReceipt("accepted", 20*time.Millisecond, 2) @@ -34,6 +43,9 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { metrics.ObserveIntrospection("failed") metrics.ObserveJWKSRefreshFailure("ssf") metrics.SetAsyncWorkerCapacity(96, 128, 96, true) + metrics.SetWorkerLoad(12, 3, 8, 2, 5, 1, "busy") + metrics.ObserveProviderQuotaWait("rpm") + metrics.ObserveProviderQuotaWait("tpm_total") metrics.ObserveAsyncWorkerResize("success") metrics.ObserveConcurrencyLeaseRenewal("success") metrics.ObserveConcurrencyLeaseRenewal("lost") @@ -65,6 +77,16 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { `easyai_gateway_async_worker_capacity 96`, `easyai_gateway_async_worker_desired_capacity 128`, `easyai_gateway_async_worker_capacity_capped 1`, + `easyai_gateway_worker_safe_capacity 12`, + `easyai_gateway_worker_heavy_capacity 3`, + `easyai_gateway_worker_active_tasks 8`, + `easyai_gateway_worker_preparing_tasks 2`, + `easyai_gateway_worker_waiting_upstream_tasks 5`, + `easyai_gateway_worker_finalizing_tasks 1`, + `easyai_gateway_worker_pressure_state 1`, + `easyai_gateway_provider_quota_waits_total{outcome="rpm"} 1`, + `easyai_gateway_provider_quota_waits_total{outcome="tpm"} 1`, + `easyai_gateway_platform_model_rate_limit_utilization{platform_model_id="platform-model-1",metric="concurrent"} 0.750000`, `easyai_gateway_async_worker_resizes_total{outcome="success"} 1`, `easyai_gateway_concurrency_lease_renewals_total{outcome="success"} 1`, `easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`, diff --git a/apps/api/internal/store/acceptance.go b/apps/api/internal/store/acceptance.go index c5a50c6..321bc5e 100644 --- a/apps/api/internal/store/acceptance.go +++ b/apps/api/internal/store/acceptance.go @@ -689,7 +689,7 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, runID); err != nil { } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases lease -SET released_at = now() +SET released_at = statement_timestamp() FROM gateway_tasks task WHERE lease.task_id = task.id AND lease.released_at IS NULL diff --git a/apps/api/internal/store/admission_queue.go b/apps/api/internal/store/admission_queue.go index 96c15de..aed5945 100644 --- a/apps/api/internal/store/admission_queue.go +++ b/apps/api/internal/store/admission_queue.go @@ -747,14 +747,14 @@ SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL - AND expires_at > now()`, taskID).Scan(&count) + AND expires_at > statement_timestamp()`, taskID).Scan(&count) return count, err } func resetTaskAdmissionToWaitingTx(ctx context.Context, tx pgx.Tx, input TaskAdmissionInput) (TaskAdmission, error) { if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, input.TaskID); err != nil { return TaskAdmission{}, err @@ -862,7 +862,7 @@ func (s *Store) deleteTaskAdmissionOnce(ctx context.Context, taskID string) erro } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, taskID); err != nil { return err @@ -884,7 +884,7 @@ SELECT id::text, FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL - AND expires_at > now() + AND expires_at > statement_timestamp() ORDER BY scope_type, scope_key, id`, taskID) if err != nil { return nil, err @@ -924,7 +924,7 @@ WHERE admission.mode = 'async' FROM gateway_concurrency_leases lease WHERE lease.task_id = admission.task_id AND lease.released_at IS NULL - AND lease.expires_at > now() + AND lease.expires_at > statement_timestamp() ) ) ) @@ -980,7 +980,7 @@ WHERE admission.mode = 'async' FROM gateway_concurrency_leases lease WHERE lease.task_id = admission.task_id AND lease.released_at IS NULL - AND lease.expires_at > now() + AND lease.expires_at > statement_timestamp() ) ) ) @@ -1006,6 +1006,29 @@ LIMIT $1`, limit) return admissions, rows.Err() } +// RequestWaitingTaskAdmissionReselect marks every queued task bound to a +// saturated platform model so the dispatcher can route the next batch to a +// different eligible candidate. The durable marker lets multiple dispatchers +// observe the same decision without assigning tasks to a specific Worker. +func (s *Store) RequestWaitingTaskAdmissionReselect(ctx context.Context, platformModelID string) (int64, error) { + result, err := s.pool.Exec(ctx, ` +UPDATE gateway_task_admissions admission +SET reselect_requested_at = now(), + updated_at = now() +FROM gateway_tasks task +WHERE admission.task_id = task.id + AND admission.platform_model_id = $1::uuid + AND admission.mode = 'async' + AND admission.status = 'waiting' + AND task.status = 'queued' + AND task.next_run_at <= now() + AND admission.reselect_requested_at IS NULL`, platformModelID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + // ListWaitingTaskAdmissionIDs returns one FIFO leader for every independent // platform-model queue, additionally requiring the task to lead its user-group // queue when one exists. It is used after capacity is released so API @@ -1268,12 +1291,12 @@ func admissionScopeStatesTx(ctx context.Context, tx pgx.Tx, scopes []AdmissionSc if scope.ConcurrentLimit > 0 { if err := tx.QueryRow(ctx, ` SELECT COALESCE(SUM(lease_value), 0)::float8, - COALESCE(MIN(expires_at), now() + interval '1 second') + COALESCE(MIN(expires_at), statement_timestamp() + interval '1 second') FROM gateway_concurrency_leases WHERE scope_type = $1 AND scope_key = $2 AND released_at IS NULL - AND expires_at > now()`, scope.ScopeType, scope.ScopeKey).Scan(&state.Active, &state.NextLeaseExpiration); err != nil { + AND expires_at > statement_timestamp()`, scope.ScopeType, scope.ScopeKey).Scan(&state.Active, &state.NextLeaseExpiration); err != nil { return nil, err } state.Saturated = state.Active+scope.Amount > scope.ConcurrentLimit diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index bc96041..ce80ae3 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -451,6 +451,20 @@ WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID) if listedSnapshot == nil || len(listedSnapshot.Scopes) != len(queuedAdmission.Scopes) { t.Fatalf("listed admission snapshot=%+v, want %d scopes", listedSnapshot, len(queuedAdmission.Scopes)) } + markedForReselect, err := first.RequestWaitingTaskAdmissionReselect(ctx, platformModelID) + if err != nil { + t.Fatalf("request waiting admission reselection: %v", err) + } + if markedForReselect < 1 { + t.Fatalf("marked admissions=%d, want at least the queued task", markedForReselect) + } + reselectAdmission, err := first.GetTaskAdmission(ctx, queuedAtomicTask.ID) + if err != nil { + t.Fatalf("read admission reselection marker: %v", err) + } + if reselectAdmission.ReselectRequestedAt.IsZero() { + t.Fatal("queued admission was not marked for candidate reselection") + } var riverJobID int64 if err := first.pool.QueryRow(ctx, ` SELECT @@ -1296,4 +1310,35 @@ WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).Str if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 { t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err) } + first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{ + InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 10, + LoadMode: "adaptive", SafeCapacity: 0, HeavyCapacity: 1, + ActiveTasks: 2, WaitingUpstreamTasks: 2, + PressureState: "critical", PressureReason: "memory", + }) + if err != nil || first.Allocated != 0 { + t.Fatalf("critical worker allocation = %+v, err=%v", first, err) + } + second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{ + InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 10, + LoadMode: "adaptive", SafeCapacity: 5, HeavyCapacity: 2, + ActiveTasks: 3, PreparingTasks: 1, WaitingUpstreamTasks: 1, FinalizingTasks: 1, + PressureState: "normal", + }) + if err != nil || second.Allocated != 5 || second.GlobalAllocated != 5 { + t.Fatalf("adaptive redistribution allocation = %+v, err=%v", second, err) + } + instances, err := db.ListWorkerInstanceRuntime(ctx) + if err != nil { + t.Fatalf("list adaptive worker runtime: %v", err) + } + foundCritical := false + for _, instance := range instances { + if instance.InstanceID == firstID { + foundCritical = instance.SafeCapacity == 0 && instance.ReportedActiveTasks == 2 && instance.WaitingUpstreamTasks == 2 && instance.PressureState == "critical" + } + } + if !foundCritical { + t.Fatalf("adaptive runtime did not expose the critical worker: %+v", instances) + } } diff --git a/apps/api/internal/store/async_worker_capacity.go b/apps/api/internal/store/async_worker_capacity.go index 725d7ef..6a706a2 100644 --- a/apps/api/internal/store/async_worker_capacity.go +++ b/apps/api/internal/store/async_worker_capacity.go @@ -6,19 +6,28 @@ import ( ) type AsyncWorkerCapacitySnapshot struct { - Capacity int - GlobalCapacity int - Desired int - HardLimit int - Capped bool - EnabledModels int - UnlimitedModels int - EnabledGroups int - UnlimitedGroups int - ModelDesired int - GroupDesired int - ActiveInstances int - InstanceID string + Capacity int + GlobalCapacity int + Desired int + HardLimit int + Capped bool + EnabledModels int + UnlimitedModels int + EnabledGroups int + UnlimitedGroups int + ModelDesired int + GroupDesired int + ActiveInstances int + InstanceID string + LoadMode string + LocalSafeCapacity int + LocalHeavyCapacity int + LocalActiveTasks int + LocalPreparingTasks int + LocalWaitingTasks int + LocalFinalizingTasks int + LocalPressureState string + LocalPressureReason string } func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) { diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index d532cfb..9ad6200 100644 --- a/apps/api/internal/store/candidates.go +++ b/apps/api/internal/store/candidates.go @@ -100,7 +100,7 @@ LEFT JOIN ( FROM gateway_concurrency_leases WHERE scope_type = 'platform_model' AND released_at IS NULL - AND expires_at > now() + AND expires_at > statement_timestamp() GROUP BY scope_key ) con ON con.scope_key = m.id::text LEFT JOIN ( diff --git a/apps/api/internal/store/rate_limit_status.go b/apps/api/internal/store/rate_limit_status.go index 72caadd..b39308d 100644 --- a/apps/api/internal/store/rate_limit_status.go +++ b/apps/api/internal/store/rate_limit_status.go @@ -169,7 +169,7 @@ LEFT JOIN ( FROM gateway_concurrency_leases WHERE scope_type = 'platform_model' AND released_at IS NULL - AND expires_at > now() + AND expires_at > statement_timestamp() GROUP BY scope_key ) con ON con.scope_key = m.id::text LEFT JOIN ( diff --git a/apps/api/internal/store/rate_limits.go b/apps/api/internal/store/rate_limits.go index 9e9a12a..3cefc37 100644 --- a/apps/api/internal/store/rate_limits.go +++ b/apps/api/internal/store/rate_limits.go @@ -184,12 +184,12 @@ func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, atte var nextAvailableAt time.Time if err := tx.QueryRow(ctx, ` SELECT COALESCE(SUM(lease_value), 0)::float8, - COALESCE(MIN(expires_at), now() + ($3::int * interval '1 second')) + COALESCE(MIN(expires_at), statement_timestamp() + ($3::int * interval '1 second')) FROM gateway_concurrency_leases WHERE scope_type = $1 AND scope_key = $2 AND released_at IS NULL - AND expires_at > now()`, + AND expires_at > statement_timestamp()`, reservation.ScopeType, reservation.ScopeKey, reservation.LeaseTTLSeconds, @@ -218,14 +218,20 @@ WHERE scope_type = $1 } var leaseID string if err := tx.QueryRow(ctx, ` -INSERT INTO gateway_concurrency_leases (task_id, attempt_id, scope_type, scope_key, lease_value, expires_at) -VALUES ($1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, now() + ($6::int * interval '1 second')) +INSERT INTO gateway_concurrency_leases ( + task_id, attempt_id, scope_type, scope_key, lease_value, limit_value, acquired_at, expires_at +) +VALUES ( + $1::uuid, NULLIF($2, '')::uuid, $3, $4, $5, $6, + statement_timestamp(), statement_timestamp() + ($7::int * interval '1 second') +) RETURNING id::text`, taskID, attemptID, reservation.ScopeType, reservation.ScopeKey, reservation.Amount, + reservation.Limit, reservation.LeaseTTLSeconds, ).Scan(&leaseID); err != nil { return ConcurrencyLease{}, err @@ -381,7 +387,7 @@ func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []Concurren } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs); err != nil { return err } @@ -411,7 +417,7 @@ func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []Concurrency } tag, err := s.pool.Exec(ctx, ` UPDATE gateway_concurrency_leases lease -SET expires_at = now() + (renewal.ttl_seconds * interval '1 second') +SET expires_at = statement_timestamp() + (renewal.ttl_seconds * interval '1 second') FROM unnest($1::uuid[], $2::int[]) AS renewal(id, ttl_seconds) WHERE lease.id = renewal.id AND lease.released_at IS NULL @@ -705,7 +711,7 @@ WHERE attempt.task_id = task.id result.FailedAttempts = tag.RowsAffected() tag, err = tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL`, taskIDs) if err != nil { @@ -772,7 +778,7 @@ FOR UPDATE OF task SKIP LOCKED`, runtimeRecoveryBatchSize) var result RuntimeRecoveryResult tag, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL`, taskIDs) if err != nil { diff --git a/apps/api/internal/store/rate_limits_integration_test.go b/apps/api/internal/store/rate_limits_integration_test.go index 572d591..421f9ef 100644 --- a/apps/api/internal/store/rate_limits_integration_test.go +++ b/apps/api/internal/store/rate_limits_integration_test.go @@ -98,23 +98,168 @@ WHERE scope_type = 'platform_model' } var active int64 + var storedLimit float64 if err := first.Pool().QueryRow(ctx, ` -SELECT COUNT(*) +SELECT COUNT(*), COALESCE(MAX(limit_value), 0)::float8 FROM gateway_concurrency_leases WHERE scope_type = 'platform_model' AND scope_key = $1 AND released_at IS NULL - AND expires_at > now()`, scopeKey).Scan(&active); err != nil { + AND expires_at > now()`, scopeKey).Scan(&active, &storedLimit); err != nil { t.Fatalf("count active leases: %v", err) } - if successes.Load() != 64 || active != 64 { - t.Fatalf("successful reservations=%d active leases=%d, want exactly 64", successes.Load(), active) + if successes.Load() != 64 || active != 64 || storedLimit != 64 { + t.Fatalf("successful reservations=%d active leases=%d stored limit=%.0f, want exactly 64", successes.Load(), active, storedLimit) } if peak.Load() > 64 { t.Fatalf("active lease peak=%d, want <=64", peak.Load()) } } +func TestConcurrencyLeaseTimestampStartsAtReservationStatement(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run concurrency lease PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + database, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect store: %v", err) + } + defer database.Close() + + scopeKey := "statement-clock-" + time.Now().UTC().Format("20060102150405.000000000") + taskIDs := createLeaseTestTasks(t, ctx, database, 1, scopeKey) + defer deleteLeaseTestTasks(t, database, taskIDs) + + tx, err := database.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin reservation transaction: %v", err) + } + defer rollbackTransaction(tx) + var transactionStartedAt time.Time + if err := tx.QueryRow(ctx, `SELECT now()`).Scan(&transactionStartedAt); err != nil { + t.Fatalf("read transaction start: %v", err) + } + time.Sleep(1100 * time.Millisecond) + lease, err := reserveConcurrencyLease(ctx, tx, taskIDs[0], "", RateLimitReservation{ + ScopeType: "platform_model", + ScopeKey: scopeKey, + Metric: "concurrent", + Limit: 1, + Amount: 1, + LeaseTTLSeconds: 30, + }) + if err != nil { + t.Fatalf("reserve concurrency lease: %v", err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatalf("commit reservation transaction: %v", err) + } + + var acquiredAt, expiresAt time.Time + if err := database.pool.QueryRow(ctx, ` +SELECT acquired_at, expires_at +FROM gateway_concurrency_leases +WHERE id = $1::uuid`, lease.ID).Scan(&acquiredAt, &expiresAt); err != nil { + t.Fatalf("read lease timestamps: %v", err) + } + if elapsed := acquiredAt.Sub(transactionStartedAt); elapsed < time.Second { + t.Fatalf("lease acquired_at advanced by %s, want at least 1s after transaction start", elapsed) + } + if ttl := expiresAt.Sub(acquiredAt); ttl != 30*time.Second { + t.Fatalf("lease ttl=%s, want 30s", ttl) + } +} + +func TestCounterWindowReservationIsAtomicAcrossPools(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run rate limit PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) + defer cancel() + first, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect first store: %v", err) + } + defer first.Close() + second, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect second store: %v", err) + } + defer second.Close() + + tests := []struct { + metric string + limit float64 + amount float64 + wantSuccesses int64 + }{ + {metric: "rpm", limit: 37, amount: 1, wantSuccesses: 37}, + {metric: "tpm_total", limit: 1_000, amount: 25, wantSuccesses: 40}, + } + for _, test := range tests { + t.Run(test.metric, func(t *testing.T) { + scopeKey := "atomic-" + test.metric + "-" + time.Now().UTC().Format("20060102150405.000000000") + taskIDs := createLeaseTestTasks(t, ctx, first, 128, scopeKey) + defer deleteLeaseTestTasks(t, first, taskIDs) + + var successes atomic.Int64 + var wg sync.WaitGroup + errs := make(chan error, len(taskIDs)) + for index, taskID := range taskIDs { + wg.Add(1) + go func(index int, taskID string) { + defer wg.Done() + target := first + if index%2 == 1 { + target = second + } + _, err := target.ReserveRateLimits(ctx, taskID, "", []RateLimitReservation{{ + ScopeType: "platform_model", + ScopeKey: scopeKey, + Metric: test.metric, + Limit: test.limit, + Amount: test.amount, + WindowSeconds: 3600, + }}) + if err == nil { + successes.Add(1) + return + } + if !errors.Is(err, ErrRateLimited) { + errs <- err + } + }(index, taskID) + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatalf("unexpected reservation error: %v", err) + } + + var current float64 + if err := first.Pool().QueryRow(ctx, ` +SELECT COALESCE(MAX(used_value + reserved_value), 0)::float8 +FROM gateway_rate_limit_counters +WHERE scope_type = 'platform_model' + AND scope_key = $1 + AND metric = $2`, scopeKey, test.metric).Scan(¤t); err != nil { + t.Fatalf("read %s counter: %v", test.metric, err) + } + if successes.Load() != test.wantSuccesses { + t.Fatalf("successful %s reservations=%d, want exactly %d", test.metric, successes.Load(), test.wantSuccesses) + } + wantCurrent := float64(test.wantSuccesses) * test.amount + if current != wantCurrent || current > test.limit { + t.Fatalf("%s current=%.0f, want %.0f and <= %.0f", test.metric, current, wantCurrent, test.limit) + } + }) + } +} + func TestConcurrencyLeaseRenewalExtendsAndReleases(t *testing.T) { databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) if databaseURL == "" { diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index d3601c0..5064b6d 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -757,7 +757,7 @@ WHERE task_id = $1::uuid } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, taskID); err != nil { return err @@ -922,7 +922,7 @@ RETURNING `+gatewayTaskColumns, taskID, message)) changed = true if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, taskID); err != nil { return err @@ -1011,7 +1011,7 @@ WHERE task_id = $1::uuid } if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, taskID); err != nil { return err @@ -1720,7 +1720,7 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, if input.FinalizeAdmission { if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, input.TaskID); err != nil { return err @@ -2096,7 +2096,7 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON) if input.FinalizeAdmission { if _, err := tx.Exec(ctx, ` UPDATE gateway_concurrency_leases -SET released_at = now() +SET released_at = statement_timestamp() WHERE task_id = $1::uuid AND released_at IS NULL`, input.TaskID); err != nil { return err diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index eb3c896..fc399d2 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -12,14 +12,24 @@ import ( const workerHeartbeatStaleAfter = 30 * time.Second type WorkerRegistrationInput struct { - InstanceID string - PodUID string - PodName string - Site string - Revision string - DesiredCapacity int - CapacityLimit int - HeartbeatStaleAfter time.Duration + InstanceID string + PodUID string + PodName string + Site string + Revision string + DesiredCapacity int + CapacityLimit int + LoadMode string + SafeCapacity int + HeavyCapacity int + ActiveTasks int + PreparingTasks int + WaitingUpstreamTasks int + FinalizingTasks int + PressureState string + PressureReason string + LoadSampledAt time.Time + HeartbeatStaleAfter time.Duration } type WorkerAllocation struct { @@ -59,9 +69,28 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra if input.CapacityLimit < 0 { return WorkerAllocation{}, errors.New("worker capacity limit cannot be negative") } + if input.SafeCapacity < 0 || input.HeavyCapacity < 0 || input.ActiveTasks < 0 || input.PreparingTasks < 0 || input.WaitingUpstreamTasks < 0 || input.FinalizingTasks < 0 { + return WorkerAllocation{}, errors.New("worker load values cannot be negative") + } + if input.ActiveTasks != input.PreparingTasks+input.WaitingUpstreamTasks+input.FinalizingTasks { + return WorkerAllocation{}, errors.New("worker active task count must equal phase task counts") + } if input.CapacityLimit == 0 { input.CapacityLimit = input.DesiredCapacity } + hardCapacityLimit := input.CapacityLimit + if strings.EqualFold(strings.TrimSpace(input.LoadMode), "adaptive") { + input.CapacityLimit = min(input.CapacityLimit, input.SafeCapacity) + } + pressureState := strings.ToLower(strings.TrimSpace(input.PressureState)) + switch pressureState { + case "normal", "busy", "critical": + default: + pressureState = "unknown" + } + if input.LoadSampledAt.IsZero() { + input.LoadSampledAt = time.Now() + } staleAfter := input.HeartbeatStaleAfter if staleAfter < workerHeartbeatStaleAfter { staleAfter = workerHeartbeatStaleAfter @@ -83,9 +112,16 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra if _, err := tx.Exec(ctx, ` INSERT INTO gateway_worker_instances ( instance_id, pod_uid, pod_name, site, revision, status, - desired_capacity, capacity_limit, allocated_capacity, started_at, heartbeat_at, updated_at + desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity, + active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks, + pressure_state, pressure_reason, load_sampled_at, + allocated_capacity, started_at, heartbeat_at, updated_at +) +VALUES ( + $1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10, + $11, $12, $13, $14, $15, $16, $17, + 0, now(), now(), now() ) -VALUES ($1, $2, $3, $4, $5, 'active', $6, $7, 0, now(), now(), now()) ON CONFLICT (instance_id) DO UPDATE SET pod_uid = EXCLUDED.pod_uid, pod_name = EXCLUDED.pod_name, @@ -97,6 +133,16 @@ SET pod_uid = EXCLUDED.pod_uid, END, desired_capacity = EXCLUDED.desired_capacity, capacity_limit = EXCLUDED.capacity_limit, + hard_capacity_limit = EXCLUDED.hard_capacity_limit, + safe_capacity = EXCLUDED.safe_capacity, + heavy_capacity = EXCLUDED.heavy_capacity, + active_tasks = EXCLUDED.active_tasks, + preparing_tasks = EXCLUDED.preparing_tasks, + waiting_upstream_tasks = EXCLUDED.waiting_upstream_tasks, + finalizing_tasks = EXCLUDED.finalizing_tasks, + pressure_state = EXCLUDED.pressure_state, + pressure_reason = EXCLUDED.pressure_reason, + load_sampled_at = EXCLUDED.load_sampled_at, heartbeat_at = now(), updated_at = now()`, input.InstanceID, @@ -106,6 +152,16 @@ SET pod_uid = EXCLUDED.pod_uid, strings.TrimSpace(input.Revision), input.DesiredCapacity, input.CapacityLimit, + hardCapacityLimit, + input.SafeCapacity, + input.HeavyCapacity, + input.ActiveTasks, + input.PreparingTasks, + input.WaitingUpstreamTasks, + input.FinalizingTasks, + pressureState, + strings.TrimSpace(input.PressureReason), + input.LoadSampledAt, ); err != nil { return WorkerAllocation{}, err } @@ -198,7 +254,7 @@ func allocateWorkerCapacities(workers []activeWorkerCapacity, desired int) (map[ for _, worker := range workers { limit := worker.CapacityLimit if limit <= 0 { - limit = desired + continue } if allocations[worker.InstanceID] >= limit { continue @@ -254,24 +310,52 @@ WHERE instance_id = $1 } type WorkerInstanceRuntime struct { - InstanceID string - PodUID string - PodName string - Site string - Revision string - Status string - Allocated int - CapacityLimit int - RunningTasks int - ActiveLeases int - HeartbeatAt time.Time - DrainingAt *time.Time + InstanceID string `json:"instanceId"` + PodUID string `json:"podUid,omitempty"` + PodName string `json:"podName,omitempty"` + Site string `json:"site,omitempty"` + Revision string `json:"revision,omitempty"` + Status string `json:"status"` + Allocated int `json:"allocatedCapacity"` + CapacityLimit int `json:"capacityLimit"` + HardCapacityLimit int `json:"hardCapacityLimit"` + SafeCapacity int `json:"safeCapacity"` + HeavyCapacity int `json:"heavyCapacity"` + ReportedActiveTasks int `json:"reportedActiveTasks"` + PreparingTasks int `json:"preparingTasks"` + WaitingUpstreamTasks int `json:"waitingUpstreamTasks"` + FinalizingTasks int `json:"finalizingTasks"` + PressureState string `json:"pressureState"` + PressureReason string `json:"pressureReason,omitempty"` + LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"` + RunningTasks int `json:"runningTasks"` + ActiveLeases int `json:"activeLeases"` + HeartbeatAt time.Time `json:"heartbeatAt"` + DrainingAt *time.Time `json:"drainingAt,omitempty"` } type WorkerQueueRuntime struct { - Queued int - Running int - OldestWaitSeconds float64 + Queued int `json:"queued"` + Running int `json:"running"` + OldestWaitSeconds float64 `json:"oldestWaitSeconds"` +} + +type WorkerClusterRuntime struct { + Workers []WorkerInstanceRuntime `json:"workers"` + Queue WorkerQueueRuntime `json:"queue"` + CapturedAt time.Time `json:"capturedAt"` +} + +func (s *Store) GetWorkerClusterRuntime(ctx context.Context) (WorkerClusterRuntime, error) { + workers, err := s.ListWorkerInstanceRuntime(ctx) + if err != nil { + return WorkerClusterRuntime{}, err + } + queue, err := s.WorkerQueueRuntime(ctx) + if err != nil { + return WorkerClusterRuntime{}, err + } + return WorkerClusterRuntime{Workers: workers, Queue: queue, CapturedAt: time.Now()}, nil } type CapacityDatabaseHealth struct { @@ -306,6 +390,16 @@ SELECT worker.instance_id, worker.status, worker.allocated_capacity, worker.capacity_limit, + worker.hard_capacity_limit, + worker.safe_capacity, + worker.heavy_capacity, + worker.active_tasks, + worker.preparing_tasks, + worker.waiting_upstream_tasks, + worker.finalizing_tasks, + worker.pressure_state, + worker.pressure_reason, + worker.load_sampled_at, count(DISTINCT task.id) FILTER (WHERE task.status = 'running')::int, count(DISTINCT lease.id) FILTER (WHERE lease.released_at IS NULL)::int, worker.heartbeat_at, @@ -339,6 +433,16 @@ ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`, &instance.Status, &instance.Allocated, &instance.CapacityLimit, + &instance.HardCapacityLimit, + &instance.SafeCapacity, + &instance.HeavyCapacity, + &instance.ReportedActiveTasks, + &instance.PreparingTasks, + &instance.WaitingUpstreamTasks, + &instance.FinalizingTasks, + &instance.PressureState, + &instance.PressureReason, + &instance.LoadSampledAt, &instance.RunningTasks, &instance.ActiveLeases, &instance.HeartbeatAt, @@ -476,7 +580,7 @@ WITH orphaned AS MATERIALIZED ( ), released_leases AS ( UPDATE gateway_concurrency_leases lease - SET released_at = now() + SET released_at = statement_timestamp() FROM orphaned WHERE lease.task_id = orphaned.task_id AND lease.released_at IS NULL diff --git a/apps/api/internal/store/worker_registry_test.go b/apps/api/internal/store/worker_registry_test.go index 73efa33..1aa125d 100644 --- a/apps/api/internal/store/worker_registry_test.go +++ b/apps/api/internal/store/worker_registry_test.go @@ -42,3 +42,15 @@ func TestAllocateWorkerCapacitiesSupportsUnequalLimits(t *testing.T) { t.Fatalf("allocations=%v global=%d, want 1/4 and 5", allocations, global) } } + +func TestAllocateWorkerCapacitiesRedistributesFromPressuredWorker(t *testing.T) { + workers := []activeWorkerCapacity{ + {InstanceID: "worker-a", CapacityLimit: 0}, + {InstanceID: "worker-b", CapacityLimit: 6}, + {InstanceID: "worker-c", CapacityLimit: 6}, + } + allocations, global := allocateWorkerCapacities(workers, 8) + if global != 8 || allocations["worker-a"] != 0 || allocations["worker-b"] != 4 || allocations["worker-c"] != 4 { + t.Fatalf("allocations=%v global=%d, want 0/4/4 and 8", allocations, global) + } +} diff --git a/apps/api/internal/workerload/controller.go b/apps/api/internal/workerload/controller.go new file mode 100644 index 0000000..9d85f56 --- /dev/null +++ b/apps/api/internal/workerload/controller.go @@ -0,0 +1,381 @@ +package workerload + +import ( + "context" + "errors" + "strings" + "sync" + "time" +) + +type Phase string + +const ( + PhasePreparing Phase = "preparing" + PhaseWaitingUpstream Phase = "waiting_upstream" + PhaseFinalizing Phase = "finalizing" +) + +type PressureState string + +const ( + PressureNormal PressureState = "normal" + PressureBusy PressureState = "busy" + PressureCritical PressureState = "critical" +) + +const ( + ModeAdaptive = "adaptive" + ModeLegacy = "legacy" +) + +var ErrReleased = errors.New("worker load lease already released") + +type Config struct { + Mode string + HardLimit int + InitialActive int + InitialHeavy int + HealthySamples int +} + +type ResourceSample struct { + MemoryCurrentBytes int64 + MemoryLimitBytes int64 + CPUUtilization float64 + CPUThrottled bool + DBConnections int32 + DBMaxConnections int32 + SampledAt time.Time +} + +type Snapshot struct { + Mode string + ActiveLimit int + HeavyLimit int + ClaimLimit int + SafeCapacity int + ActiveTasks int + PreparingTasks int + WaitingUpstreamTasks int + FinalizingTasks int + PressureState PressureState + PressureReason string + MemoryUtilization float64 + CPUUtilization float64 + DBUtilization float64 + SampledAt time.Time +} + +type Controller struct { + mu sync.Mutex + + mode string + hardLimit int + activeLimit int + heavyLimit int + claimLimit int + healthySamples int + healthyCount int + + preparing int + waiting int + finalizing int + + last Snapshot + wake chan struct{} +} + +type Lease struct { + controller *Controller + phase Phase + released bool +} + +func New(config Config) *Controller { + mode := strings.ToLower(strings.TrimSpace(config.Mode)) + if mode != ModeLegacy { + mode = ModeAdaptive + } + if config.HardLimit < 1 { + config.HardLimit = 1 + } + if config.InitialActive < 1 { + config.InitialActive = 4 + } + if config.InitialHeavy < 1 { + config.InitialHeavy = 1 + } + if config.HealthySamples < 1 { + config.HealthySamples = 3 + } + active := min(config.InitialActive, config.HardLimit) + heavy := min(config.InitialHeavy, active) + if mode == ModeLegacy { + active = config.HardLimit + heavy = config.HardLimit + } + controller := &Controller{ + mode: mode, hardLimit: config.HardLimit, + activeLimit: active, heavyLimit: heavy, claimLimit: active, + healthySamples: config.HealthySamples, + wake: make(chan struct{}, 1), + } + controller.last = controller.snapshotLocked(ResourceSample{SampledAt: time.Now()}) + return controller +} + +func (c *Controller) Observe(sample ResourceSample) Snapshot { + c.mu.Lock() + defer c.mu.Unlock() + if sample.SampledAt.IsZero() { + sample.SampledAt = time.Now() + } + memory := utilization(sample.MemoryCurrentBytes, sample.MemoryLimitBytes) + database := utilization(int64(sample.DBConnections), int64(sample.DBMaxConnections)) + cpu := bounded(sample.CPUUtilization) + state, reason := pressure(memory, cpu, database, sample.CPUThrottled) + if c.mode == ModeLegacy { + c.activeLimit = c.hardLimit + c.heavyLimit = c.hardLimit + state = PressureNormal + reason = "legacy" + } else { + c.adjustLocked(state, memory, cpu, database) + } + c.last = c.snapshotLocked(sample) + c.last.PressureState = state + c.last.PressureReason = reason + c.last.MemoryUtilization = memory + c.last.CPUUtilization = cpu + c.last.DBUtilization = database + c.last.SafeCapacity = c.activeLimit + if state == PressureCritical { + c.last.SafeCapacity = 0 + } + c.signalLocked() + return c.last +} + +func (c *Controller) SetClaimLimit(limit int) Snapshot { + c.mu.Lock() + defer c.mu.Unlock() + if limit < 0 { + limit = 0 + } + c.claimLimit = min(limit, c.hardLimit) + c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) + c.signalLocked() + return c.last +} + +func (c *Controller) Snapshot() Snapshot { + c.mu.Lock() + defer c.mu.Unlock() + return c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) +} + +func (c *Controller) TryStart() (*Lease, bool) { + c.mu.Lock() + defer c.mu.Unlock() + activeLimit := min(c.activeLimit, c.claimLimit) + if activeLimit <= 0 || c.activeLocked() >= activeLimit || c.preparing+c.finalizing >= c.heavyLimit { + return nil, false + } + c.preparing++ + c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) + return &Lease{controller: c, phase: PhasePreparing}, true +} + +func (l *Lease) EnterWaiting() error { + if l == nil || l.controller == nil { + return nil + } + c := l.controller + c.mu.Lock() + defer c.mu.Unlock() + if l.released { + return ErrReleased + } + if l.phase == PhaseWaitingUpstream { + return nil + } + c.decrementPhaseLocked(l.phase) + c.waiting++ + l.phase = PhaseWaitingUpstream + c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) + c.signalLocked() + return nil +} + +func (l *Lease) EnterFinalizing(ctx context.Context) error { + if l == nil || l.controller == nil { + return nil + } + c := l.controller + for { + c.mu.Lock() + if l.released { + c.mu.Unlock() + return ErrReleased + } + if l.phase == PhaseFinalizing { + c.mu.Unlock() + return nil + } + // Even under critical pressure, let one submitted task at a time finish + // and release its provider lease instead of deadlocking the drain path. + heavyLimit := max(c.heavyLimit, 1) + if c.preparing+c.finalizing < heavyLimit { + c.decrementPhaseLocked(l.phase) + c.finalizing++ + l.phase = PhaseFinalizing + c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) + c.mu.Unlock() + return nil + } + wake := c.wake + c.mu.Unlock() + select { + case <-ctx.Done(): + return ctx.Err() + case <-wake: + } + } +} + +func (l *Lease) Release() { + if l == nil || l.controller == nil { + return + } + c := l.controller + c.mu.Lock() + defer c.mu.Unlock() + if l.released { + return + } + c.decrementPhaseLocked(l.phase) + l.released = true + c.last = c.snapshotLocked(ResourceSample{SampledAt: c.last.SampledAt}) + c.signalLocked() +} + +func (l *Lease) Phase() Phase { + if l == nil || l.controller == nil { + return "" + } + l.controller.mu.Lock() + defer l.controller.mu.Unlock() + return l.phase +} + +func (c *Controller) adjustLocked(state PressureState, memory, cpu, database float64) { + switch state { + case PressureCritical: + c.healthyCount = 0 + c.activeLimit = max(1, c.activeLimit/2) + c.heavyLimit = 1 + case PressureBusy: + c.healthyCount = 0 + step := max(1, c.activeLimit/4) + c.activeLimit = max(1, c.activeLimit-step) + c.heavyLimit = min(c.heavyLimit, max(1, (c.activeLimit+3)/4)) + default: + if memory > .60 || cpu > .65 || database > .65 || c.activeLocked() < min(c.activeLimit, c.claimLimit) { + c.healthyCount = 0 + return + } + c.healthyCount++ + if c.healthyCount < c.healthySamples { + return + } + c.healthyCount = 0 + step := max(1, c.activeLimit/4) + c.activeLimit = min(c.hardLimit, c.activeLimit+step) + c.heavyLimit = min(c.activeLimit, max(1, (c.activeLimit+3)/4)) + } +} + +func (c *Controller) snapshotLocked(sample ResourceSample) Snapshot { + sampledAt := sample.SampledAt + if sampledAt.IsZero() { + sampledAt = c.last.SampledAt + } + safeCapacity := c.activeLimit + if c.last.PressureState == PressureCritical { + safeCapacity = 0 + } + return Snapshot{ + Mode: c.mode, + ActiveLimit: c.activeLimit, + HeavyLimit: c.heavyLimit, + ClaimLimit: c.claimLimit, + SafeCapacity: safeCapacity, + ActiveTasks: c.activeLocked(), + PreparingTasks: c.preparing, + WaitingUpstreamTasks: c.waiting, + FinalizingTasks: c.finalizing, + PressureState: c.last.PressureState, + PressureReason: c.last.PressureReason, + MemoryUtilization: c.last.MemoryUtilization, + CPUUtilization: c.last.CPUUtilization, + DBUtilization: c.last.DBUtilization, + SampledAt: sampledAt, + } +} + +func (c *Controller) activeLocked() int { return c.preparing + c.waiting + c.finalizing } + +func (c *Controller) decrementPhaseLocked(phase Phase) { + switch phase { + case PhasePreparing: + c.preparing = max(0, c.preparing-1) + case PhaseWaitingUpstream: + c.waiting = max(0, c.waiting-1) + case PhaseFinalizing: + c.finalizing = max(0, c.finalizing-1) + } +} + +func (c *Controller) signalLocked() { + select { + case c.wake <- struct{}{}: + default: + } +} + +func pressure(memory, cpu, database float64, throttled bool) (PressureState, string) { + switch { + case memory >= .90: + return PressureCritical, "memory" + case database >= .90: + return PressureCritical, "database" + case cpu >= .95 && throttled: + return PressureCritical, "cpu_throttled" + case memory >= .75: + return PressureBusy, "memory" + case database >= .80: + return PressureBusy, "database" + case cpu >= .80 || throttled: + return PressureBusy, "cpu" + default: + return PressureNormal, "" + } +} + +func utilization(current, limit int64) float64 { + if current <= 0 || limit <= 0 { + return 0 + } + return bounded(float64(current) / float64(limit)) +} + +func bounded(value float64) float64 { + if value < 0 { + return 0 + } + if value > 1 { + return 1 + } + return value +} diff --git a/apps/api/internal/workerload/controller_test.go b/apps/api/internal/workerload/controller_test.go new file mode 100644 index 0000000..17d9de7 --- /dev/null +++ b/apps/api/internal/workerload/controller_test.go @@ -0,0 +1,95 @@ +package workerload + +import ( + "context" + "testing" + "time" +) + +func TestControllerStartsConservativelyAndGrowsUnderSustainedDemand(t *testing.T) { + controller := New(Config{Mode: ModeAdaptive, HardLimit: 16, InitialActive: 4, InitialHeavy: 1, HealthySamples: 2}) + leases := make([]*Lease, 0, 4) + for range 4 { + lease, ok := controller.TryStart() + if !ok { + t.Fatal("initial task was not admitted") + } + _ = lease.EnterWaiting() + leases = append(leases, lease) + } + for range 2 { + controller.Observe(ResourceSample{MemoryCurrentBytes: 40, MemoryLimitBytes: 100, CPUUtilization: .4, DBConnections: 4, DBMaxConnections: 20}) + } + snapshot := controller.Snapshot() + if snapshot.ActiveLimit != 5 || snapshot.HeavyLimit != 2 { + t.Fatalf("grown snapshot=%+v, want active=5 heavy=2", snapshot) + } + for _, lease := range leases { + lease.Release() + } +} + +func TestControllerBusyAndCriticalPressureReduceNewClaims(t *testing.T) { + controller := New(Config{Mode: ModeAdaptive, HardLimit: 16, InitialActive: 8, InitialHeavy: 2}) + busy := controller.Observe(ResourceSample{MemoryCurrentBytes: 80, MemoryLimitBytes: 100}) + if busy.PressureState != PressureBusy || busy.SafeCapacity >= 8 { + t.Fatalf("busy snapshot=%+v", busy) + } + critical := controller.Observe(ResourceSample{MemoryCurrentBytes: 95, MemoryLimitBytes: 100}) + if critical.PressureState != PressureCritical || critical.SafeCapacity != 0 { + t.Fatalf("critical snapshot=%+v", critical) + } + controller.SetClaimLimit(0) + if _, ok := controller.TryStart(); ok { + t.Fatal("critical controller admitted a new task") + } +} + +func TestWaitingReleasesHeavyPermitAndFinalizingReacquiresIt(t *testing.T) { + controller := New(Config{Mode: ModeAdaptive, HardLimit: 4, InitialActive: 4, InitialHeavy: 1}) + first, ok := controller.TryStart() + if !ok { + t.Fatal("first lease unavailable") + } + if _, ok := controller.TryStart(); ok { + t.Fatal("second preparing task bypassed heavy limit") + } + if err := first.EnterWaiting(); err != nil { + t.Fatal(err) + } + second, ok := controller.TryStart() + if !ok { + t.Fatal("waiting task did not release heavy permit") + } + if err := second.EnterWaiting(); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := first.EnterFinalizing(ctx); err != nil { + t.Fatal(err) + } + blocked := make(chan error, 1) + go func() { blocked <- second.EnterFinalizing(ctx) }() + select { + case err := <-blocked: + t.Fatalf("second finalizer did not wait: %v", err) + case <-time.After(20 * time.Millisecond): + } + first.Release() + if err := <-blocked; err != nil { + t.Fatal(err) + } + second.Release() + if snapshot := controller.Snapshot(); snapshot.ActiveTasks != 0 { + t.Fatalf("active tasks=%d, want 0", snapshot.ActiveTasks) + } +} + +func TestLegacyModeUsesHardLimit(t *testing.T) { + controller := New(Config{Mode: ModeLegacy, HardLimit: 7}) + snapshot := controller.Observe(ResourceSample{MemoryCurrentBytes: 99, MemoryLimitBytes: 100, CPUUtilization: 1, CPUThrottled: true}) + if snapshot.ActiveLimit != 7 || snapshot.HeavyLimit != 7 || snapshot.SafeCapacity != 7 || snapshot.PressureReason != "legacy" { + t.Fatalf("legacy snapshot=%+v", snapshot) + } +} diff --git a/apps/api/internal/workerload/sampler.go b/apps/api/internal/workerload/sampler.go new file mode 100644 index 0000000..c8fa6de --- /dev/null +++ b/apps/api/internal/workerload/sampler.go @@ -0,0 +1,131 @@ +package workerload + +import ( + "errors" + "os" + "runtime" + "strconv" + "strings" + "sync" + "time" +) + +type SystemSampler struct { + mu sync.Mutex + + root string + lastCPUUsage int64 + lastCPUTime time.Time + lastThrottled int64 +} + +func NewSystemSampler() *SystemSampler { + return &SystemSampler{root: "/sys/fs/cgroup"} +} + +func NewSystemSamplerAt(root string) *SystemSampler { + return &SystemSampler{root: strings.TrimRight(root, "/")} +} + +func (s *SystemSampler) Sample(databaseConnections, databaseMax int32) ResourceSample { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + memoryCurrent, memoryLimit := s.memory() + cpuUsage, throttledCount := s.cpuStat() + cpuLimit := s.cpuLimit() + cpuUtilization := 0.0 + if s.lastCPUUsage > 0 && cpuUsage >= s.lastCPUUsage && !s.lastCPUTime.IsZero() { + elapsed := now.Sub(s.lastCPUTime).Seconds() + if elapsed > 0 { + cpuUtilization = float64(cpuUsage-s.lastCPUUsage) / 1_000_000 / elapsed / cpuLimit + } + } + throttled := s.lastThrottled > 0 && throttledCount > s.lastThrottled + s.lastCPUUsage = cpuUsage + s.lastThrottled = throttledCount + s.lastCPUTime = now + return ResourceSample{ + MemoryCurrentBytes: memoryCurrent, + MemoryLimitBytes: memoryLimit, + CPUUtilization: bounded(cpuUtilization), + CPUThrottled: throttled, + DBConnections: databaseConnections, + DBMaxConnections: databaseMax, + SampledAt: now, + } +} + +func (s *SystemSampler) memory() (int64, int64) { + current, currentErr := readIntFile(s.root + "/memory.current") + limit, limitErr := readLimitFile(s.root + "/memory.max") + if currentErr == nil && limitErr == nil { + return current, limit + } + current, _ = readIntFile(s.root + "/memory/memory.usage_in_bytes") + limit, _ = readLimitFile(s.root + "/memory/memory.limit_in_bytes") + return current, limit +} + +func (s *SystemSampler) cpuStat() (usageUsec, throttled int64) { + data, err := os.ReadFile(s.root + "/cpu.stat") + if err != nil { + return 0, 0 + } + for _, line := range strings.Split(string(data), "\n") { + fields := strings.Fields(line) + if len(fields) != 2 { + continue + } + value, parseErr := strconv.ParseInt(fields[1], 10, 64) + if parseErr != nil { + continue + } + switch fields[0] { + case "usage_usec": + usageUsec = value + case "nr_throttled": + throttled = value + } + } + return usageUsec, throttled +} + +func (s *SystemSampler) cpuLimit() float64 { + data, err := os.ReadFile(s.root + "/cpu.max") + if err == nil { + fields := strings.Fields(string(data)) + if len(fields) == 2 && fields[0] != "max" { + quota, quotaErr := strconv.ParseFloat(fields[0], 64) + period, periodErr := strconv.ParseFloat(fields[1], 64) + if quotaErr == nil && periodErr == nil && quota > 0 && period > 0 { + return max(quota/period, .001) + } + } + } + return max(float64(runtime.GOMAXPROCS(0)), 1) +} + +func readIntFile(path string) (int64, error) { + data, err := os.ReadFile(path) + if err != nil { + return 0, err + } + return strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64) +} + +func readLimitFile(path string) (int64, error) { + data, err := os.ReadFile(path) + if err != nil { + return 0, err + } + value := strings.TrimSpace(string(data)) + if value == "" || value == "max" { + return 0, errors.New("cgroup limit is unlimited") + } + limit, err := strconv.ParseInt(value, 10, 64) + if err != nil || limit <= 0 || limit > 1<<60 { + return 0, errors.New("cgroup limit is not finite") + } + return limit, nil +} diff --git a/apps/api/internal/workerload/sampler_test.go b/apps/api/internal/workerload/sampler_test.go new file mode 100644 index 0000000..5e896c0 --- /dev/null +++ b/apps/api/internal/workerload/sampler_test.go @@ -0,0 +1,34 @@ +package workerload + +import ( + "os" + "path/filepath" + "testing" + "time" +) + +func TestSystemSamplerReadsCgroupV2(t *testing.T) { + root := t.TempDir() + writeSamplerFixture(t, filepath.Join(root, "memory.current"), "50\n") + writeSamplerFixture(t, filepath.Join(root, "memory.max"), "100\n") + writeSamplerFixture(t, filepath.Join(root, "cpu.max"), "100000 100000\n") + writeSamplerFixture(t, filepath.Join(root, "cpu.stat"), "usage_usec 100000\nnr_throttled 1\n") + sampler := NewSystemSamplerAt(root) + first := sampler.Sample(2, 10) + if first.MemoryCurrentBytes != 50 || first.MemoryLimitBytes != 100 || first.DBConnections != 2 { + t.Fatalf("first sample=%+v", first) + } + time.Sleep(10 * time.Millisecond) + writeSamplerFixture(t, filepath.Join(root, "cpu.stat"), "usage_usec 105000\nnr_throttled 2\n") + second := sampler.Sample(3, 10) + if second.CPUUtilization <= 0 || !second.CPUThrottled { + t.Fatalf("second sample=%+v", second) + } +} + +func writeSamplerFixture(t *testing.T, path, value string) { + t.Helper() + if err := os.WriteFile(path, []byte(value), 0o600); err != nil { + t.Fatal(err) + } +} diff --git a/apps/api/migrations/0101_worker_adaptive_load.sql b/apps/api/migrations/0101_worker_adaptive_load.sql new file mode 100644 index 0000000..52cd884 --- /dev/null +++ b/apps/api/migrations/0101_worker_adaptive_load.sql @@ -0,0 +1,44 @@ +ALTER TABLE gateway_worker_instances + ADD COLUMN hard_capacity_limit integer NOT NULL DEFAULT 0, + ADD COLUMN safe_capacity integer NOT NULL DEFAULT 0, + ADD COLUMN heavy_capacity integer NOT NULL DEFAULT 0, + ADD COLUMN active_tasks integer NOT NULL DEFAULT 0, + ADD COLUMN preparing_tasks integer NOT NULL DEFAULT 0, + ADD COLUMN waiting_upstream_tasks integer NOT NULL DEFAULT 0, + ADD COLUMN finalizing_tasks integer NOT NULL DEFAULT 0, + ADD COLUMN pressure_state text NOT NULL DEFAULT 'unknown', + ADD COLUMN pressure_reason text NOT NULL DEFAULT '', + ADD COLUMN load_sampled_at timestamptz, + ADD CONSTRAINT gateway_worker_instances_adaptive_capacity_check + CHECK ( + hard_capacity_limit >= 0 + AND safe_capacity >= 0 + AND heavy_capacity >= 0 + AND active_tasks >= 0 + AND preparing_tasks >= 0 + AND waiting_upstream_tasks >= 0 + AND finalizing_tasks >= 0 + AND active_tasks = preparing_tasks + waiting_upstream_tasks + finalizing_tasks + ) NOT VALID, + ADD CONSTRAINT gateway_worker_instances_pressure_state_check + CHECK (pressure_state IN ('unknown', 'normal', 'busy', 'critical')) NOT VALID; + +ALTER TABLE gateway_concurrency_leases + ADD COLUMN limit_value numeric, + ADD CONSTRAINT gateway_concurrency_leases_limit_value_check + CHECK (limit_value IS NULL OR limit_value > 0) NOT VALID; + +UPDATE gateway_worker_instances +SET hard_capacity_limit = capacity_limit, + safe_capacity = capacity_limit, + heavy_capacity = capacity_limit +WHERE hard_capacity_limit = 0; + +ALTER TABLE gateway_worker_instances + VALIDATE CONSTRAINT gateway_worker_instances_adaptive_capacity_check; + +ALTER TABLE gateway_worker_instances + VALIDATE CONSTRAINT gateway_worker_instances_pressure_state_check; + +ALTER TABLE gateway_concurrency_leases + VALIDATE CONSTRAINT gateway_concurrency_leases_limit_value_check; diff --git a/apps/web/src/App.tsx b/apps/web/src/App.tsx index 1989903..6020de1 100644 --- a/apps/web/src/App.tsx +++ b/apps/web/src/App.tsx @@ -37,6 +37,7 @@ import type { UserGroupUpsertRequest, UserGroup, WalletRechargeRequest, + WorkerClusterRuntime, } from '@easyai-ai-gateway/contracts'; import { batchAccessRules, @@ -67,6 +68,7 @@ import { getRunnerPolicy, getSecurityEventConnection, getWalletSummary, + getWorkerClusterRuntime, listAccessRules, listAdminTasks, listAuditLogs, @@ -201,6 +203,7 @@ type DataKey = | 'runtimePolicySets' | 'rateLimitWindows' | 'modelRateLimits' + | 'workerClusterRuntime' | 'tenants' | 'users' | 'userGroups' @@ -255,6 +258,7 @@ export function App() { const [rateLimitWindows, setRateLimitWindows] = useState([]); const [modelRateLimits, setModelRateLimits] = useState([]); const [modelRateLimitsUpdatedAt, setModelRateLimitsUpdatedAt] = useState(null); + const [workerClusterRuntime, setWorkerClusterRuntime] = useState(null); const [tenants, setTenants] = useState([]); const [users, setUsers] = useState([]); const [userGroups, setUserGroups] = useState([]); @@ -381,18 +385,21 @@ export function App() { useEffect(() => { if (!token || activePage !== 'admin' || adminSection !== 'realtimeLoad') return undefined; const timer = window.setInterval(() => { - void Promise.all([listModelRateLimitStatuses(token), listPlatforms(token)]) - .then(([rateLimitResponse, platformResponse]) => { + void Promise.all([listModelRateLimitStatuses(token), listPlatforms(token), getWorkerClusterRuntime(token)]) + .then(([rateLimitResponse, platformResponse, workerRuntime]) => { setModelRateLimits(rateLimitResponse.items); setModelRateLimitsUpdatedAt(Date.now()); setPlatforms(platformResponse.items); + setWorkerClusterRuntime(workerRuntime); loadedDataKeysRef.current.add('modelRateLimits'); loadedDataKeysRef.current.add('platforms'); + loadedDataKeysRef.current.add('workerClusterRuntime'); }) .catch((err) => { if (handleAuthExpired(err, token)) return; loadedDataKeysRef.current.delete('modelRateLimits'); loadedDataKeysRef.current.delete('platforms'); + loadedDataKeysRef.current.delete('workerClusterRuntime'); }); }, 3000); return () => window.clearInterval(timer); @@ -446,6 +453,7 @@ export function App() { rateLimitWindows, modelRateLimits, modelRateLimitsUpdatedAt, + workerClusterRuntime, runtimePolicySets, securityEventConnection, taskResult, @@ -455,7 +463,7 @@ export function App() { users, walletAccounts, walletTransactions, - }), [accessRules, adminTasks, apiKeys, auditLogs, baseModels, clientCustomizationSettings, currentUser, currentUserGroups, fileStorageChannels, fileStorageSettings, modelCatalog, modelRateLimits, modelRateLimitsUpdatedAt, models, networkProxyConfig, platforms, pricingRuleSets, pricingRules, providers, rateLimitWindows, runnerPolicy, runtimePolicySets, securityEventConnection, taskResult, tasks, tenants, userGroups, users, walletAccounts, walletTransactions]); + }), [accessRules, adminTasks, apiKeys, auditLogs, baseModels, clientCustomizationSettings, currentUser, currentUserGroups, fileStorageChannels, fileStorageSettings, modelCatalog, modelRateLimits, modelRateLimitsUpdatedAt, models, networkProxyConfig, platforms, pricingRuleSets, pricingRules, providers, rateLimitWindows, runnerPolicy, runtimePolicySets, securityEventConnection, taskResult, tasks, tenants, userGroups, users, walletAccounts, walletTransactions, workerClusterRuntime]); async function refresh(nextToken = token) { await ensureRouteData(nextToken, true); @@ -593,6 +601,9 @@ export function App() { setModelRateLimitsUpdatedAt(Date.now()); } return; + case 'workerClusterRuntime': + setWorkerClusterRuntime(await getWorkerClusterRuntime(nextToken)); + return; case 'tenants': setTenants((await listTenants(nextToken)).items); return; @@ -1255,6 +1266,7 @@ export function App() { setAuditLogs([]); setRateLimitWindows([]); setModelRateLimits([]); + setWorkerClusterRuntime(null); setTenants([]); setUsers([]); setUserGroups([]); @@ -1744,7 +1756,7 @@ function dataKeysForRoute( case 'platforms': return ['platforms', 'models', 'providers', 'baseModels', 'pricingRuleSets', 'networkProxyConfig']; case 'realtimeLoad': - return ['platforms', 'modelRateLimits']; + return ['platforms', 'modelRateLimits', 'workerClusterRuntime']; case 'tasks': return ['adminTasks', 'tenants', 'users', 'userGroups', 'platforms', 'models']; case 'tenants': diff --git a/apps/web/src/api.ts b/apps/web/src/api.ts index 4f96374..fcd3046 100644 --- a/apps/web/src/api.ts +++ b/apps/web/src/api.ts @@ -59,6 +59,7 @@ import type { WalletBalanceAdjustmentRequest, WalletRechargeRequest, WalletSummaryResponse, + WorkerClusterRuntime, } from '@easyai-ai-gateway/contracts'; import type { AdminTaskQuery, PlatformCreateInput, PlatformModelBindingInput, WorkspaceTaskQuery } from './types'; @@ -1037,6 +1038,10 @@ export async function listModelRateLimitStatuses(token: string): Promise>('/api/admin/runtime/model-rate-limits', { token }); } +export async function getWorkerClusterRuntime(token: string): Promise { + return request('/api/admin/runtime/workers', { token }); +} + export async function restoreModelRuntimeStatus(token: string, platformModelId: string): Promise { return request(`/api/admin/runtime/model-rate-limits/${platformModelId}/restore`, { method: 'POST', diff --git a/apps/web/src/app-state.ts b/apps/web/src/app-state.ts index 77301f7..194a29e 100644 --- a/apps/web/src/app-state.ts +++ b/apps/web/src/app-state.ts @@ -26,6 +26,7 @@ import type { RuntimePolicySet, SecurityEventConnectionResponse, UserGroup, + WorkerClusterRuntime, } from '@easyai-ai-gateway/contracts'; export interface ConsoleData { @@ -50,6 +51,7 @@ export interface ConsoleData { rateLimitWindows: RateLimitWindow[]; modelRateLimits: ModelRateLimitStatus[]; modelRateLimitsUpdatedAt: number | null; + workerClusterRuntime: WorkerClusterRuntime | null; runtimePolicySets: RuntimePolicySet[]; securityEventConnection: SecurityEventConnectionResponse | null; taskResult: GatewayTask | null; diff --git a/apps/web/src/navigation.ts b/apps/web/src/navigation.ts index b4798ed..7bc4344 100644 --- a/apps/web/src/navigation.ts +++ b/apps/web/src/navigation.ts @@ -50,7 +50,7 @@ export const adminPages = [ { title: '用户组策略', path: '/admin/user-groups', description: '用户组成员、充值折扣、调用折扣、TPM/RPM/并发和队列优先级。' }, { title: '全局模型配置', path: '/admin/models/global', description: '基准模型库、能力 schema、基准定价和默认限流模板。' }, { title: '平台管理', path: '/admin/platforms', description: '平台 CRUD、凭证、默认折扣、平台模型、限流和重试策略。' }, - { title: '实时负载', path: '/admin/realtime-load', description: '按平台模型查看实时 RPM、TPM、并发、排队和冷却状态。' }, + { title: '实时负载', path: '/admin/realtime-load', description: '查看平台模型 RPM、TPM、并发，以及 Worker 自适应容量和压力状态。' }, { title: '任务记录', path: '/admin/tasks', description: '跨租户查询任务、执行链路、参数转换、计费和原始详情。' }, { title: '计费结算', path: '/admin/billing-settlements', description: '查询计费结算队列，批量处理等待重试和人工复核记录。' }, { title: '运行与队列', path: '/admin/runtime/queues', description: 'TPM/RPM 窗口、并发 lease、cooldown、任务恢复和队列积压。' }, diff --git a/apps/web/src/pages/AdminPage.tsx b/apps/web/src/pages/AdminPage.tsx index 6c52eae..9eb3a26 100644 --- a/apps/web/src/pages/AdminPage.tsx +++ b/apps/web/src/pages/AdminPage.tsx @@ -179,6 +179,7 @@ export function AdminPage(props: { modelRateLimits={props.data.modelRateLimits} modelRateLimitsUpdatedAt={props.data.modelRateLimitsUpdatedAt} platforms={props.data.platforms} + workerClusterRuntime={props.data.workerClusterRuntime} onSavePlatformDynamicPriority={props.onSavePlatformDynamicPriority} onRestoreRuntimeModel={props.onRestoreRuntimeModel} /> diff --git a/apps/web/src/pages/admin/RealtimeLoadPanel.tsx b/apps/web/src/pages/admin/RealtimeLoadPanel.tsx index 7f5fff8..6ac375d 100644 --- a/apps/web/src/pages/admin/RealtimeLoadPanel.tsx +++ b/apps/web/src/pages/admin/RealtimeLoadPanel.tsx @@ -1,13 +1,14 @@ import { useEffect, useMemo, useState, type FormEvent } from 'react'; import { Popover as AntPopover } from 'antd'; import { CheckCircle2, History, RotateCcw, Search, SlidersHorizontal } from 'lucide-react'; -import type { IntegrationPlatform, ModelRateLimitStatus, PlatformDynamicPriorityUpdateRequest, PlatformPolicyEvent, PriorityDemotionRecord } from '@easyai-ai-gateway/contracts'; +import type { IntegrationPlatform, ModelRateLimitStatus, PlatformDynamicPriorityUpdateRequest, PlatformPolicyEvent, PriorityDemotionRecord, WorkerClusterRuntime, WorkerInstanceRuntime } from '@easyai-ai-gateway/contracts'; import { Badge, Button, Card, CardContent, EmptyState, FormDialog, Input, Label, Select, Table, TableCell, TableHead, TableRow } from '../../components/ui'; export function RealtimeLoadPanel(props: { modelRateLimits: ModelRateLimitStatus[]; modelRateLimitsUpdatedAt: number | null; platforms: IntegrationPlatform[]; + workerClusterRuntime: WorkerClusterRuntime | null; onSavePlatformDynamicPriority: (platformId: string, input: PlatformDynamicPriorityUpdateRequest) => Promise; onRestoreRuntimeModel: (platformModelId: string) => Promise; }) { @@ -119,6 +120,7 @@ export function RealtimeLoadPanel(props: { return (
+
@@ -241,6 +243,95 @@ export function RealtimeLoadPanel(props: { ); } +function WorkerRuntimeTable(props: { runtime: WorkerClusterRuntime | null }) { + const workers = props.runtime?.workers ?? []; + const queue = props.runtime?.queue; + return ( + + +
+ + Worker 自适应负载 + + {queue + ? `共享队列 ${queue.queued}，运行 ${queue.running}，最老等待 ${Math.round(queue.oldestWaitSeconds)} 秒` + : '等待集群负载快照'} + + +
+ {!workers.length ? ( + + ) : ( +
+ + + Worker + 压力 + 领取 / 安全 / Hard + 重负载许可 + 准备 / 等待 / 收尾 + 任务 / 厂商租约 + 最近采样 + + {workers.map((worker) => ( + + + + {worker.podName || worker.instanceId} + {[worker.site, shortId(worker.revision)].filter(Boolean).join(' · ') || shortId(worker.instanceId)} + + + {workerPressureCell(worker)} + + {worker.allocatedCapacity} / {worker.safeCapacity} / {worker.hardCapacityLimit} + + {worker.heavyCapacity} + + + {worker.preparingTasks} / {worker.waitingUpstreamTasks} / {worker.finalizingTasks} + 上报活跃 {worker.reportedActiveTasks} + + + {worker.runningTasks} / {worker.activeLeases} + + + {formatDateTime(worker.loadSampledAt) || '-'} + 心跳 {formatDateTime(worker.heartbeatAt)} + + + + ))} +
+
+ )} + + + ); +} + +function workerPressureCell(worker: WorkerInstanceRuntime) { + const variant = worker.pressureState === 'critical' + ? 'destructive' + : worker.pressureState === 'busy' + ? 'warning' + : worker.pressureState === 'normal' + ? 'success' + : 'secondary'; + const label = worker.pressureState === 'critical' + ? '临界' + : worker.pressureState === 'busy' + ? '繁忙' + : worker.pressureState === 'normal' + ? '正常' + : '未知'; + return ( + + {label} + {worker.pressureReason || '无压力原因'} + + ); +} + type PriorityDialogState = { platform: IntegrationPlatform | undefined; status: ModelRateLimitStatus; diff --git a/deploy/kubernetes/local-acceptance/dependencies.lock b/deploy/kubernetes/local-acceptance/dependencies.lock index ab84853..7cd29ef 100644 --- a/deploy/kubernetes/local-acceptance/dependencies.lock +++ b/deploy/kubernetes/local-acceptance/dependencies.lock @@ -4,6 +4,8 @@ CRANE_VERSION=v0.21.7 K3D_DARWIN_ARM64_SHA256=fe106541d5d0a3f18debcd4d432a16f8c0ce3e6ddc06f8fbb6f696a122313e00 K3D_DARWIN_AMD64_SHA256=b4aabc37534f95b9c764e7823f2df923f50d57600837aa60a06266cce47db732 K3S_IMAGE=rancher/k3s:v1.36.2-k3s1@sha256:6a47cea22c4b834d4ba72c89d291696b79ebe406251f90b446e4dff03513dd87 +K3S_PAUSE_IMAGE=rancher/mirrored-pause:3.6@sha256:74c4244427b7312c5b901fe0f67cbc53683d06f4f24c6faee65d4182bf0fa893 +K3S_BUSYBOX_IMAGE=rancher/mirrored-library-busybox:1.37.0@sha256:101b4afd76732482eff9b95cae5f94bcf295e521fbec4e01b69c5421f3f3f3e5 CNPG_VERSION=1.29.1 CNPG_MANIFEST_URL=https://raw.githubusercontent.com/cloudnative-pg/cloudnative-pg/release-1.29/releases/cnpg-1.29.1.yaml CNPG_MANIFEST_SHA256=e0c5ff41bf5b01c0775bf225929b8423d543cd23e9742e4274fe7de8499cf93a diff --git a/deploy/kubernetes/local-acceptance/k3d.yaml b/deploy/kubernetes/local-acceptance/k3d.yaml index 4e78b58..f878e5c 100644 --- a/deploy/kubernetes/local-acceptance/k3d.yaml +++ b/deploy/kubernetes/local-acceptance/k3d.yaml @@ -57,6 +57,10 @@ options: nodeFilters: - server:0 - server:1 + - arg: --node-label=easyai.io/worker=true + nodeFilters: + - server:0 + - server:1 - arg: --node-label=easyai.io/database=true nodeFilters: - server:0 diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index c441414..a1cfa8c 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -961,6 +961,43 @@ export interface ModelRateLimitStatus { recentPriorityDemotions?: PriorityDemotionRecord[]; } +export interface WorkerInstanceRuntime { + instanceId: string; + podUid?: string; + podName?: string; + site?: string; + revision?: string; + status: string; + allocatedCapacity: number; + capacityLimit: number; + hardCapacityLimit: number; + safeCapacity: number; + heavyCapacity: number; + reportedActiveTasks: number; + preparingTasks: number; + waitingUpstreamTasks: number; + finalizingTasks: number; + pressureState: 'normal' | 'busy' | 'critical' | 'unknown' | string; + pressureReason?: string; + loadSampledAt?: string; + runningTasks: number; + activeLeases: number; + heartbeatAt: string; + drainingAt?: string; +} + +export interface WorkerQueueRuntime { + queued: number; + running: number; + oldestWaitSeconds: number; +} + +export interface WorkerClusterRuntime { + workers: WorkerInstanceRuntime[]; + queue: WorkerQueueRuntime; + capturedAt: string; +} + export interface GatewayNetworkProxyConfig { globalHttpProxy?: string; globalHttpProxySet: boolean; diff --git a/scripts/acceptance/local-cluster.sh b/scripts/acceptance/local-cluster.sh index 5463389..709a169 100755 --- a/scripts/acceptance/local-cluster.sh +++ b/scripts/acceptance/local-cluster.sh @@ -41,7 +41,8 @@ load_lock() { } # shellcheck source=/dev/null source "$lock_file" - : "${K3D_VERSION:?}" "${CRANE_VERSION:?}" "${K3S_IMAGE:?}" "${CNPG_MANIFEST_URL:?}" "${CNPG_MANIFEST_SHA256:?}" + : "${K3D_VERSION:?}" "${CRANE_VERSION:?}" "${K3S_IMAGE:?}" "${K3S_PAUSE_IMAGE:?}" "${K3S_BUSYBOX_IMAGE:?}" + : "${CNPG_MANIFEST_URL:?}" "${CNPG_MANIFEST_SHA256:?}" : "${CNPG_CONTROLLER_IMAGE:?}" "${CNPG_POSTGRES_IMAGE:?}" "${K3S_COREDNS_IMAGE:?}" : "${K3S_METRICS_SERVER_IMAGE:?}" "${K3S_LOCAL_PATH_IMAGE:?}" } @@ -123,6 +124,9 @@ host_available_disk_bytes() { } preflight() { + local required_disk_gib=${1:-30} required_disk_bytes + [[ $required_disk_gib =~ ^[0-9]+$ && $required_disk_gib -ge 10 ]] + required_disk_bytes=$((required_disk_gib * 1024 * 1024 * 1024)) load_lock require_command docker require_command kubectl @@ -165,12 +169,12 @@ preflight() { exit 1 } available_disk_bytes=$(host_available_disk_bytes) - [[ $available_disk_bytes =~ ^[0-9]+$ && $available_disk_bytes -ge 32212254720 ]] || { + [[ $available_disk_bytes =~ ^[0-9]+$ && $available_disk_bytes -ge $required_disk_bytes ]] || { available_disk_gib=$(awk -v bytes="${available_disk_bytes:-0}" 'BEGIN {printf "%.1f", bytes/1024/1024/1024}') - echo "host free disk is ${available_disk_gib} GiB; local acceptance requires at least 30 GiB" >&2 + echo "host free disk is ${available_disk_gib} GiB; local acceptance requires at least ${required_disk_gib} GiB" >&2 exit 1 } - echo "local_acceptance_preflight=PASS docker_memory_bytes=$memory_bytes docker_cpus=$cpu_count host_available_disk_bytes=$available_disk_bytes architecture=$architecture k3d=$K3D_VERSION" + echo "local_acceptance_preflight=PASS docker_memory_bytes=$memory_bytes docker_cpus=$cpu_count host_available_disk_bytes=$available_disk_bytes required_disk_gib=$required_disk_gib architecture=$architecture k3d=$K3D_VERSION" } ensure_private_material() { @@ -224,8 +228,8 @@ render_k3d_config() { host_memory_mib=$((host_memory_bytes / 1024 / 1024)) site_system_cpu="$(((host_cpu - 4) * 1000 + 750))m" witness_system_cpu="$(((host_cpu - 2) * 1000 + 750))m" - site_system_memory="$((host_memory_mib - 8192 + 1024))Mi" - witness_system_memory="$((host_memory_mib - 4096 + 1024))Mi" + site_system_memory="$((host_memory_mib - 8192 + 512))Mi" + witness_system_memory="$((host_memory_mib - 4096 + 512))Mi" sed \ -e "s|EASYAI_ACCEPTANCE_MEDIA_DIR|$media_root|g" \ -e "s|EASYAI_ACCEPTANCE_K3S_IMAGE|${K3S_IMAGE%%@*}|g" \ @@ -294,7 +298,8 @@ remove_cluster_network() { pull_dependency_images() { local image - for image in "$K3S_IMAGE" "$CNPG_CONTROLLER_IMAGE" "$CNPG_POSTGRES_IMAGE" \ + for image in "$K3S_IMAGE" "$K3S_PAUSE_IMAGE" "$K3S_BUSYBOX_IMAGE" \ + "$CNPG_CONTROLLER_IMAGE" "$CNPG_POSTGRES_IMAGE" \ "$K3S_COREDNS_IMAGE" "$K3S_METRICS_SERVER_IMAGE" "$K3S_LOCAL_PATH_IMAGE"; do pull_dependency_image "$image" done @@ -596,6 +601,10 @@ render_and_apply_application() { -f "$repository_root/deploy/kubernetes/production/service-account-rbac.yaml" >/dev/null kubectl --context "$context" -n "$namespace" apply -f "$manifest_root/local-config.yaml" >/dev/null kubectl --context "$context" -n "$namespace" apply -f "$rendered" >/dev/null + for workload in easyai-worker-ningbo easyai-worker-hongkong; do + kubectl --context "$context" -n "$namespace" patch deployment "$workload" \ + --type=merge -p='{"spec":{"template":{"spec":{"affinity":null}}}}' >/dev/null + done for workload in easyai-api-ningbo easyai-worker-ningbo; do kubectl --context "$context" -n "$namespace" set env deployment/"$workload" \ @@ -791,10 +800,14 @@ up_cluster() { echo 'local K3s node allocatable capacity does not match 4/8, 4/8, 2/4 resource envelope' >&2 exit 1 } - "$k3d" image import -c "$cluster_name" \ + local cluster_image + for cluster_image in \ "$api_image" "$web_image" "$netem_image" \ + "${K3S_PAUSE_IMAGE%%@*}" "${K3S_BUSYBOX_IMAGE%%@*}" \ "${CNPG_CONTROLLER_IMAGE%%@*}" "${CNPG_POSTGRES_IMAGE%%@*}" \ - "${K3S_COREDNS_IMAGE%%@*}" "${K3S_METRICS_SERVER_IMAGE%%@*}" "${K3S_LOCAL_PATH_IMAGE%%@*}" + "${K3S_COREDNS_IMAGE%%@*}" "${K3S_METRICS_SERVER_IMAGE%%@*}" "${K3S_LOCAL_PATH_IMAGE%%@*}"; do + "$k3d" image import -c "$cluster_name" "$cluster_image" + done local cnpg_manifest="$state_root/cnpg.yaml" curl -fsSL "$CNPG_MANIFEST_URL" -o "$cnpg_manifest" @@ -844,7 +857,7 @@ up_cluster() { new_run() { load_lock - preflight + preflight 10 [[ -z $(git -C "$repository_root" status --short) ]] || { echo 'local acceptance requires a clean source tree for a fresh Run ID' >&2 exit 1 diff --git a/scripts/acceptance/local-control-plane.sh b/scripts/acceptance/local-control-plane.sh index 7bb8ad8..4413b4d 100644 --- a/scripts/acceptance/local-control-plane.sh +++ b/scripts/acceptance/local-control-plane.sh @@ -61,11 +61,12 @@ verify_local_node_capacity() { elif endswith("Gi") then rtrimstr("Gi") | tonumber * 1024 else 0 end; ([.items[] | select(.metadata.labels["easyai.io/workload"] == "true") | + select(.metadata.labels["easyai.io/worker"] == "true") | select((.status.allocatable.cpu | cpu_m) == 3000 and - (.status.allocatable.memory | memory_mi) == 6656)] | length) == 2 and + (.status.allocatable.memory | memory_mi | floor) == 6656)] | length) == 2 and ([.items[] | select(.metadata.labels["easyai.io/site"] == "los-angeles") | select((.status.allocatable.cpu | cpu_m) == 1000 and - (.status.allocatable.memory | memory_mi) == 2560)] | length) == 1 and + (.status.allocatable.memory | memory_mi | floor) == 2560)] | length) == 1 and ([.items[] | select(any(.status.conditions[]; .type == "Ready" and .status == "True"))] | length) == 3 and ([.items[] | select(any(.status.conditions[]; .type == "MemoryPressure" and .status == "True"))] | length) == 0 ' >/dev/null diff --git a/scripts/acceptance/provider-burst.sh b/scripts/acceptance/provider-burst.sh new file mode 100644 index 0000000..777b670 --- /dev/null +++ b/scripts/acceptance/provider-burst.sh @@ -0,0 +1,515 @@ +#!/usr/bin/env bash + +# This library is sourced by run-local-acceptance.sh and intentionally uses its +# validated runtime globals and lifecycle state. +# shellcheck disable=SC2034,SC2154 + +provider_burst_platform_ids=( + f1000000-0000-4000-8000-000000000002 + f1000000-0000-4000-8000-000000000004 + f1000000-0000-4000-8000-000000000006 +) + +provider_burst_model_ids=( + f2000000-0000-4000-8000-000000000002 + f2000000-0000-4000-8000-000000000004 + f2000000-0000-4000-8000-000000000006 +) + +restore_provider_burst_platforms() { + database_query " +UPDATE integration_platforms +SET status = config->>'acceptanceBurstPreviousStatus', + config = config - 'acceptanceBurstPreviousStatus', + updated_at = statement_timestamp() +WHERE config ? 'acceptanceBurstPreviousStatus'; + +UPDATE integration_platforms +SET status = 'disabled', + updated_at = statement_timestamp() +WHERE COALESCE((config->>'acceptanceProviderBurst')::boolean, false);" >/dev/null +} + +setup_provider_burst_platforms() { + [[ $run_id =~ ^[0-9a-f-]{36}$ ]] + [[ $acceptance_video_model =~ ^[A-Za-z0-9._:-]+$ ]] + restore_provider_burst_platforms + database_query " +WITH source AS ( + SELECT platform.* + FROM integration_platforms platform + JOIN platform_models model ON model.platform_id = platform.id + JOIN base_model_catalog base_model ON base_model.id = model.base_model_id + WHERE base_model.invocation_name = '$acceptance_video_model' + AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) + ORDER BY platform.priority, platform.created_at + LIMIT 1 +), variants(platform_id, platform_key, platform_name, concurrency_limit) AS ( + VALUES + ('${provider_burst_platform_ids[0]}'::uuid, 'acceptance-provider-burst-2', 'Acceptance Provider Burst C2', 2), + ('${provider_burst_platform_ids[1]}'::uuid, 'acceptance-provider-burst-4', 'Acceptance Provider Burst C4', 4), + ('${provider_burst_platform_ids[2]}'::uuid, 'acceptance-provider-burst-6', 'Acceptance Provider Burst C6', 6) +) +INSERT INTO integration_platforms ( + id, provider, platform_key, name, base_url, auth_type, credentials, config, + visibility_scope, tenant_id, tenant_key, default_pricing_mode, + default_discount_factor, pricing_rule_set_id, retry_policy, rate_limit_policy, + priority, dynamic_priority, status, disabled_reason, cooldown_until, internal_name +) +SELECT variants.platform_id, source.provider, variants.platform_key, variants.platform_name, + source.base_url, source.auth_type, source.credentials, + source.config || jsonb_build_object( + 'acceptanceProviderBurst', true, + 'acceptanceBurstRunId', '$run_id', + 'acceptanceConcurrencyLimit', variants.concurrency_limit + ), + source.visibility_scope, source.tenant_id, source.tenant_key, + source.default_pricing_mode, source.default_discount_factor, + source.pricing_rule_set_id, source.retry_policy, '{\"rules\":[]}'::jsonb, + 100, NULL, 'enabled', NULL, NULL, variants.platform_name +FROM source CROSS JOIN variants +ON CONFLICT (id) DO UPDATE SET + provider = EXCLUDED.provider, + platform_key = EXCLUDED.platform_key, + name = EXCLUDED.name, + base_url = EXCLUDED.base_url, + auth_type = EXCLUDED.auth_type, + credentials = EXCLUDED.credentials, + config = EXCLUDED.config, + visibility_scope = EXCLUDED.visibility_scope, + tenant_id = EXCLUDED.tenant_id, + tenant_key = EXCLUDED.tenant_key, + default_pricing_mode = EXCLUDED.default_pricing_mode, + default_discount_factor = EXCLUDED.default_discount_factor, + pricing_rule_set_id = EXCLUDED.pricing_rule_set_id, + retry_policy = EXCLUDED.retry_policy, + rate_limit_policy = EXCLUDED.rate_limit_policy, + priority = EXCLUDED.priority, + dynamic_priority = EXCLUDED.dynamic_priority, + status = 'enabled', + disabled_reason = NULL, + cooldown_until = NULL, + internal_name = EXCLUDED.internal_name, + deleted_at = NULL, + updated_at = statement_timestamp(); + +WITH source AS ( + SELECT model.* + FROM platform_models model + JOIN integration_platforms platform ON platform.id = model.platform_id + JOIN base_model_catalog base_model ON base_model.id = model.base_model_id + WHERE base_model.invocation_name = '$acceptance_video_model' + AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) + ORDER BY platform.priority, platform.created_at + LIMIT 1 +), variants(model_id, platform_id, concurrency_limit) AS ( + VALUES + ('${provider_burst_model_ids[0]}'::uuid, '${provider_burst_platform_ids[0]}'::uuid, 2), + ('${provider_burst_model_ids[1]}'::uuid, '${provider_burst_platform_ids[1]}'::uuid, 4), + ('${provider_burst_model_ids[2]}'::uuid, '${provider_burst_platform_ids[2]}'::uuid, 6) +) +INSERT INTO platform_models ( + id, platform_id, base_model_id, model_name, model_alias, model_type, + display_name, capability_override, capabilities, pricing_mode, + discount_factor, pricing_rule_set_id, billing_config_override, billing_config, + permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, + runtime_policy_set_id, runtime_policy_override, provider_model_name, + cooldown_until, enabled +) +SELECT variants.model_id, variants.platform_id, source.base_model_id, + source.model_name, source.model_alias, source.model_type, source.display_name, + source.capability_override, source.capabilities, source.pricing_mode, + source.discount_factor, source.pricing_rule_set_id, + source.billing_config_override, source.billing_config, + source.permission_config, source.retry_policy, + jsonb_build_object('rules', jsonb_build_array( + jsonb_build_object( + 'metric', 'concurrent', 'limit', variants.concurrency_limit, + 'leaseTtlSeconds', 120 + ), + jsonb_build_object('metric', 'rpm', 'limit', 600, 'windowSeconds', 60) + )), + 'override', NULL, '{}'::jsonb, source.provider_model_name, NULL, true +FROM source CROSS JOIN variants +ON CONFLICT (id) DO UPDATE SET + platform_id = EXCLUDED.platform_id, + base_model_id = EXCLUDED.base_model_id, + model_name = EXCLUDED.model_name, + model_alias = EXCLUDED.model_alias, + model_type = EXCLUDED.model_type, + display_name = EXCLUDED.display_name, + capability_override = EXCLUDED.capability_override, + capabilities = EXCLUDED.capabilities, + pricing_mode = EXCLUDED.pricing_mode, + discount_factor = EXCLUDED.discount_factor, + pricing_rule_set_id = EXCLUDED.pricing_rule_set_id, + billing_config_override = EXCLUDED.billing_config_override, + billing_config = EXCLUDED.billing_config, + permission_config = EXCLUDED.permission_config, + retry_policy = EXCLUDED.retry_policy, + rate_limit_policy = EXCLUDED.rate_limit_policy, + rate_limit_policy_mode = 'override', + runtime_policy_set_id = NULL, + runtime_policy_override = '{}'::jsonb, + provider_model_name = EXCLUDED.provider_model_name, + cooldown_until = NULL, + enabled = true, + updated_at = statement_timestamp(); + +WITH selected_platforms AS ( + SELECT platform.id + FROM integration_platforms platform + JOIN platform_models model ON model.platform_id = platform.id + JOIN base_model_catalog base_model ON base_model.id = model.base_model_id + WHERE base_model.invocation_name = '$acceptance_video_model' + AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false) + AND NOT COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) +) +UPDATE integration_platforms platform +SET config = jsonb_set( + platform.config, + '{acceptanceBurstPreviousStatus}', + to_jsonb(platform.status), + true + ), + status = 'disabled', + updated_at = statement_timestamp() +FROM selected_platforms +WHERE platform.id = selected_platforms.id; + +WITH acceptance_group AS ( + SELECT api_key.user_group_id AS id + FROM gateway_acceptance_runs run + JOIN gateway_api_keys api_key ON api_key.id::text = run.api_key_id + WHERE run.id = '$run_id'::uuid +), resources(resource_type, resource_id) AS ( + SELECT 'platform'::text, platform.id + FROM integration_platforms platform + WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId' = '$run_id' + UNION ALL + SELECT 'platform_model', model.id + FROM platform_models model + JOIN integration_platforms platform ON platform.id = model.platform_id + WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId' = '$run_id' + UNION + SELECT 'base_model', model.base_model_id + FROM platform_models model + JOIN integration_platforms platform ON platform.id = model.platform_id + WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId' = '$run_id' +) +INSERT INTO gateway_access_rules ( + subject_type, subject_id, resource_type, resource_id, effect, + priority, min_permission_level, conditions, metadata, status +) +SELECT 'user_group', acceptance_group.id, resources.resource_type, + resources.resource_id, 'allow', 1, 0, '{}'::jsonb, + jsonb_build_object('purpose', 'provider_burst_acceptance', 'runId', '$run_id'), + 'active' +FROM acceptance_group CROSS JOIN resources +WHERE acceptance_group.id IS NOT NULL +ON CONFLICT (subject_type, subject_id, resource_type, resource_id, effect) +DO UPDATE SET status = 'active', metadata = EXCLUDED.metadata, + updated_at = statement_timestamp();" >/dev/null + + local configured + configured=$(database_query " +SELECT jsonb_build_object( + 'platforms', count(*), + 'limits', jsonb_agg((platform.config->>'acceptanceConcurrencyLimit')::int ORDER BY (platform.config->>'acceptanceConcurrencyLimit')::int), + 'enabled', count(*) FILTER (WHERE platform.status='enabled' AND model.enabled) +) +FROM integration_platforms platform +JOIN platform_models model ON model.platform_id = platform.id +WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId' = '$run_id';") + jq -e '.platforms == 3 and .enabled == 3 and .limits == [2,4,6]' <<<"$configured" >/dev/null +} + +sample_provider_burst() { + local output=$1 stop_file=$2 + : >"$output" + while [[ ! -e $stop_file ]]; do + database_query " +WITH configured AS ( + SELECT platform.id platform_id, platform.platform_key, model.id platform_model_id, + (platform.config->>'acceptanceConcurrencyLimit')::int concurrency_limit + FROM integration_platforms platform + JOIN platform_models model ON model.platform_id = platform.id + WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId' = '$run_id' +), platform_state AS ( + SELECT configured.*, + (SELECT count(*) FROM gateway_concurrency_leases lease + WHERE lease.scope_type='platform_model' + AND lease.scope_key='acceptance:$run_id:'||configured.platform_model_id::text + AND lease.released_at IS NULL + AND lease.expires_at > statement_timestamp()) active_leases, + (SELECT count(*) FROM gateway_task_attempts attempt + WHERE attempt.platform_model_id=configured.platform_model_id + AND attempt.task_id IN (SELECT id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid)) attempts + FROM configured +) +SELECT jsonb_build_object( + 'sampledAt', to_char(clock_timestamp() AT TIME ZONE 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"'), + 'queuedTasks', (SELECT count(*) FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND status='queued'), + 'runningTasks', (SELECT count(*) FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND status='running'), + 'succeededTasks', (SELECT count(*) FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND status='succeeded'), + 'waitingAdmissions', (SELECT count(*) FROM gateway_task_admissions admission JOIN gateway_tasks task ON task.id=admission.task_id WHERE task.acceptance_run_id='$run_id'::uuid AND admission.status='waiting'), + 'platforms', COALESCE((SELECT jsonb_agg(jsonb_build_object( + 'platformId', platform_id, + 'platformModelId', platform_model_id, + 'platformKey', platform_key, + 'limit', concurrency_limit, + 'activeLeases', active_leases, + 'attempts', attempts + ) ORDER BY concurrency_limit) FROM platform_state), '[]'::jsonb) +);" >>"$output" || true + sleep 1 + done +} + +provider_burst_resource_summary() { + local resources=$1 + jq -Rn ' + def number_without($suffix): rtrimstr($suffix) | tonumber; + def memory_mib: + if endswith("Gi") then number_without("Gi") * 1024 + elif endswith("Mi") then number_without("Mi") + elif endswith("Ki") then number_without("Ki") / 1024 + else tonumber end; + [inputs | split(",") | select(.[0] != "timestamp") | { + scope:.[1],name:.[2],cpu:.[3],memory:.[4] + }] + | group_by([.scope,.name]) + | map(if .[0].scope == "node" then { + scope:"node",name:.[0].name, + maxCpuPercent:(map(.cpu | number_without("%")) | max), + maxMemoryPercent:(map(.memory | number_without("%")) | max) + } else { + scope:"pod",name:.[0].name, + maxCpuMillicores:(map(.cpu | if endswith("m") then number_without("m") else tonumber * 1000 end) | max), + maxMemoryMiB:(map(.memory | memory_mib) | max) + } end)' <"$resources" +} + +build_provider_burst_report() { + local output=$1 load_report=$2 samples=$3 worker_samples=$4 resources=$5 requests=$6 + local platforms queue workers distribution worker_peaks limits concurrency_peaks + local leaks duplicates callbacks resource_summary passed=true + platforms=$(database_query " +WITH configured AS ( + SELECT platform.id platform_id, platform.platform_key, model.id platform_model_id, + (platform.config->>'acceptanceConcurrencyLimit')::int concurrency_limit + FROM integration_platforms platform + JOIN platform_models model ON model.platform_id=platform.id + WHERE COALESCE((platform.config->>'acceptanceProviderBurst')::boolean, false) + AND platform.config->>'acceptanceBurstRunId'='$run_id' +), events AS ( + SELECT configured.platform_model_id, lease.acquired_at event_at, lease.lease_value delta + FROM configured + JOIN gateway_concurrency_leases lease + ON lease.scope_type='platform_model' + AND lease.scope_key='acceptance:$run_id:'||configured.platform_model_id::text + UNION ALL + SELECT configured.platform_model_id, + LEAST(COALESCE(lease.released_at, lease.expires_at), lease.expires_at), + -lease.lease_value + FROM configured + JOIN gateway_concurrency_leases lease + ON lease.scope_type='platform_model' + AND lease.scope_key='acceptance:$run_id:'||configured.platform_model_id::text +), points AS ( + SELECT platform_model_id, + SUM(delta) OVER (PARTITION BY platform_model_id ORDER BY event_at,delta ROWS UNBOUNDED PRECEDING) active + FROM events +), peaks AS ( + SELECT platform_model_id, COALESCE(max(active),0) peak FROM points GROUP BY platform_model_id +), attempts AS ( + SELECT attempt.platform_model_id, count(*) attempts + FROM gateway_task_attempts attempt + WHERE attempt.task_id IN (SELECT id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) + GROUP BY attempt.platform_model_id +) +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'platformId', configured.platform_id, + 'platformModelId', configured.platform_model_id, + 'platformKey', configured.platform_key, + 'limit', configured.concurrency_limit, + 'peak', COALESCE(peaks.peak,0), + 'attempts', COALESCE(attempts.attempts,0) +) ORDER BY configured.concurrency_limit), '[]'::jsonb) +FROM configured +LEFT JOIN peaks ON peaks.platform_model_id=configured.platform_model_id +LEFT JOIN attempts ON attempts.platform_model_id=configured.platform_model_id;") + queue=$(jq -s '{ + samples:length, + maxQueuedTasks:(map(.queuedTasks)|max), + maxRunningTasks:(map(.runningTasks)|max), + maxWaitingAdmissions:(map(.waitingAdmissions)|max), + queueObserved:any(.queuedTasks > 0), + admissionWaitObserved:any(.waitingAdmissions > 0), + drainedAtEnd:(.[-1].queuedTasks == 0 and .[-1].runningTasks == 0 and .[-1].waitingAdmissions == 0) + }' "$samples") + workers=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'instanceId',instance_id,'podName',pod_name,'site',site, + 'allocatedCapacity',allocated_capacity,'safeCapacity',safe_capacity, + 'heavyCapacity',heavy_capacity,'pressureState',pressure_state +) ORDER BY site,instance_id),'[]'::jsonb) +FROM gateway_worker_instances +WHERE status='active' AND heartbeat_at > now() - interval '30 seconds';") + distribution=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object('workerInstanceId',worker_id,'tasks',tasks) ORDER BY worker_id),'[]'::jsonb) +FROM ( + SELECT COALESCE( + regexp_replace(job.attempted_by[cardinality(job.attempted_by)], '-exec-[^-]+-[^-]+$', ''), + 'unclaimed' + ) worker_id, + count(*) tasks + FROM gateway_tasks task + LEFT JOIN river_job job ON job.id=task.river_job_id + WHERE task.acceptance_run_id='$run_id'::uuid + GROUP BY worker_id +) grouped;") + worker_peaks=$(jq -s '[.[].workers[]] | group_by(.instanceId) | map({ + instanceId:.[0].instanceId, + peakActiveTasks:(map(.activeTasks)|max), + peakPreparingTasks:(map(.preparingTasks)|max), + peakWaitingUpstreamTasks:(map(.waitingUpstreamTasks)|max), + peakFinalizingTasks:(map(.finalizingTasks)|max), + minSafeCapacity:(map(.safeCapacity)|min), + maxSafeCapacity:(map(.safeCapacity)|max), + pressureStates:(map(.pressureState)|unique) + })' "$worker_samples") + limits=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'scopeType',scope_type,'scopeKey',scope_key,'metric',metric, + 'limit',limit_value,'used',used_value,'reserved',reserved_value, + 'windowStart',window_start,'resetAt',reset_at +) ORDER BY scope_type,scope_key,metric,window_start),'[]'::jsonb) +FROM gateway_rate_limit_counters +WHERE scope_key LIKE 'acceptance:$run_id:%';") + concurrency_peaks=$(database_query " +WITH events AS ( + SELECT scope_type,scope_key,limit_value,acquired_at event_at,lease_value delta + FROM gateway_concurrency_leases WHERE scope_key LIKE 'acceptance:$run_id:%' + UNION ALL + SELECT scope_type,scope_key,limit_value, + LEAST(COALESCE(released_at,expires_at),expires_at),-lease_value + FROM gateway_concurrency_leases WHERE scope_key LIKE 'acceptance:$run_id:%' +), points AS ( + SELECT scope_type,scope_key,limit_value, + SUM(delta) OVER (PARTITION BY scope_type,scope_key ORDER BY event_at,delta ROWS UNBOUNDED PRECEDING) active + FROM events +), peaks AS ( + SELECT scope_type,scope_key,max(limit_value) limit_value,COALESCE(max(active),0) peak + FROM points GROUP BY scope_type,scope_key +) +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'scopeType',scope_type,'scopeKey',scope_key,'limit',limit_value,'peak',peak +) ORDER BY scope_type,scope_key),'[]'::jsonb) FROM peaks;") + leaks=$(database_query " +SELECT jsonb_build_object( + 'activeConcurrencyLeases',(SELECT count(*) FROM gateway_concurrency_leases WHERE scope_key LIKE 'acceptance:$run_id:%' AND released_at IS NULL AND expires_at>statement_timestamp()), + 'activeRateReservations',(SELECT count(*) FROM gateway_rate_limit_reservations WHERE scope_key LIKE 'acceptance:$run_id:%' AND status='reserved'), + 'activeRiverJobs',(SELECT count(*) FROM river_job job JOIN gateway_tasks task ON task.river_job_id=job.id WHERE task.acceptance_run_id='$run_id'::uuid AND job.state IN ('available','pending','retryable','running','scheduled')) +);") + duplicates=$(database_query " +SELECT jsonb_build_object( + 'remoteTaskIds',(SELECT count(*) FROM (SELECT remote_task_id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND remote_task_id IS NOT NULL GROUP BY remote_task_id HAVING count(*)>1) duplicate), + 'billingTransactions',(SELECT count(*) FROM (SELECT reference_id,transaction_type FROM gateway_wallet_transactions WHERE reference_type='gateway_task' AND reference_id IN (SELECT id::text FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) GROUP BY reference_id,transaction_type HAVING count(*)>1) duplicate) +);") + callbacks=$(database_query " +SELECT jsonb_build_object( + 'deliveries',count(*) FILTER (WHERE callback.status='delivered'), + 'pending',count(*) FILTER (WHERE callback.status<>'delivered'), + 'duplicates',(SELECT count(*) FROM ( + SELECT callback_inner.task_id,callback_inner.seq,callback_inner.callback_url + FROM gateway_task_callback_outbox callback_inner + WHERE callback_inner.task_id IN (SELECT id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) + GROUP BY callback_inner.task_id,callback_inner.seq,callback_inner.callback_url + HAVING count(*)>1 + ) duplicate) +) +FROM gateway_task_callback_outbox callback +WHERE callback.task_id IN (SELECT id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid);") + resource_summary=$(provider_burst_resource_summary "$resources") + if ! jq -e 'length==3 and all(.[]; .peak == .limit and .attempts > 0)' <<<"$platforms" >/dev/null; then passed=false; fi + if ! jq -e '.queueObserved and .admissionWaitObserved and .drainedAtEnd and .maxQueuedTasks > 0 and .maxWaitingAdmissions > 0' <<<"$queue" >/dev/null; then passed=false; fi + if ! jq -e --argjson requests "$requests" 'length==3 and all(.[]; .tasks>0 and .workerInstanceId!="unclaimed") and ([.[].tasks]|add)==$requests' <<<"$distribution" >/dev/null; then passed=false; fi + if ! jq -e 'length==3 and all(.[]; .peakActiveTasks>0 and (.pressureStates|index("critical")|not))' <<<"$worker_peaks" >/dev/null; then passed=false; fi + if ! jq -e 'length>0 and all(.[]; (.used+.reserved)<=.limit)' <<<"$limits" >/dev/null; then passed=false; fi + if ! jq -e 'all(.[]; .peak<=.limit)' <<<"$concurrency_peaks" >/dev/null; then passed=false; fi + if ! jq -e '.activeConcurrencyLeases==0 and .activeRateReservations==0 and .activeRiverJobs==0' <<<"$leaks" >/dev/null; then passed=false; fi + if ! jq -e '.remoteTaskIds==0 and .billingTransactions==0' <<<"$duplicates" >/dev/null; then passed=false; fi + if ! jq -e '.deliveries>0 and .pending==0 and .duplicates==0' <<<"$callbacks" >/dev/null; then passed=false; fi + if ! jq -e --argjson requests "$requests" '.passed and ([.phases[].completed]|add)==$requests and ([.phases[].failed]|add)==0' "$load_report" >/dev/null; then passed=false; fi + jq -n \ + --arg runId "$run_id" --arg model "$acceptance_video_model" \ + --argjson passed "$passed" --argjson requests "$requests" \ + --argjson phases "$(jq '.phases' "$load_report")" \ + --argjson platforms "$platforms" --argjson queue "$queue" \ + --argjson workers "$workers" --argjson distribution "$distribution" \ + --argjson workerPeaks "$worker_peaks" --argjson rateLimits "$limits" \ + --argjson concurrencyPeaks "$concurrency_peaks" \ + --argjson resources "$resource_summary" --argjson leaks "$leaks" \ + --argjson duplicates "$duplicates" --argjson callbacks "$callbacks" \ + '{ + schemaVersion:"acceptance-provider-burst-report/v1", + runId:$runId,profile:"provider-burst",model:$model, + startedAt:(now|todateiso8601),finishedAt:(now|todateiso8601), + passed:$passed,secretSafe:true,requests:$requests,phases:$phases, + providerRouting:{platforms:$platforms,totalConfiguredConcurrency:([$platforms[].limit]|add)}, + queue:$queue, + cluster:{workers:$workers,workerPeaks:$workerPeaks,taskDistribution:$distribution, + rateLimits:$rateLimits,concurrencyPeaks:$concurrencyPeaks, + resourcePeaks:$resources,leaks:$leaks,duplicates:$duplicates,callbacks:$callbacks} + }' >"$output" + chmod 0600 "$output" + jq -e '.passed==true and .secretSafe==true' "$output" >/dev/null +} + +provider_burst() { + local requests=${AI_GATEWAY_LOCAL_PROVIDER_BURST_REQUESTS:-48} + local load_report="$report_root/provider-burst-load.json" + local samples="$report_root/provider-burst-samples.ndjson" + local worker_samples="$report_root/provider-burst-worker-samples.ndjson" + local resources="$report_root/provider-burst-resources.csv" + local report="$report_root/provider-burst.json" + worker_sample_stop="$report_root/provider-burst-samples.stop" + [[ $requests =~ ^[0-9]+$ ]] && ((requests >= 24 && requests <= 256)) + [[ ! -e $load_report && ! -e $samples && ! -e $worker_samples && ! -e $resources && ! -e $report && ! -e $worker_sample_stop ]] || + fail_gate acceptance_report_exists "provider burst evidence already exists" + current_phase=provider_burst_platforms + setup_provider_burst_platforms + provider_burst_configured=true + current_phase=provider_burst_workers + configure_adaptive_workers + sample_provider_burst "$samples" "$worker_sample_stop" & + active_provider_sampler_pid=$! + sample_worker_runtime "$worker_samples" "$worker_sample_stop" & + active_worker_sampler_pid=$! + sample_resources "$resources" "$worker_sample_stop" & + active_resource_sampler_pid=$! + current_phase=provider_burst_load + run_load video-throughput "$load_report" -requests "$requests" + sleep 2 + touch "$worker_sample_stop" + wait "$active_provider_sampler_pid" + wait "$active_worker_sampler_pid" + wait "$active_resource_sampler_pid" + active_provider_sampler_pid= + active_worker_sampler_pid= + active_resource_sampler_pid= + worker_sample_stop= + current_phase=provider_burst_hard_gates + verify_hard_gates + current_phase=provider_burst_report + build_provider_burst_report "$report" "$load_report" "$samples" "$worker_samples" "$resources" "$requests" + restore_provider_burst_platforms + provider_burst_configured=false + echo "local_acceptance_provider_burst=PASS run_id=$run_id requests=$requests report=$report" +} diff --git a/scripts/acceptance/report.mjs b/scripts/acceptance/report.mjs index 84b032b..561b797 100755 --- a/scripts/acceptance/report.mjs +++ b/scripts/acceptance/report.mjs @@ -154,7 +154,7 @@ async function loadReports(directory) { const reports = []; for (const name of entries) { const report = await regularJSON(resolve(directory, name)); - if (report.schemaVersion === 'acceptance-load-report/v1') { + if (report.schemaVersion === 'acceptance-load-report/v1' || report.schemaVersion === 'acceptance-load-report/v2') { if (report.secretSafe !== true) throw new Error(`load report ${name} is not secret-safe`); reports.push({ file: name, @@ -163,6 +163,8 @@ async function loadReports(directory) { phases: report.phases, failure: report.failure, failureOperation: report.failureOperation, + capacity: report.capacity, + cluster: report.cluster, startedAt: report.startedAt, finishedAt: report.finishedAt }); diff --git a/scripts/acceptance/run-local-acceptance.sh b/scripts/acceptance/run-local-acceptance.sh index 26a9961..0f54bc8 100755 --- a/scripts/acceptance/run-local-acceptance.sh +++ b/scripts/acceptance/run-local-acceptance.sh @@ -13,20 +13,32 @@ snapshot="$state_root/snapshot.json" load_binary="$state_root/easyai-ai-gateway-acceptance-load" identity_file="$state_root/control-plane-identity.json" active_load_pid= +active_worker_sampler_pid= +active_resource_sampler_pid= +active_provider_sampler_pid= +worker_sample_stop= netem_active=false +provider_burst_configured=false current_phase=initialization stable_profile= failure_gate_id=local_execution_incomplete # shellcheck source=scripts/acceptance/local-control-plane.sh source "$script_dir/local-control-plane.sh" +# shellcheck source=scripts/acceptance/provider-burst.sh +source "$script_dir/provider-burst.sh" usage() { cat <<'EOF' Usage: scripts/acceptance/run-local-acceptance.sh quick + scripts/acceptance/run-local-acceptance.sh adaptive + scripts/acceptance/run-local-acceptance.sh provider-burst scripts/acceptance/run-local-acceptance.sh full --release-manifest dist/releases/.json scripts/acceptance/run-local-acceptance.sh artifact-smoke --release-manifest dist/releases/.json +`adaptive` runs a three-Worker emulator-only concurrency ladder without changing +provider limits. `provider-burst` binds one video model to three emulator-only +platforms with concurrency limits 2/4/6 and validates burst queueing and routing. `full` executes P24/P28/P32 three times, the fault matrix, autoscaling/drain, 80% soak, 120% overload, and exact linux/amd64 artifact smoke. The load process runs outside K3s and splits requests 50/50 across both TLS entrances. @@ -51,9 +63,24 @@ cleanup() { kill "$active_load_pid" >/dev/null 2>&1 || true wait "$active_load_pid" >/dev/null 2>&1 || true fi + if [[ -n $active_worker_sampler_pid ]]; then + [[ -z $worker_sample_stop ]] || touch "$worker_sample_stop" + wait "$active_worker_sampler_pid" >/dev/null 2>&1 || true + fi + if [[ -n $active_resource_sampler_pid ]]; then + [[ -z $worker_sample_stop ]] || touch "$worker_sample_stop" + wait "$active_resource_sampler_pid" >/dev/null 2>&1 || true + fi + if [[ -n $active_provider_sampler_pid ]]; then + [[ -z $worker_sample_stop ]] || touch "$worker_sample_stop" + wait "$active_provider_sampler_pid" >/dev/null 2>&1 || true + fi if [[ $netem_active == true ]]; then "$script_dir/network-fault.sh" reset >/dev/null 2>&1 || true fi + if [[ $provider_burst_configured == true && -n ${runtime:-} ]]; then + restore_provider_burst_platforms >/dev/null 2>&1 || true + fi if [[ $status -ne 0 && -n ${runtime:-} && -n ${report_root:-} && -f ${runtime:-} && -f $snapshot ]]; then restore_profile_best_effort "${stable_profile:-P24}" if [[ -f ${failure_gate_file:-} && ! -L ${failure_gate_file:-} ]]; then @@ -206,10 +233,281 @@ run_load() { return 1 fi ((load_status == 0)) || return "$load_status" - jq -e '.schemaVersion == "acceptance-load-report/v1" and .secretSafe == true and .passed == true' \ + jq -e '.schemaVersion == "acceptance-load-report/v2" and .secretSafe == true and .passed == true' \ "$report_path" >/dev/null } +sample_worker_runtime() { + local output=$1 stop_file=$2 + : >"$output" + while [[ ! -e $stop_file ]]; do + database_query " +SELECT jsonb_build_object( + 'sampledAt', to_char(clock_timestamp() AT TIME ZONE 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.MS\"Z\"'), + 'workers', COALESCE(jsonb_agg(jsonb_build_object( + 'instanceId', instance_id, + 'allocatedCapacity', allocated_capacity, + 'safeCapacity', safe_capacity, + 'heavyCapacity', heavy_capacity, + 'activeTasks', active_tasks, + 'preparingTasks', preparing_tasks, + 'waitingUpstreamTasks', waiting_upstream_tasks, + 'finalizingTasks', finalizing_tasks, + 'pressureState', pressure_state, + 'pressureReason', pressure_reason + ) ORDER BY instance_id), '[]'::jsonb) +) +FROM gateway_worker_instances +WHERE status='active' AND heartbeat_at > now() - interval '30 seconds';" >>"$output" || true + sleep 1 + done +} + +wait_for_adaptive_workloads_settled() { + local deadline=$((SECONDS + 180)) stable_samples=0 pod_state + local api_total api_ready api_terminating worker_total worker_ready worker_terminating + while ((SECONDS < deadline)); do + pod_state=$(kubectl --context "$context" -n "$namespace" get pods \ + -l 'app.kubernetes.io/name in (easyai-api,easyai-worker)' -o json) || { + stable_samples=0 + sleep 2 + continue + } + read -r api_total api_ready api_terminating worker_total worker_ready worker_terminating < <( + jq -r ' + def total($name): [.items[] | select(.metadata.labels["app.kubernetes.io/name"] == $name and .metadata.deletionTimestamp == null)] | length; + def ready($name): [.items[] | select( + .metadata.labels["app.kubernetes.io/name"] == $name and + .metadata.deletionTimestamp == null and + .status.phase == "Running" and + ((.status.containerStatuses // []) | length) > 0 and + all(.status.containerStatuses[]; .ready == true) + )] | length; + def terminating($name): [.items[] | select(.metadata.labels["app.kubernetes.io/name"] == $name and .metadata.deletionTimestamp != null)] | length; + [total("easyai-api"), ready("easyai-api"), terminating("easyai-api"), total("easyai-worker"), ready("easyai-worker"), terminating("easyai-worker")] | @tsv + ' <<<"$pod_state" + ) + if [[ $api_total == 2 && $api_ready == 2 && $api_terminating == 0 && + $worker_total == 3 && $worker_ready == 3 && $worker_terminating == 0 ]]; then + stable_samples=$((stable_samples + 1)) + if ((stable_samples >= 3)); then + return 0 + fi + else + stable_samples=0 + fi + sleep 2 + done + fail_gate adaptive_workload_settle_timeout \ + "adaptive workloads did not settle at two Ready APIs and three Ready Workers without terminating Pods" +} + +configure_adaptive_workers() { + local hard_limit=${AI_GATEWAY_LOCAL_ADAPTIVE_HARD_LIMIT:-32} + local global_limit=$((hard_limit * 3)) + [[ $hard_limit =~ ^[0-9]+$ ]] && ((hard_limit >= 4 && hard_limit <= 256)) + kubectl --context "$context" -n "$namespace" patch configmap easyai-ai-gateway-config \ + --type=merge -p "$(jq -cn --arg hard "$hard_limit" --arg global "$global_limit" '{data:{ + AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:"false", + AI_GATEWAY_WORKER_LOAD_MODE:"adaptive", + AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT:$hard, + AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT:$global, + AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:$global + }}')" >/dev/null + kubectl --context "$context" -n "$namespace" rollout restart \ + deployment/easyai-capacity-controller >/dev/null + kubectl --context "$context" -n "$namespace" rollout status \ + deployment/easyai-capacity-controller --timeout=10m + kubectl --context "$context" -n "$namespace" scale deployment/easyai-worker-ningbo --replicas=2 >/dev/null + kubectl --context "$context" -n "$namespace" scale deployment/easyai-worker-hongkong --replicas=1 >/dev/null + for deployment in easyai-worker-ningbo easyai-worker-hongkong; do + kubectl --context "$context" -n "$namespace" set env deployment/"$deployment" \ + AI_GATEWAY_WORKER_LOAD_MODE=adaptive \ + AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT="$hard_limit" >/dev/null + kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=10m + done + for deployment in easyai-api-ningbo easyai-api-hongkong; do + kubectl --context "$context" -n "$namespace" rollout restart deployment/"$deployment" >/dev/null + kubectl --context "$context" -n "$namespace" rollout status deployment/"$deployment" --timeout=10m + done + wait_for_adaptive_workloads_settled + local ready + ready=$(kubectl --context "$context" -n "$namespace" get pods -l app.kubernetes.io/name=easyai-worker -o json | + jq '[.items[] | select(.status.phase=="Running" and any(.status.containerStatuses[]?; .ready==true))] | length') + [[ $ready == 3 ]] || fail_gate adaptive_worker_count "adaptive validation requires exactly three Ready Workers" +} + +run_adaptive_level() { + local concurrency=$1 passed=true etcd_healthy=true + local gemini_report="$report_root/adaptive-load-${concurrency}-gemini.json" + local video_report="$report_root/adaptive-load-${concurrency}-video.json" + verify_local_control_plane_identity "$context" "$cluster_name" "$identity_file" || + fail_gate local_control_plane_restarted "local K3s server identity changed before adaptive level $concurrency" + record_local_control_plane_identity "$context" "$cluster_name" "$identity_file" || + fail_gate local_control_plane_identity "failed to record adaptive level $concurrency control-plane baseline" + if ! run_load gemini-multi-image "$gemini_report" -requests "$concurrency"; then + passed=false + fi + if [[ $passed == true ]] && ! run_load video-throughput "$video_report" -requests "$concurrency"; then + passed=false + fi + local critical + critical=$(database_query "SELECT count(*) FROM gateway_worker_instances WHERE status='active' AND pressure_state='critical' AND heartbeat_at > now() - interval '30 seconds';") + if ((critical > 0)); then + passed=false + fi + if ! verify_local_etcd_runtime_logs "$cluster_name" "$identity_file"; then + passed=false + etcd_healthy=false + record_local_control_plane_identity "$context" "$cluster_name" "$identity_file" || return 1 + fi + jq -n --argjson concurrency "$concurrency" --argjson passed "$passed" \ + --argjson etcdHealthy "$etcd_healthy" \ + --arg gemini "$(basename "$gemini_report")" --arg video "$(basename "$video_report")" \ + '{concurrency:$concurrency,passed:$passed,etcdHealthy:$etcdHealthy,reports:[$gemini,$video]}' \ + >"$report_root/adaptive-level-${concurrency}.json" + chmod 0600 "$report_root/adaptive-level-${concurrency}.json" + [[ $passed == true ]] +} + +build_adaptive_report() { + local output=$1 samples=$2 resources=$3 highest=$4 failed=${5:-0} + local levels workers distribution routing limits concurrency_peaks leaks duplicates sample_summary resource_summary phases + levels=$(jq -s 'sort_by(.concurrency)' "$report_root"/adaptive-level-*.json) + workers=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'instanceId',instance_id,'podName',pod_name,'site',site,'status',status, + 'allocatedCapacity',allocated_capacity,'hardCapacityLimit',hard_capacity_limit, + 'safeCapacity',safe_capacity,'heavyCapacity',heavy_capacity, + 'activeTasks',active_tasks,'preparingTasks',preparing_tasks, + 'waitingUpstreamTasks',waiting_upstream_tasks,'finalizingTasks',finalizing_tasks, + 'pressureState',pressure_state,'pressureReason',pressure_reason, + 'loadSampledAt',load_sampled_at +) ORDER BY site,instance_id),'[]'::jsonb) +FROM gateway_worker_instances +WHERE heartbeat_at > now() - interval '30 seconds';") + distribution=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object('workerInstanceId',worker_id,'tasks',tasks) ORDER BY worker_id),'[]'::jsonb) +FROM ( + SELECT COALESCE( + regexp_replace(job.attempted_by[cardinality(job.attempted_by)], '-exec-[^-]+-[^-]+$', ''), + 'unclaimed' + ) worker_id, + count(*) tasks + FROM gateway_tasks task + LEFT JOIN river_job job ON job.id = task.river_job_id + WHERE task.acceptance_run_id='$run_id'::uuid + GROUP BY worker_id +) grouped;") + routing=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'platformId',platform_id,'platformModelId',platform_model_id,'attempts',attempts +) ORDER BY platform_id,platform_model_id),'[]'::jsonb) +FROM (SELECT platform_id::text,platform_model_id::text,count(*) attempts + FROM gateway_task_attempts WHERE task_id IN + (SELECT id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) + GROUP BY platform_id,platform_model_id) grouped;") + limits=$(database_query " +SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'scopeType',scope_type,'scopeKey',scope_key,'metric',metric, + 'limit',limit_value,'used',used_value,'reserved',reserved_value, + 'windowStart',window_start,'resetAt',reset_at +) ORDER BY scope_type,scope_key,metric,window_start),'[]'::jsonb) +FROM gateway_rate_limit_counters +WHERE scope_key LIKE 'acceptance:$run_id:%';") + concurrency_peaks=$(database_query " +WITH events AS ( + SELECT scope_type,scope_key,limit_value,acquired_at event_at,lease_value delta FROM gateway_concurrency_leases + WHERE scope_key LIKE 'acceptance:$run_id:%' + UNION ALL + SELECT scope_type,scope_key,limit_value,LEAST(COALESCE(released_at,expires_at),expires_at),-lease_value + FROM gateway_concurrency_leases WHERE scope_key LIKE 'acceptance:$run_id:%' +), points AS ( + SELECT scope_type,scope_key,limit_value, + SUM(delta) OVER (PARTITION BY scope_type,scope_key ORDER BY event_at,delta ROWS UNBOUNDED PRECEDING) active + FROM events +), peaks AS ( + SELECT scope_type,scope_key,MAX(limit_value) limit_value,COALESCE(MAX(active),0) peak + FROM points GROUP BY scope_type,scope_key +) +SELECT COALESCE(jsonb_agg(jsonb_build_object('scopeType',scope_type,'scopeKey',scope_key,'limit',limit_value,'peak',peak) + ORDER BY scope_type,scope_key),'[]'::jsonb) FROM peaks;") + leaks=$(database_query " +SELECT jsonb_build_object( + 'activeConcurrencyLeases',(SELECT count(*) FROM gateway_concurrency_leases WHERE scope_key LIKE 'acceptance:$run_id:%' AND released_at IS NULL AND expires_at>now()), + 'activeRateReservations',(SELECT count(*) FROM gateway_rate_limit_reservations WHERE scope_key LIKE 'acceptance:$run_id:%' AND status='reserved'), + 'activeRiverJobs',(SELECT count(*) FROM river_job job JOIN gateway_tasks task ON task.river_job_id=job.id WHERE task.acceptance_run_id='$run_id'::uuid AND job.state IN ('available','pending','retryable','running','scheduled')) +);") + duplicates=$(database_query " +SELECT jsonb_build_object( + 'remoteTaskIds',(SELECT count(*) FROM (SELECT remote_task_id FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid AND remote_task_id IS NOT NULL GROUP BY remote_task_id HAVING count(*)>1) duplicate), + 'billingTransactions',(SELECT count(*) FROM (SELECT reference_id,transaction_type FROM gateway_wallet_transactions WHERE reference_type='gateway_task' AND reference_id IN (SELECT id::text FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid) GROUP BY reference_id,transaction_type HAVING count(*)>1) duplicate) +);") + sample_summary=$(jq -s '[.[].workers[]] | group_by(.instanceId) | map({ + instanceId:.[0].instanceId, + peakActiveTasks:(map(.activeTasks)|max), + peakPreparingTasks:(map(.preparingTasks)|max), + peakWaitingUpstreamTasks:(map(.waitingUpstreamTasks)|max), + peakFinalizingTasks:(map(.finalizingTasks)|max), + minSafeCapacity:(map(.safeCapacity)|min), + maxSafeCapacity:(map(.safeCapacity)|max), + pressureStates:(map(.pressureState)|unique) + })' "$samples") + resource_summary=$(jq -Rn ' + def number_without($suffix): rtrimstr($suffix) | tonumber; + def memory_mib: + if endswith("Gi") then number_without("Gi") * 1024 + elif endswith("Mi") then number_without("Mi") + elif endswith("Ki") then number_without("Ki") / 1024 + else tonumber end; + [inputs | split(",") | select(.[0] != "timestamp") | { + scope:.[1],name:.[2],cpu:.[3],memory:.[4] + }] + | group_by([.scope,.name]) + | map(if .[0].scope == "node" then { + scope:"node",name:.[0].name, + maxCpuPercent:(map(.cpu | number_without("%")) | max), + maxMemoryPercent:(map(.memory | number_without("%")) | max) + } else { + scope:"pod",name:.[0].name, + maxCpuMillicores:(map(.cpu | if endswith("m") then number_without("m") else tonumber * 1000 end) | max), + maxMemoryMiB:(map(.memory | memory_mib) | max) + } end)' <"$resources") + phases=$(jq -s '[.[] | .phases[]]' "$report_root"/adaptive-load-*.json) + local passed=true bottleneck=configured_probe_ceiling + if ((highest < 1)); then passed=false; fi + if ((failed > 0)); then bottleneck=worker_or_gateway_limit; fi + if jq -e 'any(.[]; .etcdHealthy == false)' <<<"$levels" >/dev/null; then + bottleneck=local_control_plane_limit + fi + if ! jq -e 'length == 3' <<<"$workers" >/dev/null; then passed=false; fi + if ! jq -e 'all(.[]; ((.used + .reserved) <= .limit))' <<<"$limits" >/dev/null; then passed=false; fi + if ! jq -e 'all(.[]; (.limit == null or .peak <= .limit))' <<<"$concurrency_peaks" >/dev/null; then passed=false; fi + if ! jq -e '.activeConcurrencyLeases==0 and .activeRateReservations==0 and .activeRiverJobs==0' <<<"$leaks" >/dev/null; then passed=false; fi + if ! jq -e '.remoteTaskIds==0 and .billingTransactions==0' <<<"$duplicates" >/dev/null; then passed=false; fi + jq -n \ + --arg runId "$run_id" --arg startedAt "$(date -u '+%Y-%m-%dT%H:%M:%SZ')" \ + --arg bottleneck "$bottleneck" --argjson passed "$passed" \ + --argjson highest "$highest" --argjson failed "$failed" \ + --argjson levels "$levels" --argjson phases "$phases" \ + --argjson workers "$workers" --argjson distribution "$distribution" \ + --argjson routing "$routing" --argjson limits "$limits" \ + --argjson concurrencyPeaks "$concurrency_peaks" --argjson workerPeaks "$sample_summary" \ + --argjson resourcePeaks "$resource_summary" \ + --argjson leaks "$leaks" --argjson duplicates "$duplicates" \ + '{ + schemaVersion:"acceptance-load-report/v2",runId:$runId,profile:"adaptive-capacity", + startedAt:$startedAt,finishedAt:(now|todateiso8601),passed:$passed,secretSafe:true, + phases:$phases, + capacity:{levels:$levels,highestStableConcurrency:$highest,firstFailedConcurrency:$failed,bottleneck:$bottleneck}, + cluster:{workers:$workers,workerPeaks:$workerPeaks,taskDistribution:$distribution, + platformRouting:$routing,rateLimits:$limits,concurrencyPeaks:$concurrencyPeaks, + resourcePeaks:$resourcePeaks,leaks:$leaks,duplicates:$duplicates} + }' >"$output" + chmod 0600 "$output" + jq -e '.passed==true and .secretSafe==true' "$output" >/dev/null +} + sample_resources() { local output=$1 stop_file=$2 printf 'timestamp,scope,name,cpu,memory\n' >"$output" @@ -567,6 +865,64 @@ quick() { echo "local_acceptance_quick=PASS run_id=$run_id report=$report_root/quick.json" } +adaptive() { + local hard_limit=${AI_GATEWAY_LOCAL_ADAPTIVE_HARD_LIMIT:-32} + local max_probe=${AI_GATEWAY_LOCAL_ADAPTIVE_MAX_CONCURRENCY:-$((hard_limit * 3))} + local level=4 highest=0 failed=0 low high mid + local samples="$report_root/adaptive-worker-samples.ndjson" + local resources="$report_root/adaptive-resources.csv" + worker_sample_stop="$report_root/adaptive-worker-samples.stop" + [[ ! -e $samples && ! -e $resources && ! -e $worker_sample_stop ]] || fail_gate acceptance_report_exists "adaptive evidence already exists" + current_phase=adaptive_workers + configure_adaptive_workers + sample_worker_runtime "$samples" "$worker_sample_stop" & + active_worker_sampler_pid=$! + sample_resources "$resources" "$worker_sample_stop" & + active_resource_sampler_pid=$! + while ((level <= max_probe)); do + current_phase="adaptive_capacity_$level" + if run_adaptive_level "$level"; then + highest=$level + level=$((level * 2)) + else + failed=$level + break + fi + done + if ((failed > 0 && highest > 0)); then + low=$((highest + 1)) + high=$((failed - 1)) + while ((low <= high)); do + mid=$(((low + high) / 2)) + current_phase="adaptive_capacity_refine_$mid" + if run_adaptive_level "$mid"; then + highest=$mid + low=$((mid + 1)) + else + failed=$mid + high=$((mid - 1)) + fi + done + fi + touch "$worker_sample_stop" + wait "$active_worker_sampler_pid" + wait "$active_resource_sampler_pid" + active_worker_sampler_pid= + active_resource_sampler_pid= + worker_sample_stop= + local ready_workers registered_workers + ready_workers=$(kubectl --context "$context" -n "$namespace" get pods \ + -l app.kubernetes.io/name=easyai-worker -o json | + jq '[.items[] | select(.status.phase=="Running" and any(.status.containerStatuses[]?; .ready==true))] | length') + registered_workers=$(database_query "SELECT count(*) FROM gateway_worker_instances WHERE status='active' AND heartbeat_at > now() - interval '30 seconds';") + [[ $ready_workers == 3 && $registered_workers == 3 ]] || + fail_gate adaptive_worker_count "adaptive validation requires three Ready and registered Workers throughout the final gate" + verify_hard_gates + current_phase=adaptive_report + build_adaptive_report "$report_root/adaptive-capacity.json" "$samples" "$resources" "$highest" "$failed" + echo "local_acceptance_adaptive=PASS run_id=$run_id highest_stable_concurrency=$highest first_failed_concurrency=$failed report=$report_root/adaptive-capacity.json" +} + full() { local manifest=$1 profile repetition quick @@ -610,7 +966,7 @@ full() { command=${1:-} shift || true case $command in - quick) + quick|adaptive|provider-burst) [[ $# -eq 0 ]] || { usage >&2; exit 64; } ;; artifact-smoke|full) @@ -628,6 +984,12 @@ case $command in quick) quick ;; + adaptive) + adaptive + ;; + provider-burst) + provider_burst + ;; artifact-smoke) artifact_smoke "$2" ;; diff --git a/scripts/cluster/run-production-acceptance.sh b/scripts/cluster/run-production-acceptance.sh index 4a03539..846fe2b 100755 --- a/scripts/cluster/run-production-acceptance.sh +++ b/scripts/cluster/run-production-acceptance.sh @@ -15,6 +15,8 @@ Usage: scripts/cluster/run-production-acceptance.sh \ --execute dist/releases/.json \ --skip-local-acceptance + scripts/cluster/run-production-acceptance.sh \ + --execute-single-node dist/releases/.json scripts/cluster/run-production-acceptance.sh \ --promote dist/releases/.json --run-id @@ -29,6 +31,10 @@ The explicit --skip-local-acceptance form records both local stages as skipped with a user-directed waiver. It never records them as passed and preserves all online simulation, real canary, resource, consistency, and release CAS gates. +The --execute-single-node form is a Ningbo-only diagnostic. It keeps production +traffic in validation, routes acceptance candidates only to the protocol +emulator, never runs real-canary, and leaves the Run non-promotable. + Required private environment values: AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN, or AI_GATEWAY_ONLINE_ACCOUNT/PASSWORD @@ -54,6 +60,7 @@ if [[ ${1:-} == --promote ]]; then fi skip_local_acceptance=false +single_node_acceptance=false local_acceptance_report= if [[ ${1:-} == --execute && ${3:-} == --local-report && $# -eq 4 ]]; then release_manifest=$2 @@ -61,6 +68,10 @@ if [[ ${1:-} == --execute && ${3:-} == --local-report && $# -eq 4 ]]; then elif [[ ${1:-} == --execute && ${3:-} == --skip-local-acceptance && $# -eq 3 ]]; then release_manifest=$2 skip_local_acceptance=true +elif [[ ${1:-} == --execute-single-node && $# -eq 2 ]]; then + release_manifest=$2 + skip_local_acceptance=true + single_node_acceptance=true else usage >&2 exit 64 @@ -79,6 +90,18 @@ fi load_cluster_env require_commands curl git go jq node openssl sed shasum +if [[ $single_node_acceptance == true ]]; then + AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO=1 + AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG=0 + AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MIN_REPLICAS_NINGBO=1 + AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MIN_REPLICAS_HONGKONG=0 + AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_NINGBO=1 + AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_HONGKONG=0 + AI_GATEWAY_ACCEPTANCE_IDENTITY_SHARDS=${AI_GATEWAY_ACCEPTANCE_IDENTITY_SHARDS:-16} + AI_GATEWAY_ACCEPTANCE_WORKER_MEMORY_REQUEST_MIB=512 + AI_GATEWAY_ACCEPTANCE_DATABASE_MAX_CONN_IDLE_SECONDS=30 +fi + : "${AI_GATEWAY_ACCEPTANCE_ADMIN_TOKEN:=}" : "${AI_GATEWAY_ACCEPTANCE_API_KEY:=}" : "${AI_GATEWAY_ACCEPTANCE_API_KEY_ID:=}" @@ -109,15 +132,45 @@ require_commands curl git go jq node openssl sed shasum : "${AI_GATEWAY_ACCEPTANCE_OVERLOAD_DURATION:=10m}" : "${AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME:=}" : "${AI_GATEWAY_ACCEPTANCE_GATEWAYS:=}" +: "${AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_BASELINE_SLOTS:=8}" +: "${AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_SLOTS:=8}" +: "${AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_REQUESTS:=96}" +: "${AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_VIDEO_SLOTS:=8 12 16 24 32 40 48}" : "${AI_GATEWAY_ACCEPTANCE_REAL_IMAGE_URLS:=}" : "${AI_GATEWAY_ACCEPTANCE_SNAPSHOT_DATABASE_URL:=}" : "${AI_GATEWAY_ONLINE_ACCOUNT:=}" : "${AI_GATEWAY_ONLINE_PASSWORD:=}" : "${AI_GATEWAY_ONLINE_BASE_URL:=}" if [[ -z $AI_GATEWAY_ACCEPTANCE_GATEWAYS ]]; then - AI_GATEWAY_ACCEPTANCE_GATEWAYS="https://${CLUSTER_NINGBO_HOST#root@},https://${CLUSTER_HONGKONG_HOST#root@}" - : "${AI_GATEWAY_DEPLOY_DOMAIN:?}" - AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME=${AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME:-$AI_GATEWAY_DEPLOY_DOMAIN} + if [[ $single_node_acceptance == true ]]; then + AI_GATEWAY_ACCEPTANCE_GATEWAYS='http://10.77.0.1:18089' + else + AI_GATEWAY_ACCEPTANCE_GATEWAYS="https://${CLUSTER_NINGBO_HOST#root@},https://${CLUSTER_HONGKONG_HOST#root@}" + : "${AI_GATEWAY_DEPLOY_DOMAIN:?}" + AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME=${AI_GATEWAY_ACCEPTANCE_GATEWAY_TLS_SERVER_NAME:-$AI_GATEWAY_DEPLOY_DOMAIN} + fi +fi +single_node_video_slots=() +if [[ $single_node_acceptance == true ]]; then + [[ $AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_BASELINE_SLOTS =~ ^[1-9][0-9]*$ && + $AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_BASELINE_SLOTS -le 128 && + $AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_SLOTS =~ ^[1-9][0-9]*$ && + $AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_SLOTS -le 128 && + $AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_REQUESTS =~ ^[1-9][0-9]*$ ]] || { + echo 'single-node baseline, GEMINI slots, and GEMINI requests must be positive bounded integers' >&2 + exit 1 + } + read -r -a single_node_video_slots <<<"$AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_VIDEO_SLOTS" + (( ${#single_node_video_slots[@]} > 0 )) || { + echo 'AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_VIDEO_SLOTS must contain at least one slot count' >&2 + exit 1 + } + for single_node_slot in "${single_node_video_slots[@]}"; do + [[ $single_node_slot =~ ^[1-9][0-9]*$ && $single_node_slot -le 128 ]] || { + echo 'single-node video slot counts must be integers between 1 and 128' >&2 + exit 1 + } + done fi [[ $AI_GATEWAY_ACCEPTANCE_API_DATABASE_MAX_CONNS =~ ^[1-9][0-9]*$ && $AI_GATEWAY_ACCEPTANCE_API_DATABASE_MAX_CONNS -le 256 ]] || { @@ -189,16 +242,27 @@ fi $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MIN_REPLICAS_HONGKONG =~ ^[0-9]+$ && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_NINGBO =~ ^[0-9]+$ && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_HONGKONG =~ ^[0-9]+$ && - $((AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO + AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG)) -eq 2 && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MIN_REPLICAS_NINGBO -le $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO && $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO -le $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_NINGBO && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MIN_REPLICAS_HONGKONG -le $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG && $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG -le $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_HONGKONG && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_NINGBO -le 16 && $AI_GATEWAY_ACCEPTANCE_AUTOSCALING_MAX_REPLICAS_HONGKONG -le 16 ]] || { - echo 'acceptance requires exactly two baseline Workers and ordered 0..16 per-site autoscaling bounds' >&2 + echo 'acceptance baseline Worker topology or per-site autoscaling bounds are invalid' >&2 exit 1 } +if [[ $single_node_acceptance == true ]]; then + [[ $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO == 1 && + $AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG == 0 ]] || { + echo 'single-node acceptance requires one Ningbo Worker and zero Hong Kong Workers' >&2 + exit 1 + } +else + (( AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO + AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG == 2 )) || { + echo 'production acceptance requires exactly two baseline Workers' >&2 + exit 1 + } +fi node "$cluster_root/scripts/release-manifest.mjs" validate "$release_manifest" >/dev/null release_sha=$(node "$cluster_root/scripts/release-manifest.mjs" get "$release_manifest" sourceSha) @@ -223,24 +287,30 @@ if [[ $skip_local_acceptance != true ]]; then exit 1 } fi -[[ -z $(git -C "$cluster_root" status --short) ]] || { - echo 'production acceptance requires a clean release working copy' >&2 - exit 1 -} +if [[ $single_node_acceptance != true ]]; then + [[ -z $(git -C "$cluster_root" status --short) ]] || { + echo 'production acceptance requires a clean release working copy' >&2 + exit 1 + } +else + echo 'single_node_acceptance_tool=working_tree diagnostic_only=true certification=false' +fi acceptance_tool_sha=$(git -C "$cluster_root" rev-parse HEAD) if [[ $acceptance_tool_sha != "$release_sha" ]]; then git -C "$cluster_root" merge-base --is-ancestor "$release_sha" "$acceptance_tool_sha" || { echo 'production release must be an ancestor of the acceptance tool HEAD' >&2 exit 1 } - acceptance_tool_delta=$(node "$cluster_root/scripts/release-components.mjs" \ - "$release_sha" "$acceptance_tool_sha") - [[ $(jq -r '.components' <<<"$acceptance_tool_delta") == none && - $(jq -r '.migrationsChanged' <<<"$acceptance_tool_delta") == false ]] || { - echo 'acceptance tool HEAD contains runtime or migration changes beyond the production release' >&2 - exit 1 - } - echo "acceptance_tool_delta=PASS release=$release_sha tool_sha=$acceptance_tool_sha runtime_changes=false" + if [[ $single_node_acceptance != true ]]; then + acceptance_tool_delta=$(node "$cluster_root/scripts/release-components.mjs" \ + "$release_sha" "$acceptance_tool_sha") + [[ $(jq -r '.components' <<<"$acceptance_tool_delta") == none && + $(jq -r '.migrationsChanged' <<<"$acceptance_tool_delta") == false ]] || { + echo 'acceptance tool HEAD contains runtime or migration changes beyond the production release' >&2 + exit 1 + } + echo "acceptance_tool_delta=PASS release=$release_sha tool_sha=$acceptance_tool_sha runtime_changes=false" + fi fi namespace=${AI_GATEWAY_K3S_NAMESPACE:-easyai} @@ -262,6 +332,7 @@ run_id= report_root= stable_profile=P24 active_profile=P24 +single_node_stable_slots=4 failure_reason= failure_gate_id= failure_recorded=false @@ -272,6 +343,7 @@ acceptance_participants_json='[]' AI_GATEWAY_ACCEPTANCE_API_KEYS=$AI_GATEWAY_ACCEPTANCE_API_KEY acceptance_load_binary=$temporary_root/easyai-ai-gateway-acceptance-load acceptance_load_linux_binary=$temporary_root/easyai-ai-gateway-acceptance-load-linux-amd64 +acceptance_emulator_linux_binary=$temporary_root/easyai-ai-gateway-acceptance-emulator-linux-amd64 acceptance_snapshot_binary=$temporary_root/easyai-ai-gateway-acceptance-snapshot current_production_snapshot=$temporary_root/current-production-snapshot.json local_snapshot_config_hash= @@ -291,6 +363,8 @@ video_admitted_throughput= certified_max_replicas_ningbo=$AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_NINGBO certified_max_replicas_hongkong=$AI_GATEWAY_ACCEPTANCE_BASE_REPLICAS_HONGKONG runtime_observation_started_at= +single_node_emulator_url= +single_node_emulator_installed=false cleanup() { local status=$? @@ -306,9 +380,16 @@ cleanup() { if [[ $remote_load_drivers_installed == true ]]; then cleanup_remote_load_drivers >/dev/null 2>&1 || true fi + if [[ $single_node_emulator_installed == true ]]; then + cleanup_single_node_emulator >/dev/null 2>&1 || true + fi if (( status != 0 )) && [[ -n $run_id && $failure_recorded != true ]]; then set +e - apply_capacity_profile "$stable_profile" >/dev/null 2>&1 + if [[ $single_node_acceptance == true ]]; then + apply_single_node_capacity "$single_node_stable_slots" >/dev/null 2>&1 + else + apply_capacity_profile "$stable_profile" >/dev/null 2>&1 + fi [[ -n $failure_gate_id ]] || failure_gate_id=workflow_unexpected_exit mark_run_failed "${failure_reason:-acceptance workflow exited unexpectedly}" set -e @@ -361,6 +442,10 @@ bootstrap_acceptance_admin_token -trimpath -o "$acceptance_load_binary" ./cmd/acceptance-load env -u AI_GATEWAY_TEST_DATABASE_URL CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \ -trimpath -o "$acceptance_load_linux_binary" ./cmd/acceptance-load + if [[ $single_node_acceptance == true ]]; then + env -u AI_GATEWAY_TEST_DATABASE_URL CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build \ + -trimpath -o "$acceptance_emulator_linux_binary" ./cmd/acceptance-emulator + fi env -u AI_GATEWAY_TEST_DATABASE_URL go build \ -trimpath -o "$acceptance_snapshot_binary" ./cmd/acceptance-snapshot ) @@ -1386,14 +1471,93 @@ WHERE NOT EXISTS ( echo "acceptance_model_access=PASS resources=$expected" } +cleanup_single_node_emulator() { + [[ $run_token =~ ^[0-9a-f]{64}$ ]] || return 0 + local suffix=${run_token:0:12} + local binary_path=/root/easyai-acceptance-emulator-$suffix + local pid_path=$binary_path.pid + local log_path=$binary_path.log + cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$binary_path" "$pid_path" "$log_path" <<'REMOTE' || true +set -euo pipefail +binary_path=$1 +pid_path=$2 +log_path=$3 +if [[ -f $pid_path && ! -L $pid_path ]]; then + pid=$(<"$pid_path") + if [[ $pid =~ ^[0-9]+$ && -e /proc/$pid/exe && $(readlink -f "/proc/$pid/exe") == "$binary_path" ]]; then + kill -TERM "$pid" + for _ in {1..20}; do + kill -0 "$pid" >/dev/null 2>&1 || break + sleep 0.1 + done + fi +fi +for path in "$binary_path" "$pid_path" "$log_path"; do + [[ ! -e $path ]] || unlink "$path" +done +REMOTE + single_node_emulator_installed=false +} + +start_single_node_emulator() { + [[ $run_token =~ ^[0-9a-f]{64}$ ]] + local suffix=${run_token:0:12} + local binary_path=/root/easyai-acceptance-emulator-$suffix + local pid_path=$binary_path.pid + local log_path=$binary_path.log + local local_digest remote_digest node_ip + local_digest=$(shasum -a 256 "$acceptance_emulator_linux_binary" | awk '{print $1}') + cluster_scp "$acceptance_emulator_linux_binary" "$CLUSTER_NINGBO_HOST:$binary_path" >/dev/null + remote_digest=$(cluster_ssh "$CLUSTER_NINGBO_HOST" "chmod 0755 '$binary_path'; sha256sum '$binary_path'" | awk '{print $1}') + [[ $remote_digest == "$local_digest" ]] + single_node_emulator_installed=true + cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$binary_path" "$pid_path" "$log_path" <<'REMOTE' +set -euo pipefail +binary_path=$1 +pid_path=$2 +log_path=$3 +if ss -ltn '( sport = :18090 )' | tail -n +2 | grep -q .; then + echo 'single-node emulator port 18090 is already in use' >&2 + exit 1 +fi +nohup env HTTP_ADDR=:18090 "$binary_path" "$log_path" 2>&1 & +pid=$! +printf '%s\n' "$pid" >"$pid_path" +chmod 0600 "$pid_path" "$log_path" +for _ in {1..50}; do + if curl -fsS --max-time 2 http://127.0.0.1:18090/healthz >/dev/null; then + exit 0 + fi + kill -0 "$pid" >/dev/null 2>&1 || exit 1 + sleep 0.2 +done +exit 1 +REMOTE + node_ip=$(remote_kubectl get node easyai-ningbo -o 'jsonpath={.status.addresses[?(@.type=="InternalIP")].address}') + [[ $node_ip =~ ^[0-9a-fA-F:.]+$ ]] + single_node_emulator_url=http://$node_ip:18090 + remote_kubectl exec -n "$namespace" deployment/easyai-api-ningbo -- \ + wget -qO- "$single_node_emulator_url/healthz" >/dev/null + echo 'single_node_protocol_emulator=PASS placement=ningbo-host digest_verified=true pod_reachable=true' +} + deploy_protocol_emulator() { sed "s|image: easyai-api|image: $api_image|" \ "$cluster_root/deploy/kubernetes/acceptance/protocol-emulator.yaml" | cluster_ssh "$CLUSTER_NINGBO_HOST" 'k3s kubectl apply -f -' >/dev/null + if [[ $single_node_acceptance == true ]]; then + remote_kubectl patch deployment easyai-acceptance-emulator -n "$namespace" \ + --type=json -p='[{"op":"replace","path":"/spec/template/spec/nodeSelector","value":{"kubernetes.io/hostname":"easyai-ningbo"}}]' >/dev/null + remote_kubectl patch deployment easyai-acceptance-callback-collector -n "$namespace" \ + --type=json -p='[{"op":"replace","path":"/spec/template/spec/nodeSelector","value":{"kubernetes.io/hostname":"easyai-ningbo"}}]' >/dev/null + fi remote_kubectl rollout status deployment/easyai-acceptance-emulator \ -n "$namespace" --timeout=300s remote_kubectl rollout status deployment/easyai-acceptance-callback-collector \ -n "$namespace" --timeout=300s + if [[ $single_node_acceptance == true ]]; then + start_single_node_emulator + fi } is_release_ancestor() { @@ -1408,7 +1572,11 @@ create_and_activate_run() { local activate_response=$temporary_root/activate-run.json local traffic_response=$temporary_root/pre-activate-traffic.json local previous_run_response=$temporary_root/previous-run.json - local body + local body emulator_url='http://easyai-acceptance-emulator.easyai.svc.cluster.local:8090' + if [[ $single_node_acceptance == true ]]; then + [[ $single_node_emulator_url == http://* ]] + emulator_url=$single_node_emulator_url + fi body=$(jq -cn \ --arg releaseSha "$release_sha" \ --arg apiDigest "$api_digest" \ @@ -1416,7 +1584,7 @@ create_and_activate_run() { --arg apiKeyID "$AI_GATEWAY_ACCEPTANCE_API_KEY_ID" \ --arg userID "$AI_GATEWAY_ACCEPTANCE_USER_ID" \ --arg token "$run_token" \ - --arg emulatorURL 'http://easyai-acceptance-emulator.easyai.svc.cluster.local:8090' \ + --arg emulatorURL "$emulator_url" \ --arg callbackURL 'http://easyai-acceptance-callback-collector.easyai.svc.cluster.local:8091/callbacks' \ --argjson participants "$acceptance_participants_json" \ '{ @@ -2271,8 +2439,10 @@ cleanup_remote_load_drivers() { hongkong_env=$(remote_load_env_path hongkong) cluster_ssh "$CLUSTER_NINGBO_HOST" \ "pkill -TERM -f '^$binary_path( |$)' >/dev/null 2>&1 || true; [[ ! -e '$binary_path' ]] || unlink '$binary_path'; [[ ! -e '$ningbo_env' ]] || unlink '$ningbo_env'" || true - cluster_ssh "$CLUSTER_HONGKONG_HOST" \ - "pkill -TERM -f '^$binary_path( |$)' >/dev/null 2>&1 || true; [[ ! -e '$binary_path' ]] || unlink '$binary_path'; [[ ! -e '$hongkong_env' ]] || unlink '$hongkong_env'" || true + if [[ $single_node_acceptance != true ]]; then + cluster_ssh "$CLUSTER_HONGKONG_HOST" \ + "pkill -TERM -f '^$binary_path( |$)' >/dev/null 2>&1 || true; [[ ! -e '$binary_path' ]] || unlink '$binary_path'; [[ ! -e '$hongkong_env' ]] || unlink '$hongkong_env'" || true + fi remote_load_drivers_installed=false } @@ -3268,6 +3438,363 @@ certify_worker_resource_requests() { }' >"$report_root/certified-worker-resources.json" } +apply_single_node_capacity() { + local slots=$1 + [[ $slots =~ ^[0-9]+$ ]] && (( slots >= 4 && slots <= 48 )) || return 1 + local worker_pool=$((slots + 12)) + local target_outstanding=$((slots * 2)) + local config_patch + config_patch=$(jq -nc \ + --arg slots "$slots" \ + --arg workerPool "$worker_pool" \ + --arg targetOutstanding "$target_outstanding" \ + '{data:{ + AI_GATEWAY_WORKER_REPLICAS_NINGBO:"1", + AI_GATEWAY_WORKER_REPLICAS_HONGKONG:"0", + AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO:"1", + AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG:"0", + AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO:"1", + AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG:"0", + AI_GATEWAY_WORKER_AUTOSCALING_ENABLED:"false", + AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT:$slots, + AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT:$slots, + AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT:$slots, + AI_GATEWAY_WORKER_DATABASE_MAX_CONNS:$workerPool, + AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY:$slots, + AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY:$slots, + AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA:$targetOutstanding + }}') + remote_kubectl label node easyai-ningbo easyai.io/worker=true --overwrite >/dev/null + remote_kubectl scale deployment/easyai-capacity-controller -n "$namespace" --replicas=0 >/dev/null + remote_kubectl scale deployment/easyai-worker-hongkong -n "$namespace" --replicas=0 >/dev/null + remote_kubectl patch configmap easyai-ai-gateway-config -n "$namespace" \ + --type=merge -p "$config_patch" >/dev/null + remote_kubectl set env deployment/easyai-worker-ningbo -n "$namespace" \ + "AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT=$slots" \ + "AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=$slots" \ + "AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$slots" \ + "AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE=$AI_GATEWAY_ACCEPTANCE_ASYNC_ADMISSION_MICROBATCH_SIZE" \ + "AI_GATEWAY_DATABASE_MAX_CONNS=$worker_pool" \ + 'AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS=4' \ + "AI_GATEWAY_DATABASE_RIVER_MAX_CONNS=$AI_GATEWAY_ACCEPTANCE_WORKER_RIVER_MAX_CONNS" \ + 'AI_GATEWAY_DATABASE_MIN_IDLE_CONNS=4' \ + 'AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS=30' \ + "AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY=$slots" \ + "AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY=$slots" \ + 'AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY=2' >/dev/null + remote_kubectl set resources deployment/easyai-worker-ningbo -n "$namespace" \ + --containers=worker \ + --requests="cpu=${AI_GATEWAY_ACCEPTANCE_WORKER_CPU_REQUEST_MILLICORES}m,memory=${AI_GATEWAY_ACCEPTANCE_WORKER_MEMORY_REQUEST_MIB}Mi" \ + --limits='cpu=2,memory=2Gi' >/dev/null + remote_kubectl scale deployment/easyai-worker-ningbo -n "$namespace" --replicas=1 >/dev/null + remote_kubectl rollout status deployment/easyai-worker-ningbo -n "$namespace" --timeout=300s + remote_kubectl set env deployment/easyai-api-ningbo -n "$namespace" \ + "AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=$slots" \ + "AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT=$slots" >/dev/null + remote_kubectl rollout status deployment/easyai-api-ningbo -n "$namespace" --timeout=300s + + local deadline=$((SECONDS + 90)) allocation + while (( SECONDS < deadline )); do + allocation=$(database_query " +SELECT count(*)||':'||COALESCE(max(allocated_capacity),0) +FROM gateway_worker_instances +WHERE status='active' AND heartbeat_at > now()-interval '30 seconds';") + if [[ $allocation == "1:$slots" ]]; then + local metrics_ready=false + for _ in {1..30}; do + if remote_kubectl top node easyai-ningbo --no-headers >/dev/null 2>&1 && + [[ -n $(remote_kubectl top pods -n "$namespace" \ + -l 'app.kubernetes.io/name=easyai-worker,easyai.io/site=ningbo' \ + --no-headers 2>/dev/null) ]]; then + metrics_ready=true + break + fi + sleep 2 + done + [[ $metrics_ready == true ]] || { + echo 'single-node metrics did not become ready after Worker rollout' >&2 + return 1 + } + active_profile=S$slots + echo "single_node_capacity=PASS slots=$slots worker_pool=$worker_pool active_instances=1" + return 0 + fi + sleep 2 + done + echo "single-node Worker allocation did not converge: expected=1:$slots actual=${allocation:-missing}" >&2 + return 1 +} + +install_single_node_load_driver() { + [[ $remote_load_drivers_installed == false ]] || return 0 + local binary_path env_file remote_env local_digest remote_digest + binary_path=$(remote_load_binary_path) + remote_env=$(remote_load_env_path ningbo) + env_file=$temporary_root/remote-load-ningbo.env + write_remote_load_env "$env_file" "$AI_GATEWAY_ACCEPTANCE_GATEWAYS" + local_digest=$(shasum -a 256 "$acceptance_load_linux_binary" | awk '{print $1}') + cluster_scp "$acceptance_load_linux_binary" "$CLUSTER_NINGBO_HOST:$binary_path" >/dev/null + cluster_scp "$env_file" "$CLUSTER_NINGBO_HOST:$remote_env" >/dev/null + remote_digest=$(cluster_ssh "$CLUSTER_NINGBO_HOST" \ + "chmod 0755 '$binary_path'; chmod 0600 '$remote_env'; sha256sum '$binary_path'" | awk '{print $1}') + [[ $remote_digest == "$local_digest" ]] + remote_load_drivers_installed=true + echo 'acceptance_load_driver=PASS sites=1 placement=ningbo-host digest_verified=true' +} + +run_single_node_load_profile() { + local profile=$1 + local report_path=$2 + local requests=$3 + install_single_node_load_driver + local binary_path remote_env artifact remote_report status=0 + binary_path=$(remote_load_binary_path) + remote_env=$(remote_load_env_path ningbo) + artifact=$(basename "$report_path") + remote_report=/root/easyai-acceptance-load-"$run_id"-"$artifact" + cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- \ + "$binary_path" "$remote_env" "$profile" "$remote_report" "${artifact%.json}" "$requests" \ + >"$temporary_root/$artifact.stdout" <<'REMOTE' || status=$? +set -euo pipefail +binary_path=$1 +env_file=$2 +profile=$3 +report=$4 +execution_id=$5 +requests=$6 +[[ ! -e $report ]] || unlink "$report" +set -a +# shellcheck source=/dev/null +source "$env_file" +set +a +"$binary_path" -profile "$profile" -report "$report" \ + -shard-index 0 -shard-count 1 -execution-id "$execution_id" \ + -requests "$requests" +REMOTE + cluster_scp "$CLUSTER_NINGBO_HOST:$remote_report" "$report_path" >/dev/null || return 1 + cluster_ssh "$CLUSTER_NINGBO_HOST" "[[ ! -e '$remote_report' ]] || unlink '$remote_report'" >/dev/null || true + chmod 0600 "$report_path" + jq -e --arg runId "$run_id" --arg profile "$profile" \ + '.schemaVersion == "acceptance-load-report/v1" and .runId == $runId and .profile == $profile and .secretSafe == true' \ + "$report_path" >/dev/null || return 1 + (( status == 0 )) && jq -e '.passed == true' "$report_path" >/dev/null +} + +sample_single_node_pressure() { + local output=$1 + local stop_file=$2 + local failure_file=$3 + echo 'timestamp,queued,running,db_connections,db_max_connections,active_instances,allocated_capacity,node_memory_percent,worker_memory_mib,worker_cpu_millicores,restarts' >"$output" + while [[ ! -f $stop_file ]]; do + local state node_memory worker_resources restarts row + state=$(database_query " +SELECT + count(*) FILTER (WHERE status='queued')||','|| + count(*) FILTER (WHERE status='running')||','|| + (SELECT count(*) FROM pg_stat_activity WHERE backend_type='client backend')||','|| + (SELECT setting FROM pg_settings WHERE name='max_connections')||','|| + (SELECT count(*) FROM gateway_worker_instances WHERE status='active' AND heartbeat_at > now()-interval '30 seconds')||','|| + (SELECT COALESCE(max(allocated_capacity),0) FROM gateway_worker_instances WHERE status='active' AND heartbeat_at > now()-interval '30 seconds') +FROM gateway_tasks +WHERE acceptance_run_id='$run_id'::uuid;") || return 1 + node_memory=$(remote_kubectl top node easyai-ningbo --no-headers | + awk '{value=$5; sub(/%$/, "", value); print value}') || return 1 + worker_resources=$(remote_kubectl top pods -n "$namespace" \ + -l 'app.kubernetes.io/name=easyai-worker,easyai.io/site=ningbo' --no-headers | + awk '{ + cpu=$2; memory=$3 + if (cpu ~ /n$/) {sub(/n$/, "", cpu); cpu/=1000000} + else if (cpu ~ /u$/) {sub(/u$/, "", cpu); cpu/=1000} + else if (cpu ~ /m$/) {sub(/m$/, "", cpu)} + else {cpu*=1000} + if (memory ~ /Gi$/) {sub(/Gi$/, "", memory); memory*=1024} + else if (memory ~ /Mi$/) {sub(/Mi$/, "", memory)} + else if (memory ~ /Ki$/) {sub(/Ki$/, "", memory); memory/=1024} + print int(memory+0) "," int(cpu+0) + }') || return 1 + restarts=$(remote_kubectl get pods -n "$namespace" \ + -l 'app.kubernetes.io/part-of=easyai-ai-gateway' -o json | + jq '[.items[].status.containerStatuses[]?.restartCount] | add // 0') || return 1 + row="$(date -u '+%Y-%m-%dT%H:%M:%SZ'),$state,$node_memory,$worker_resources,$restarts" + printf '%s\n' "$row" >>"$output" + if ! awk -F',' '($5<=0)||($4*4>=$5*3)||($6!=1)||($7<1)||($8>=85)||($9>=1536)||($11>0){exit 1}' <<<"$row"; then + printf '%s\n' "$row" >"$failure_file" + return 1 + fi + sleep 1 + done +} + +record_single_node_result() { + local workload=$1 + local slots=$2 + local load_report=$3 + local pressure_report=$4 + local summary=$5 + local peaks + peaks=$(awk -F',' ' + NR>1 { + samples++ + if ($2>queued) queued=$2 + if ($3>running) running=$3 + if ($4>db) db=$4 + if ($7>allocated) allocated=$7 + if ($8>node_memory) node_memory=$8 + if ($9>worker_memory) worker_memory=$9 + if ($10>worker_cpu) worker_cpu=$10 + } + END {printf "%d,%d,%d,%d,%d,%d,%d,%d",samples,queued,running,db,allocated,node_memory,worker_memory,worker_cpu} + ' "$pressure_report") + IFS=',' read -r samples peak_queued peak_running peak_db peak_allocated peak_node_memory peak_worker_memory peak_worker_cpu <<<"$peaks" + jq -n \ + --arg workload "$workload" \ + --argjson slots "$slots" \ + --argjson samples "$samples" \ + --argjson peakQueued "$peak_queued" \ + --argjson peakRunning "$peak_running" \ + --argjson peakDatabaseConnections "$peak_db" \ + --argjson peakAllocatedCapacity "$peak_allocated" \ + --argjson peakNodeMemoryPercent "$peak_node_memory" \ + --argjson peakWorkerMemoryMiB "$peak_worker_memory" \ + --argjson peakWorkerCPUMillicores "$peak_worker_cpu" \ + --slurpfile load "$load_report" \ + '{ + workload:$workload, + configuredSlots:$slots, + samples:$samples, + peakQueued:$peakQueued, + peakRunning:$peakRunning, + peakDatabaseConnections:$peakDatabaseConnections, + peakAllocatedCapacity:$peakAllocatedCapacity, + peakNodeMemoryPercent:$peakNodeMemoryPercent, + peakWorkerMemoryMiB:$peakWorkerMemoryMiB, + peakWorkerCPUMillicores:$peakWorkerCPUMillicores, + load:$load[0] + }' >"$summary" + chmod 0600 "$summary" +} + +run_single_node_profile() { + local workload=$1 + local profile=$2 + local slots=$3 + local requests=$4 + local prefix=$report_root/single-node-$workload-s$slots + local load_report=$prefix-load.json + local pressure_report=$prefix-pressure.csv + local pressure_failure=$prefix-pressure.failure.csv + local pressure_stop=$temporary_root/single-pressure-stop + local sampler_pid status=0 sampler_status=0 + rm -f -- "$pressure_stop" "$pressure_failure" + sample_single_node_pressure "$pressure_report" "$pressure_stop" "$pressure_failure" & + sampler_pid=$! + active_pressure_pid=$sampler_pid + run_with_pressure_monitor "$sampler_pid" \ + run_single_node_load_profile "$profile" "$load_report" "$requests" || status=$? + touch "$pressure_stop" + wait "$sampler_pid" || sampler_status=$? + active_pressure_pid= + (( status == 0 && sampler_status == 0 )) || { + [[ -s $pressure_failure ]] && cp "$pressure_failure" "$report_root/last-pressure-failure.csv" + return 1 + } + record_single_node_result "$workload" "$slots" "$load_report" "$pressure_report" \ + "$prefix-summary.json" +} + +finish_single_node_diagnostic() { + local diagnostic_status=$1 + local reason=$2 + local emulator_report=$report_root/single-node-emulator-report.json + local callback_report=$report_root/single-node-callback-report.json + local finish_response=$temporary_root/single-node-finish.json + if [[ $single_node_emulator_installed == true ]]; then + cluster_ssh "$CLUSTER_NINGBO_HOST" \ + 'curl -fsS --max-time 10 http://127.0.0.1:18090/report' >"$emulator_report" || true + else + remote_kubectl exec -n "$namespace" deployment/easyai-acceptance-emulator -- \ + wget -qO- http://127.0.0.1:8090/report >"$emulator_report" || true + fi + remote_kubectl exec -n "$namespace" deployment/easyai-acceptance-callback-collector -- \ + wget -qO- http://127.0.0.1:8091/report >"$callback_report" || true + local task_state + task_state=$(database_query " +SELECT count(*)||':'||count(*) FILTER (WHERE status='succeeded')||':'||count(*) FILTER (WHERE status IN ('queued','running')) +FROM gateway_tasks WHERE acceptance_run_id='$run_id'::uuid;") + local results_json=$temporary_root/single-node-results.json + local -a result_files=() + shopt -s nullglob + result_files=("$report_root"/single-node-*-summary.json) + shopt -u nullglob + if (( ${#result_files[@]} > 0 )); then + jq -s '.' "${result_files[@]}" >"$results_json" + else + printf '[]\n' >"$results_json" + fi + jq -n \ + --arg runId "$run_id" \ + --arg releaseSha "$release_sha" \ + --arg status "$diagnostic_status" \ + --arg reason "$reason" \ + --arg taskState "$task_state" \ + --argjson stableSlots "$single_node_stable_slots" \ + --slurpfile results "$results_json" \ + '{ + schemaVersion:"acceptance-single-node-diagnostic/v1", + runId:$runId, + releaseSha:$releaseSha, + site:"ningbo", + mode:"acceptance-emulator-only", + realUpstreamRequests:0, + certification:false, + promotable:false, + status:$status, + reason:$reason, + taskState:$taskState, + stableVideoSlots:$stableSlots, + results:$results[0] + }' >"$report_root/single-node-summary.json" + chmod 0600 "$report_root/single-node-summary.json" + local finish_body + finish_body=$(jq -cn \ + --arg reason "single-node diagnostic completed; real upstream and certification intentionally skipped" \ + --arg status "$diagnostic_status" \ + --argjson stableSlots "$single_node_stable_slots" \ + '{passed:false,failureReason:$reason,report:{diagnosticStatus:$status,site:"ningbo",realUpstreamRequests:0,certification:false,stableVideoSlots:$stableSlots}}') + admin_request POST "/api/admin/system/acceptance/runs/$run_id/finish" "$finish_body" "$finish_response" + failure_recorded=true + cleanup_single_node_emulator + remote_kubectl delete deployment easyai-acceptance-emulator easyai-acceptance-callback-collector \ + -n "$namespace" --ignore-not-found >/dev/null || true + remote_kubectl delete service easyai-acceptance-emulator easyai-acceptance-callback-collector \ + -n "$namespace" --ignore-not-found >/dev/null || true +} + +run_single_node_acceptance() { + wait_for_existing_tasks_to_drain + single_node_stable_slots=$AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_BASELINE_SLOTS + apply_single_node_capacity "$AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_SLOTS" + if ! run_single_node_profile gemini-multi-image gemini-multi-image \ + "$AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_SLOTS" \ + "$AI_GATEWAY_ACCEPTANCE_SINGLE_NODE_GEMINI_REQUESTS"; then + finish_single_node_diagnostic failed 'GEMINI multi-image diagnostic failed or crossed a resource gate' + return 1 + fi + local slots requests + for slots in "${single_node_video_slots[@]}"; do + apply_single_node_capacity "$slots" || break + requests=$((slots * 3)) + (( requests < 32 )) && requests=32 + if ! run_single_node_profile image-video video-throughput "$slots" "$requests"; then + break + fi + single_node_stable_slots=$slots + done + apply_single_node_capacity "$single_node_stable_slots" + finish_single_node_diagnostic completed 'single-node emulator-only capacity ladder completed' + echo "single_node_acceptance=PASS_DIAGNOSTIC run_id=$run_id stable_video_slots=$single_node_stable_slots traffic_mode=validation certification=false report=$report_root/single-node-summary.json" +} + run_capacity_round() { local profile=$1 local repetition=$2 @@ -3651,11 +4178,17 @@ verify_release_cas false bootstrap_acceptance_primary_identity ensure_acceptance_user_group ensure_acceptance_identity_shards -ensure_acceptance_real_images +if [[ $single_node_acceptance != true ]]; then + ensure_acceptance_real_images +fi select_acceptance_models ensure_acceptance_model_access deploy_protocol_emulator create_and_activate_run +if [[ $single_node_acceptance == true ]]; then + run_single_node_acceptance + exit $? +fi wait_for_existing_tasks_to_drain snapshot_pre_acceptance_capacity if ! apply_capacity_profile P24; then