提交代码

modules/cloudbrain/cloudbrain.go

@@ -259,21 +259,7 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error {
 	}
 	//如果没有匹配到spec信息，尝试从专属资源池获取
 	if resourceSpec == nil && SpecialPools != nil {
-		for _, specialPool := range SpecialPools.Pools {
-			if resourceSpec != nil {
-				break
-			}
-			if specialPool.ResourceSpec != nil {
-				if IsElementExist(specialPool.JobType, req.JobType) && IsQueueInSpecialtPool(specialPool.Pool, req.GpuQueue) {
-					for _, spec := range specialPool.ResourceSpec {
-						if req.ResourceSpecId == spec.Id {
-							resourceSpec = spec
-							break
-						}
-					}
-				}
-			}
-		}
+		resourceSpec = geMatchResourceSpec(req.JobType, req.GpuQueue, req.ResourceSpecId)
 	}
 
 	if resourceSpec == nil {
@@ -464,6 +450,11 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 		}
 	}
 
+	//如果没有匹配到spec信息，尝试从专属资源池获取
+	if resourceSpec == nil && SpecialPools != nil {
+		resourceSpec = geMatchResourceSpec(task.JobType, task.GpuQueue, task.ResourceSpecId)
+	}
+
 	if resourceSpec == nil {
 		log.Error("no such resourceSpecId(%d)", task.ResourceSpecId, ctx.Data["MsgID"])
 		return errors.New("no such resourceSpec")
@@ -602,6 +593,23 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 	return nil
 }
 
+func geMatchResourceSpec(jobType string, gpuQueue string, resourceSpecId int) *models.ResourceSpec {
+
+	for _, specialPool := range SpecialPools.Pools {
+
+		if specialPool.ResourceSpec != nil {
+			if IsElementExist(specialPool.JobType, jobType) && IsQueueInSpecialtPool(specialPool.Pool, gpuQueue) {
+				for _, spec := range specialPool.ResourceSpec {
+					if resourceSpecId == spec.Id {
+						return spec
+					}
+				}
+			}
+		}
+	}
+	return nil
+}
+
 func DelCloudBrainJob(jobId string) string {
 	task, err := models.GetCloudbrainByJobID(jobId)
 	if err != nil {

routers/repo/cloudbrain.go

@@ -134,7 +134,7 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 	if inferenceGpuInfos == nil {
 		json.Unmarshal([]byte(setting.InferenceGpuTypes), &inferenceGpuInfos)
 	}
-	ctx.Data["train_gpu_types"] = trainGpuInfos.GpuInfo
+	ctx.Data["inference_gpu_types"] = inferenceGpuInfos.GpuInfo
 
 	if benchmarkGpuInfos == nil {
 		json.Unmarshal([]byte(setting.BenchmarkGpuTypes), &benchmarkGpuInfos)
@@ -558,7 +558,18 @@ func CloudBrainRestart(ctx *context.Context) {
 		for _, resourceType := range gpuInfos.GpuInfo {
 			if resourceType.Queue == task.GpuQueue {
 				hasSameResource = true
-				continue
+				break
+			}
+		}
+		if !hasSameResource && cloudbrain.SpecialPools != nil {
+
+			for _, specialPool := range cloudbrain.SpecialPools.Pools {
+				cloudbrain.IsElementExist(specialPool.JobType, string(models.JobTypeDebug))
+				for _, pool := range specialPool.Pool {
+					if pool.Queue == task.GpuQueue {
+						hasSameResource = true
+					}
+				}
 			}
 		}