From 4fd24d2f80bf8f9839a3c47cc8594b1a127ea794 Mon Sep 17 00:00:00 2001
From: tzwang <wtz1003@sina.com>
Date: Sat, 22 Jun 2024 21:16:51 +0800
Subject: [PATCH] fix aitask bugs

Former-commit-id: 1fbc69fee0c10fd34334338ef379d2b5843e6fad
---
 api/internal/logic/inference/imageinferencelogic.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/api/internal/logic/inference/imageinferencelogic.go b/api/internal/logic/inference/imageinferencelogic.go
index b80de115..684d6059 100644
--- a/api/internal/logic/inference/imageinferencelogic.go
+++ b/api/internal/logic/inference/imageinferencelogic.go
@@ -269,6 +269,7 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 
 	var result_ch = make(chan *types.ImageResult, len(ts))
 	var results []*types.ImageResult
+	limit := make(chan bool, 7)
 
 	var imageNumIdx int32 = 0
 	var imageNumIdxEnd int32 = 0
@@ -284,7 +285,7 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 		imageNumIdx = imageNumIdx + c.imageNum
 
 		wg.Add(len(new_images))
-		go sendInferReq(new_images, c, &wg, result_ch)
+		go sendInferReq(new_images, c, &wg, result_ch, limit)
 	}
 	wg.Wait()
 	close(result_ch)
@@ -375,8 +376,9 @@ func sendInferReq(images []struct {
 	clusterId   string
 	clusterName string
 	imageNum    int32
-}, wg *sync.WaitGroup, ch chan<- *types.ImageResult) {
+}, wg *sync.WaitGroup, ch chan<- *types.ImageResult, limit chan bool) {
 	for _, image := range images {
+		limit <- true
 		go func(t struct {
 			imageResult *types.ImageResult
 			file        multipart.File
@@ -395,6 +397,7 @@ func sendInferReq(images []struct {
 					t.imageResult.Card = c.urls[0].Card
 					ch <- t.imageResult
 					wg.Done()
+					<-limit
 					return
 				}
 				t.imageResult.ImageResult = r
@@ -404,6 +407,7 @@ func sendInferReq(images []struct {
 
 				ch <- t.imageResult
 				wg.Done()
+				<-limit
 				return
 			} else {
 				idx := rand.Intn(len(c.urls))
@@ -415,6 +419,7 @@ func sendInferReq(images []struct {
 					t.imageResult.Card = c.urls[idx].Card
 					ch <- t.imageResult
 					wg.Done()
+					<-limit
 					return
 				}
 				t.imageResult.ImageResult = r
@@ -424,9 +429,11 @@ func sendInferReq(images []struct {
 
 				ch <- t.imageResult
 				wg.Done()
+				<-limit
 				return
 			}
 		}(image, cluster)
+		<-limit
 	}
 }