From 4ed9a05c202989778b44a185b3c5197af757dc3c Mon Sep 17 00:00:00 2001
From: chenyifan01 <cyf_uestc@outlook.com>
Date: Wed, 7 Sep 2022 10:55:43 +0800
Subject: [PATCH] #2863 #2862 update

---
 models/resource_specification.go                       |  1 +
 routers/repo/modelarts.go                              |  4 ++++
 services/cloudbrain/resource/resource_specification.go | 15 +++++++++++++--
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/models/resource_specification.go b/models/resource_specification.go
index b81105220..2da8d015d 100644
--- a/models/resource_specification.go
+++ b/models/resource_specification.go
@@ -525,6 +525,7 @@ func InitCloudbrainTwoSpecs() (map[string]*Specification, error) {
 				MemGiB:       float32(i * 256),
 				GPUMemGiB:    float32(32),
 				Status:       SpecOffShelf,
+				IsAvailable:  true,
 			}
 			_, err = x.Insert(spec)
 			if err != nil {
diff --git a/routers/repo/modelarts.go b/routers/repo/modelarts.go
index 7b90372cd..f78a50318 100755
--- a/routers/repo/modelarts.go
+++ b/routers/repo/modelarts.go
@@ -971,6 +971,10 @@ func trainJobNewVersionDataPrepare(ctx *context.Context) error {
 	ctx.Data["engine_versions"] = versionInfos.Version
 
 	prepareCloudbrainTwoTrainSpecs(ctx)
+	spec, _ := resource.GetCloudbrainSpec(task.ID)
+	if spec != nil {
+		ctx.Data["spec_id"] = spec.ID
+	}
 
 	var Parameters modelarts.Parameters
 	if err = json.Unmarshal([]byte(task.Parameters), &Parameters); err != nil {
diff --git a/services/cloudbrain/resource/resource_specification.go b/services/cloudbrain/resource/resource_specification.go
index f21c12b51..462d903f1 100644
--- a/services/cloudbrain/resource/resource_specification.go
+++ b/services/cloudbrain/resource/resource_specification.go
@@ -438,8 +438,12 @@ func findCloudbrainOneSpecFromRemote(task *models.Cloudbrain) (*models.Specifica
 	var shmMB float32
 	if jobRes.Config.TaskRoles != nil && len(jobRes.Config.TaskRoles) > 0 {
 		shmMB = float32(jobRes.Config.TaskRoles[0].ShmMB) / 1024
+		if jobRes.Config.TaskRoles[0].ShmMB == 103600 {
+			shmMB = 100
+		} else if jobRes.Config.TaskRoles[0].ShmMB == 51800 {
+			shmMB = 50
+		}
 	}
-
 	opt := models.FindSpecsOptions{
 		ComputeResource: task.ComputeResource,
 		Cluster:         models.OpenICluster,
@@ -538,6 +542,12 @@ func findCloudbrainOneSpecFromConfig(task *models.Cloudbrain) (*models.Specifica
 		task.ComputeResource = models.GPU
 	}
 
+	shareMemMiB := float32(specConfig.ShareMemMiB) / 1024
+	if specConfig.ShareMemMiB == 103600 {
+		shareMemMiB = 100
+	} else if specConfig.ShareMemMiB == 51800 {
+		shareMemMiB = 50
+	}
 	opt := models.FindSpecsOptions{
 		JobType:         models.JobType(task.JobType),
 		ComputeResource: task.ComputeResource,
@@ -550,7 +560,7 @@ func findCloudbrainOneSpecFromConfig(task *models.Cloudbrain) (*models.Specifica
 		UseCpuCores:     true,
 		MemGiB:          float32(specConfig.MemMiB) / 1024,
 		UseMemGiB:       true,
-		ShareMemGiB:     float32(specConfig.ShareMemMiB) / 1024,
+		ShareMemGiB:     shareMemMiB,
 		UseShareMemGiB:  true,
 		RequestAll:      true,
 	}
@@ -639,5 +649,6 @@ func InitQueueAndSpec(opt models.FindSpecsOptions, aiCenterName string, remark s
 		GPUMemGiB:   opt.GPUMemGiB,
 		ShareMemGiB: opt.ShareMemGiB,
 		Status:      models.SpecOffShelf,
+		IsAvailable: true,
 	})
 }