From 0f271360d2425b69198c8d254614eef82c524c71 Mon Sep 17 00:00:00 2001
From: tzwang <wtz1003@sina.com>
Date: Fri, 21 Jun 2024 20:10:51 +0800
Subject: [PATCH] updated imageinfer api

Former-commit-id: 2b0e2a56398b2c10fefb359b63d0b6f1fdfa5582
---
 .../logic/inference/imageinferencelogic.go    | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/api/internal/logic/inference/imageinferencelogic.go b/api/internal/logic/inference/imageinferencelogic.go
index 2a06f78d..30567d8b 100644
--- a/api/internal/logic/inference/imageinferencelogic.go
+++ b/api/internal/logic/inference/imageinferencelogic.go
@@ -217,6 +217,20 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 		return nil, tx.Error
 
 	}
+
+	//no cluster available
+	if len(cs) == 0 {
+		for _, t := range aiTaskList {
+			t.Status = constants.Failed
+			err := svcCtx.Scheduler.AiStorages.UpdateAiTask(t)
+			if err != nil {
+				logx.Errorf(tx.Error.Error())
+			}
+		}
+		svcCtx.Scheduler.AiStorages.AddNoticeInfo(opt.AdapterId, adapterName, "", "", opt.TaskName, "failed", "任务失败")
+		return nil, errors.New("image infer task failed")
+	}
+
 	//change cluster status
 	if len(clusters) != len(cs) {
 		var acs []*strategy.AssignedCluster
@@ -275,6 +289,7 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 	})
 
 	// update succeeded cluster status
+	var successStatusCount int
 	for _, c := range cs {
 		for _, t := range aiTaskList {
 			if c.clusterId == strconv.Itoa(int(t.ClusterId)) {
@@ -283,9 +298,19 @@ func infer(opt *option.InferOption, clusters []*strategy.AssignedCluster, ts []s
 				if err != nil {
 					logx.Errorf(tx.Error.Error())
 				}
+				successStatusCount++
+			} else {
+				continue
 			}
 		}
 	}
+
+	if len(cs) == successStatusCount {
+		svcCtx.Scheduler.AiStorages.AddNoticeInfo(opt.AdapterId, adapterName, "", "", opt.TaskName, "failed", "任务失败")
+	} else {
+		svcCtx.Scheduler.AiStorages.AddNoticeInfo(opt.AdapterId, adapterName, "", "", opt.TaskName, "completed", "任务完成")
+	}
+
 	return results, nil
 }