BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(worker): 限制大图归一化峰值内存

Browse Source 
6K 多参考图在 P24 下会并发进入高内存缩放并击穿 2 GiB Worker。新增独立可配置的图片归一化并发，生产默认 2，只约束解码、缩放和重编码阶段，不占用视频上游等待槽；缩放器改为内存稳定的近似双线性实现。\n\n验证：6K 到 6000x2400 转换、信号量串行化、Go 全量测试、gofmt、Kubernetes 客户端 dry-run。
This commit is contained in:
wangbo
2026-07-31 20:54:11 +08:00
parent f190af00be
commit 29533537ec
7 changed files with 127 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/config/config.go
+16 -8
View File
@@ -67,6 +67,7 @@ type Config struct {
DatabaseLockTimeoutSeconds int
MediaRequestConcurrency int
MediaMaterializationConcurrency int
MediaImageNormalizationConcurrency int
MediaOSSDirectEnabled bool
MediaOSSEndpoint string
MediaOSSBucket string
@@ -167,14 +168,18 @@ func Load() Config {
MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16),
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
MediaImageNormalizationConcurrency: envIntValidated(
"AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY",
2,
),
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated(
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT",
envOptionalIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
@@ -252,6 +257,9 @@ func (c Config) Validate() error {
if c.MediaRequestConcurrency != 0 && (c.MediaRequestConcurrency < 1 || c.MediaRequestConcurrency > 1024) {
return errors.New("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY must be between 1 and 1024")
}
if c.MediaImageNormalizationConcurrency != 0 && (c.MediaImageNormalizationConcurrency < 1 || c.MediaImageNormalizationConcurrency > 64) {
return errors.New("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY must be between 1 and 64")
}
if c.MediaOSSDirectEnabled {
if strings.TrimSpace(c.MediaOSSAccessKeyID) == "" || strings.TrimSpace(c.MediaOSSAccessKeySecret) == "" {
return errors.New("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID and AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET are required when direct media OSS is enabled")
apps/api/internal/config/config_test.go
+9
View File
@@ -270,6 +270,15 @@ func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
if got := Load().MediaRequestConcurrency; got != 96 {
t.Fatalf("media request concurrency = %d, want 96", got)
}
cfg = Load()
cfg.MediaImageNormalizationConcurrency = 65
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_IMAGE_NORMALIZATION_CONCURRENCY") {
t.Fatalf("Validate() error = %v, want invalid image normalization concurrency", err)
}
t.Setenv("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY", "3")
if got := Load().MediaImageNormalizationConcurrency; got != 3 {
t.Fatalf("image normalization concurrency = %d, want 3", got)
}
}
func TestValidateTaskHistorySettings(t *testing.T) {
apps/api/internal/runner/input_image_normalizer.go
+22 -1
View File
@@ -14,6 +14,7 @@ import (
"net/http"
"os"
"strings"
"sync"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
@@ -195,6 +196,11 @@ func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source str
return source, false, original, original, nil
}
release, err := s.acquireImageNormalizationSlot(ctx)
if err != nil {
return "", false, original, image.Point{}, err
}
defer release()
decoded, _, err := image.Decode(bytes.NewReader(payload))
if err != nil {
return "", false, original, image.Point{}, fmt.Errorf("图片无法解码")
@@ -209,7 +215,7 @@ func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source str
contentHeight := max(1, int(math.Round(float64(sourceBounds.Dy())*scale)))
left := (target.X - contentWidth) / 2
top := (target.Y - contentHeight) / 2
draw.CatmullRom.Scale(
draw.ApproxBiLinear.Scale(
canvas,
image.Rect(left, top, left+contentWidth, top+contentHeight),
decoded,
@@ -228,6 +234,21 @@ func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source str
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(output.Bytes()), true, original, target, nil
}
func (s *Service) acquireImageNormalizationSlot(ctx context.Context) (func(), error) {
if s.imageNormalizeSlots == nil {
return func() {}, nil
}
select {
case s.imageNormalizeSlots <- struct{}{}:
var once sync.Once
return func() {
once.Do(func() { <-s.imageNormalizeSlots })
}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
func (s *Service) readVideoInputImageBytes(ctx context.Context, source string) ([]byte, error) {
lower := strings.ToLower(source)
if strings.HasPrefix(lower, "data:") || (!strings.Contains(source, "://") && !strings.HasPrefix(source, "/")) {
apps/api/internal/runner/input_image_normalizer_test.go
+62
View File
@@ -10,6 +10,7 @@ import (
"image/png"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -61,6 +62,15 @@ func decodedImageSize(t *testing.T, source string) image.Point {
return image.Pt(config.Width, config.Height)
}
func jpegImageDataURL(t *testing.T, width int, height int) string {
t.Helper()
var payload bytes.Buffer
if err := jpeg.Encode(&payload, image.NewRGBA(image.Rect(0, 0, width, height)), &jpeg.Options{Quality: 90}); err != nil {
t.Fatal(err)
}
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(payload.Bytes())
}
func TestNormalizeVideoInputImagesResizesOversizedSeedanceImage(t *testing.T) {
service := &Service{}
body := map[string]any{
@@ -85,6 +95,58 @@ func TestNormalizeVideoInputImagesResizesOversizedSeedanceImage(t *testing.T) {
}
}
func TestNormalizeVideoInputImagesBoundsSixKReferenceMemoryWork(t *testing.T) {
service := &Service{imageNormalizeSlots: make(chan struct{}, 1)}
candidate := store.RuntimeModelCandidate{
ModelType: "omni_video",
Capabilities: map[string]any{
"omni_video": map[string]any{
"input_image_resolution_range": map[string]any{
"min": map[string]any{"long_edge": float64(300), "short_edge": float64(300)},
"max": map[string]any{"long_edge": float64(6000), "short_edge": float64(6000)},
},
"input_image_aspect_ratio_range": []any{float64(0.4), float64(2.5)},
},
},
}
body := map[string]any{"first_frame": jpegImageDataURL(t, 6144, 2160)}
normalized, err := service.normalizeVideoInputImages(context.Background(), body, candidate)
if err != nil {
t.Fatal(err)
}
if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(6000, 2400) {
t.Fatalf("unexpected converted size: %v", got)
}
}
func TestImageNormalizationSlotSerializesHeavyConversions(t *testing.T) {
service := &Service{imageNormalizeSlots: make(chan struct{}, 1)}
releaseFirst, err := service.acquireImageNormalizationSlot(context.Background())
if err != nil {
t.Fatal(err)
}
acquired := make(chan func(), 1)
go func() {
release, acquireErr := service.acquireImageNormalizationSlot(context.Background())
if acquireErr == nil {
acquired <- release
}
}()
select {
case release := <-acquired:
release()
t.Fatal("second normalization entered while the only slot was held")
case <-time.After(50 * time.Millisecond):
}
releaseFirst()
select {
case release := <-acquired:
release()
case <-time.After(time.Second):
t.Fatal("second normalization did not acquire the released slot")
}
}
func TestNormalizeVideoInputImagesPadsAspectRatioViolation(t *testing.T) {
service := &Service{}
source := pngImageDataURL(t, 1530, 500)
apps/api/internal/runner/service.go
+5
View File
@@ -47,6 +47,7 @@ type Service struct {
admissionListener sync.Once
asyncClientFactory func(int) (asyncExecutionClient, error)
mediaResultSlots chan struct{}
imageNormalizeSlots chan struct{}
directOSS *directOSSUploader
taskCompletionOnce sync.Once
taskCompletionMu sync.Mutex
@@ -130,6 +131,9 @@ func NewWithStores(
if cfg.MediaMaterializationConcurrency == 0 {
cfg.MediaMaterializationConcurrency = 8
}
if cfg.MediaImageNormalizationConcurrency == 0 {
cfg.MediaImageNormalizationConcurrency = 2
}
if cfg.TaskProgressCallbackTimeoutMS == 0 {
cfg.TaskProgressCallbackTimeoutMS = 5000
}
@@ -184,6 +188,7 @@ func NewWithStores(
asyncAdmissionWake: make(chan struct{}, 1),
admissionTaskWaiters: map[string]*admissionTaskWaiter{},
mediaResultSlots: make(chan struct{}, cfg.MediaMaterializationConcurrency),
imageNormalizeSlots: make(chan struct{}, cfg.MediaImageNormalizationConcurrency),
directOSS: newDirectOSSUploader(cfg),
taskCompletionWaiters: map[string]map[chan struct{}]struct{}{},
taskCompletionPollWake: make(chan struct{}, 1),
deploy/kubernetes/production/application.yaml
+8
View File
@@ -66,6 +66,8 @@ spec:
value: "16"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8"
- name: AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY
value: "2"
- name: AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID
valueFrom:
secretKeyRef:
@@ -335,6 +337,8 @@ spec:
value: "16"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8"
- name: AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY
value: "2"
- name: AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID
valueFrom:
secretKeyRef:
@@ -492,6 +496,8 @@ spec:
value: "24"
- name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY
value: "24"
- name: AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY
value: "2"
- name: AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID
valueFrom:
secretKeyRef:
@@ -649,6 +655,8 @@ spec:
value: "24"
- name: AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY
value: "24"
- name: AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY
value: "2"
- name: AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID
valueFrom:
secretKeyRef:
docs/operations/production-acceptance.md
+5
View File
@@ -77,6 +77,7 @@ AI_GATEWAY_DATABASE_MIN_IDLE_CONNS
AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS
AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY
AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY
AI_GATEWAY_WORKER_REPLICAS_NINGBO
AI_GATEWAY_WORKER_REPLICAS_HONGKONG
```
@@ -91,6 +92,10 @@ AI_GATEWAY_WORKER_REPLICAS_HONGKONG
| P28 | 28 | 36 | 28 | 56 |
| P32 | 32 | 40 | 32 | 64 |
`AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY` 单独约束 6K 等超限参考图的解码、缩放和
重新编码，避免执行槽升档时把大图峰值内存同比放大。2 GiB Worker 的起始值为 `2`；它不占用
视频提交后的上游等待槽，只有实测单 Pod RSS 和节点预算允许时才单独升档。
API Pod 使用独立发布配置 `AI_GATEWAY_API_DATABASE_MAX_CONNS`，生产同构验收默认使用 `64`；
也可通过 `AI_GATEWAY_ACCEPTANCE_API_DATABASE_MAX_CONNS` 覆盖本次验收值，无需修改 Go 代码。
所有 API/Worker Pool 的 `AI_GATEWAY_DATABASE_MIN_IDLE_CONNS` 固定为 `4`，仅预热必要连接，