#2701

update

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
+}

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)

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)
 	}
 

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 {

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)
 
 	}