From 2d8f18129b3d6d49d20a4c258da54e8e5ef6b171 Mon Sep 17 00:00:00 2001 From: chenyifan01 Date: Mon, 5 Sep 2022 18:06:30 +0800 Subject: [PATCH] #2701 update --- models/cloudbrain.go | 45 +++++++++++++++++++++++++++++++ modules/modelarts_cd/modelarts.go | 7 ++--- routers/admin/cloudbrains.go | 4 +-- routers/repo/modelarts.go | 25 ++++++++--------- routers/user/home.go | 4 +-- 5 files changed, 65 insertions(+), 20 deletions(-) diff --git a/models/cloudbrain.go b/models/cloudbrain.go index 8072625c6..6154ced56 100755 --- a/models/cloudbrain.go +++ b/models/cloudbrain.go @@ -121,6 +121,7 @@ const ( //AI center AICenterOfCloudBrainOne = "OpenIOne" AICenterOfCloudBrainTwo = "OpenITwo" + AICenterOfChengdu = "OpenIChengdu" //ComputeResource GPU = "GPU" @@ -2421,3 +2422,47 @@ func GetCloudbrainCountByJobName(jobName, jobType string, typeCloudbrain int) (i count, err := x.Where("job_name = ? and job_type= ? and type = ?", jobName, jobType, typeCloudbrain).Count(new(Cloudbrain)) return int(count), err } + +func LoadSpecs(tasks []*Cloudbrain) error { + cloudbrainIds := make([]int64, len(tasks)) + for i, v := range tasks { + cloudbrainIds[i] = v.ID + } + specs := make([]*CloudbrainSpec, 0) + err := x.In("cloudbrain_id", cloudbrainIds).Find(&specs) + if err != nil { + return err + } + specMap := make(map[int64]*CloudbrainSpec) + for _, v := range specs { + specMap[v.SpecId] = v + } + for _, v := range tasks { + if specMap[v.ID] != nil { + v.Spec = specMap[v.ID].ConvertToSpecification() + } + } + return nil +} + +func LoadSpecs4CloudbrainInfo(tasks []*CloudbrainInfo) error { + cloudbrainIds := make([]int64, len(tasks)) + for i, v := range tasks { + cloudbrainIds[i] = v.Cloudbrain.ID + } + specs := make([]*CloudbrainSpec, 0) + err := x.In("cloudbrain_id", cloudbrainIds).Find(&specs) + if err != nil { + return err + } + specMap := make(map[int64]*CloudbrainSpec) + for _, v := range specs { + specMap[v.SpecId] = v + } + for _, v := range tasks { + if specMap[v.Cloudbrain.ID] != nil { + v.Cloudbrain.Spec = specMap[v.Cloudbrain.ID].ConvertToSpecification() + } + } + return nil +} diff --git a/modules/modelarts_cd/modelarts.go b/modules/modelarts_cd/modelarts.go index 25324dbd5..330b048ca 100755 --- a/modules/modelarts_cd/modelarts.go +++ b/modules/modelarts_cd/modelarts.go @@ -88,7 +88,7 @@ type Parameters struct { } `json:"parameter"` } -func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, description, flavor, imageId string) error { +func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, description, imageId string, spec *models.Specification) error { imageName, err := GetNotebookImageName(imageId) if err != nil { log.Error("GetNotebookImageName failed: %v", err.Error()) @@ -98,7 +98,7 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr jobResult, err := createNotebook(models.CreateNotebookWithoutPoolParams{ JobName: jobName, Description: description, - Flavor: flavor, + Flavor: spec.SourceSpecId, Duration: autoStopDurationMs, ImageID: imageId, Feature: models.NotebookFeature, @@ -134,7 +134,7 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr RepoID: ctx.Repo.Repository.ID, JobID: jobResult.ID, JobName: jobName, - FlavorCode: flavor, + FlavorCode: spec.SourceSpecId, DisplayJobName: displayJobName, JobType: string(models.JobTypeDebug), Type: models.TypeCDCenter, @@ -144,6 +144,7 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr Description: description, CreatedUnix: createTime, UpdatedUnix: createTime, + Spec: spec, } err = models.CreateCloudbrain(task) diff --git a/routers/admin/cloudbrains.go b/routers/admin/cloudbrains.go index ec0034f4f..fcb878627 100755 --- a/routers/admin/cloudbrains.go +++ b/routers/admin/cloudbrains.go @@ -92,13 +92,13 @@ func CloudBrains(ctx *context.Context) { return } + models.LoadSpecs4CloudbrainInfo(ciTasks) + for i, task := range ciTasks { ciTasks[i].CanDebug = true ciTasks[i].CanDel = true ciTasks[i].Cloudbrain.ComputeResource = task.ComputeResource ciTasks[i].Cloudbrain.AiCenter = repo.GetCloudbrainAiCenter(task.Cloudbrain, ctx) - _, cardType, _ := repo.GetCloudbrainCardNumAndType(task.Cloudbrain) - ciTasks[i].Cloudbrain.CardType = cardType ciTasks[i].Cloudbrain.Cluster = repo.GetCloudbrainCluster(task.Cloudbrain, ctx) } diff --git a/routers/repo/modelarts.go b/routers/repo/modelarts.go index 2a585432a..3acc9e7b2 100755 --- a/routers/repo/modelarts.go +++ b/routers/repo/modelarts.go @@ -6,6 +6,7 @@ import ( "code.gitea.io/gitea/services/cloudbrain/resource" "encoding/json" "errors" + "fmt" "io" "io/ioutil" "net/http" @@ -131,11 +132,6 @@ func notebookNewDataPrepare(ctx *context.Context) error { ctx.Data["attachments"] = attachs ctx.Data["images"] = setting.StImageInfos.ImageInfo - if modelarts.ImageInfos == nil { - json.Unmarshal([]byte(setting.ImageInfos), &modelarts.ImageInfos) - } - ctx.Data["images"] = modelarts.ImageInfos.ImageInfo - prepareCloudbrainTwoDebugSpecs(ctx) ctx.Data["datasetType"] = models.TypeCloudBrainTwo @@ -147,11 +143,15 @@ func notebookNewDataPrepare(ctx *context.Context) error { } func prepareCloudbrainTwoDebugSpecs(ctx *context.Context) { + aiCenterCode := models.AICenterOfCloudBrainTwo + if setting.ModelartsCD.Enabled { + aiCenterCode = models.AICenterOfChengdu + } noteBookSpecs, _ := resource.FindAvailableSpecs(ctx.User.ID, models.FindSpecsOptions{ JobType: models.JobTypeDebug, ComputeResource: models.NPU, Cluster: models.OpenICluster, - AiCenterCode: models.AICenterOfCloudBrainTwo, + AiCenterCode: aiCenterCode, }) ctx.Data["Specs"] = noteBookSpecs } @@ -240,23 +240,24 @@ func Notebook2Create(ctx *context.Context, form auth.CreateModelArtsNotebookForm return } } - + var aiCenterCode = models.AICenterOfCloudBrainTwo + if setting.ModelartsCD.Enabled { + aiCenterCode = models.AICenterOfChengdu + } spec, err := resource.GetAndCheckSpec(ctx.User.ID, form.SpecId, models.FindSpecsOptions{ JobType: models.JobTypeDebug, ComputeResource: models.NPU, Cluster: models.OpenICluster, - AiCenterCode: models.AICenterOfCloudBrainTwo}) + AiCenterCode: aiCenterCode}) if err != nil || spec == nil { notebookNewDataPrepare(ctx) ctx.RenderWithErr("Resource specification not available", tplModelArtsNotebookNew, &form) return } - err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, uuid, description, imageId, spec) - if setting.ModelartsCD.Enabled { - err = modelarts_cd.GenerateNotebook(ctx, displayJobName, jobName, uuid, description, flavor, imageId) + err = modelarts_cd.GenerateNotebook(ctx, displayJobName, jobName, uuid, description, imageId, spec) } else { - err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, uuid, description, flavor, imageId) + err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, uuid, description, imageId, spec) } if err != nil { diff --git a/routers/user/home.go b/routers/user/home.go index d8c2565c6..78e6c00e9 100755 --- a/routers/user/home.go +++ b/routers/user/home.go @@ -836,14 +836,12 @@ func Cloudbrains(ctx *context.Context) { ctx.ServerError("Get job failed:", err) return } - + models.LoadSpecs4CloudbrainInfo(ciTasks) for i, task := range ciTasks { ciTasks[i].CanDebug = true ciTasks[i].CanDel = true ciTasks[i].Cloudbrain.ComputeResource = task.ComputeResource ciTasks[i].Cloudbrain.AiCenter = repo.GetCloudbrainAiCenter(task.Cloudbrain, ctx) - _, cardType, _ := repo.GetCloudbrainCardNumAndType(task.Cloudbrain) - ciTasks[i].Cloudbrain.CardType = cardType ciTasks[i].Cloudbrain.Cluster = repo.GetCloudbrainCluster(task.Cloudbrain, ctx) }