fix(worker): 限制大图归一化峰值内存
6K 多参考图在 P24 下会并发进入高内存缩放并击穿 2 GiB Worker。新增独立可配置的图片归一化并发,生产默认 2,只约束解码、缩放和重编码阶段,不占用视频上游等待槽;缩放器改为内存稳定的近似双线性实现。\n\n验证:6K 到 6000x2400 转换、信号量串行化、Go 全量测试、gofmt、Kubernetes 客户端 dry-run。
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user