diff --git a/models/resource_specification.go b/models/resource_specification.go index 2f815818b..809a3496a 100644 --- a/models/resource_specification.go +++ b/models/resource_specification.go @@ -3,6 +3,7 @@ package models import ( "code.gitea.io/gitea/modules/timeutil" "fmt" + "strings" "xorm.io/builder" ) @@ -197,12 +198,104 @@ type Specification struct { AiCenterName string IsExclusive bool ExclusiveOrg string + //specs that have the same sourceSpecId, computeResource and cluster as current spec + RelatedSpecs []*Specification } func (Specification) TableName() string { return "resource_specification" } +func (s *Specification) loadRelatedSpecs() { + if s.RelatedSpecs != nil { + return + } + defaultSpecs := make([]*Specification, 0) + if s.SourceSpecId == "" { + s.RelatedSpecs = defaultSpecs + return + } + r, err := FindSpecs(FindSpecsOptions{ + ComputeResource: s.ComputeResource, + Cluster: s.Cluster, + SourceSpecId: s.SourceSpecId, + RequestAll: true, + SpecStatus: SpecOnShelf, + }) + if err != nil { + s.RelatedSpecs = defaultSpecs + return + } + s.RelatedSpecs = r +} +func (s *Specification) GetAvailableCenterIds(userIds ...int64) []string { + s.loadRelatedSpecs() + + if len(s.RelatedSpecs) == 0 { + return make([]string, 0) + } + + var uId int64 + if len(userIds) > 0 { + uId = userIds[0] + } + //filter exclusive specs + specs := FilterExclusiveSpecs(s.RelatedSpecs, uId) + + centerIds := make([]string, len(specs)) + for i, v := range specs { + centerIds[i] = v.AiCenterCode + } + return centerIds +} + +func FilterExclusiveSpecs(r []*Specification, userId int64) []*Specification { + if userId == 0 { + return r + } + specs := make([]*Specification, 0, len(r)) + specMap := make(map[int64]string, 0) + for i := 0; i < len(r); i++ { + spec := r[i] + if _, has := specMap[spec.ID]; has { + continue + } + if !spec.IsExclusive { + specs = append(specs, spec) + specMap[spec.ID] = "" + continue + } + orgs := strings.Split(spec.ExclusiveOrg, ";") + for _, org := range orgs { + isMember, _ := IsOrganizationMemberByOrgName(org, userId) + if isMember { + specs = append(specs, spec) + specMap[spec.ID] = "" + break + } + } + } + return specs +} + +func DistinctSpecs(r []*Specification) []*Specification { + specs := make([]*Specification, 0, len(r)) + sourceSpecIdMap := make(map[string]string, 0) + for i := 0; i < len(r); i++ { + spec := r[i] + if spec.SourceSpecId == "" { + specs = append(specs, spec) + continue + } + if _, has := sourceSpecIdMap[spec.SourceSpecId]; has { + continue + } + specs = append(specs, spec) + sourceSpecIdMap[spec.SourceSpecId] = "" + } + return specs +} + func InsertResourceSpecification(r ResourceSpecification) (int64, error) { return x.Insert(&r) } diff --git a/modules/grampus/grampus.go b/modules/grampus/grampus.go index b6f62560a..6b2ea6288 100755 --- a/modules/grampus/grampus.go +++ b/modules/grampus/grampus.go @@ -105,8 +105,6 @@ func getDatasetGrampus(datasetInfos map[string]models.DatasetInfo) []models.Gram func GenerateTrainJob(ctx *context.Context, req *GenerateTrainJobReq) (jobId string, err error) { createTime := timeutil.TimeStampNow() - centerID, centerName := getCentersParamter(ctx, req) - var datasetGrampus, modelGrampus []models.GrampusDataset var codeGrampus models.GrampusDataset if ProcessorTypeNPU == req.ProcessType { @@ -138,8 +136,7 @@ func GenerateTrainJob(ctx *context.Context, req *GenerateTrainJobReq) (jobId str ResourceSpecId: req.Spec.SourceSpecId, ImageId: req.ImageId, ImageUrl: req.ImageUrl, - CenterID: centerID, - CenterName: centerName, + CenterID: req.Spec.GetAvailableCenterIds(ctx.User.ID), ReplicaNum: 1, Datasets: datasetGrampus, Models: modelGrampus, diff --git a/services/cloudbrain/resource/resource_specification.go b/services/cloudbrain/resource/resource_specification.go index 8f4182d87..5070d7c1e 100644 --- a/services/cloudbrain/resource/resource_specification.go +++ b/services/cloudbrain/resource/resource_specification.go @@ -246,10 +246,10 @@ func FindAvailableSpecs(userId int64, opts models.FindSpecsOptions) ([]*models.S return nil, err } //filter exclusive specs - specs := filterExclusiveSpecs(r, userId) + specs := models.FilterExclusiveSpecs(r, userId) //distinct by sourceSpecId - specs = distinctSpecs(specs) + specs = models.DistinctSpecs(specs) return specs, err } @@ -265,50 +265,6 @@ func FindAvailableSpecs4Show(userId int64, opts models.FindSpecsOptions) ([]*api return result, nil } -func filterExclusiveSpecs(r []*models.Specification, userId int64) []*models.Specification { - specs := make([]*models.Specification, 0, len(r)) - specMap := make(map[int64]string, 0) - for i := 0; i < len(r); i++ { - spec := r[i] - if _, has := specMap[spec.ID]; has { - continue - } - if !spec.IsExclusive { - specs = append(specs, spec) - specMap[spec.ID] = "" - continue - } - orgs := strings.Split(spec.ExclusiveOrg, ";") - for _, org := range orgs { - isMember, _ := models.IsOrganizationMemberByOrgName(org, userId) - if isMember { - specs = append(specs, spec) - specMap[spec.ID] = "" - break - } - } - } - return specs -} - -func distinctSpecs(r []*models.Specification) []*models.Specification { - specs := make([]*models.Specification, 0, len(r)) - sourceSpecIdMap := make(map[string]string, 0) - for i := 0; i < len(r); i++ { - spec := r[i] - if spec.SourceSpecId == "" { - specs = append(specs, spec) - continue - } - if _, has := sourceSpecIdMap[spec.SourceSpecId]; has { - continue - } - specs = append(specs, spec) - sourceSpecIdMap[spec.SourceSpecId] = "" - } - return specs -} - func GetAndCheckSpec(userId int64, specId int64, opts models.FindSpecsOptions) (*models.Specification, error) { if specId == 0 { return nil, nil