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

6K 多参考图在 P24 下会并发进入高内存缩放并击穿 2 GiB Worker。新增独立可配置的图片归一化并发,生产默认 2,只约束解码、缩放和重编码阶段,不占用视频上游等待槽;缩放器改为内存稳定的近似双线性实现。\n\n验证:6K 到 6000x2400 转换、信号量串行化、Go 全量测试、gofmt、Kubernetes 客户端 dry-run。
This commit is contained in:
2026-07-31 20:54:11 +08:00
parent f190af00be
commit 29533537ec
7 changed files with 127 additions and 9 deletions
@@ -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, "/")) {
@@ -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)
+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),