diff --git a/internal/logic/schedule/schedulecreatetasklogic.go b/internal/logic/schedule/schedulecreatetasklogic.go index 80566fc0..077511b5 100644 --- a/internal/logic/schedule/schedulecreatetasklogic.go +++ b/internal/logic/schedule/schedulecreatetasklogic.go @@ -103,6 +103,7 @@ func generateFilteredDataDistributes(clusters []*strategy.AssignedCluster, distr code := &types.CodeDistribute{} code.DataName = codeDistribute.DataName code.PackageID = codeDistribute.PackageID + code.Output = codeDistribute.Output clusterScheduledList := make([]*types.ClusterScheduled, 0) if len(codeDistribute.Clusters) != 0 { diff --git a/internal/logic/schedule/scheduleruntasklogic.go b/internal/logic/schedule/scheduleruntasklogic.go index 3d41b3e7..9c844b6c 100644 --- a/internal/logic/schedule/scheduleruntasklogic.go +++ b/internal/logic/schedule/scheduleruntasklogic.go @@ -103,6 +103,7 @@ func (l *ScheduleRunTaskLogic) SaveResult(task *models.Task, results []*schedule opt.ComputeCard = strings.ToUpper(r.Card) opt.Replica = r.Replica + opt.Output = r.Output adapterName, err := l.svcCtx.Scheduler.AiStorages.GetAdapterNameById(r.AdapterId) if err != nil { @@ -274,10 +275,14 @@ func updateClustersByScheduledDatas(taskId int64, clustersWithDataDistributes *C } } - if cluster.CodeId == "" { - for _, distribute := range clustersWithDataDistributes.DataDistributes.Code { - for _, c := range distribute.Clusters { - if cluster.ClusterId == c.ClusterID { + //if cluster.CodeId == "" { + for _, distribute := range clustersWithDataDistributes.DataDistributes.Code { + for _, c := range distribute.Clusters { + if cluster.ClusterId == c.ClusterID { + + cluster.Output = distribute.Output + + if cluster.CodeId == "" { if c.JsonData == "" { continue } @@ -291,6 +296,7 @@ func updateClustersByScheduledDatas(taskId int64, clustersWithDataDistributes *C } cluster.CodeId = jsonData.Id } + } } } diff --git a/internal/scheduler/database/aiStorage.go b/internal/scheduler/database/aiStorage.go index 787c3ffb..85e2c595 100644 --- a/internal/scheduler/database/aiStorage.go +++ b/internal/scheduler/database/aiStorage.go @@ -196,6 +196,7 @@ func (s *AiStorage) SaveAiTask(taskId int64, opt option.Option, adapterName stri Strategy: aiOpt.StrategyName, Status: status, Msg: msg, + Output: aiOpt.Output, Card: aiOpt.ComputeCard, StartTime: time.Now().Format(time.RFC3339), CommitTime: time.Now(), diff --git a/internal/scheduler/schedulers/aiScheduler.go b/internal/scheduler/schedulers/aiScheduler.go index 2589790b..d701b186 100644 --- a/internal/scheduler/schedulers/aiScheduler.go +++ b/internal/scheduler/schedulers/aiScheduler.go @@ -57,6 +57,7 @@ type AiResult struct { Replica int32 Card string Msg string + Output string } func NewAiScheduler(ctx context.Context, val string, scheduler *scheduler.Scheduler, option *option.AiOption) (*AiScheduler, error) { @@ -206,6 +207,7 @@ func (as *AiScheduler) AssignTask(clusters []*strategy.AssignedCluster, mode int result.ClusterId = c.ClusterId result.Strategy = as.option.StrategyName result.Card = opt.ComputeCard + result.Output = opt.Output ch <- result wg.Done() @@ -319,9 +321,7 @@ func (as *AiScheduler) handleErrors(errs []interface{}, clusters []*strategy.Ass } //report status - if mode == executor.SUBMIT_MODE_STORAGE_SCHEDULE { - _ = jcs.StatusReport(as.AiService.Conf.JcsMiddleware.JobStatusReportUrl, report) - } + _ = jcs.StatusReport(as.AiService.Conf.JcsMiddleware.JobStatusReportUrl, report) logx.Errorf(errors.New(errmsg).Error()) return errors.New(errmsg) @@ -343,6 +343,7 @@ func updateAiOptionByMode(cluster *strategy.AssignedCluster, opt *option.AiOptio opt.ResourcesRequired = cluster.ResourcesRequired + opt.Output = cluster.Output default: } diff --git a/internal/scheduler/schedulers/option/aiOption.go b/internal/scheduler/schedulers/option/aiOption.go index 7aa08911..3dcd35d7 100644 --- a/internal/scheduler/schedulers/option/aiOption.go +++ b/internal/scheduler/schedulers/option/aiOption.go @@ -34,6 +34,8 @@ type AiOption struct { Image string Model interface{} + Output string + ResourcesRequired []map[string]interface{} } diff --git a/internal/scheduler/strategy/strategy.go b/internal/scheduler/strategy/strategy.go index 7df66927..c3754eba 100644 --- a/internal/scheduler/strategy/strategy.go +++ b/internal/scheduler/strategy/strategy.go @@ -31,6 +31,7 @@ type AssignedCluster struct { ImageId string CodeId string ModelId string + Output string } func GetStrategyNames() []string { diff --git a/pkg/models/taskaimodel_gen.go b/pkg/models/taskaimodel_gen.go index 5c58decc..4d4f72a9 100644 --- a/pkg/models/taskaimodel_gen.go +++ b/pkg/models/taskaimodel_gen.go @@ -55,6 +55,7 @@ type ( Card string `db:"card"` InferUrl string `db:"infer_url"` ModelName string `db:"model_name"` + Output string `db:"output"` } )