提交代码

models/cloudbrain.go

@@ -306,6 +306,22 @@ type CreateJobResult struct {
 	Payload map[string]interface{} `json:"payload"`
 }
 
+type QueueDetailResult struct {
+	Code    string                 `json:"code"`
+	Msg     string                 `json:"msg"`
+	Payload map[string]QueueDetail `json:"payload"`
+}
+
+type QueueDetail struct {
+	JobScheduleInfo JobScheduleInfo `json:"JobScheduleInfo"`
+}
+
+type JobScheduleInfo struct {
+	Pending                     int `json:"Pending"`
+	Running                     int `json:"Running"`
+	MedianPendingJobDurationSec int `json:"MedianPendingJobDurationSec"`
+}
+
 type GetJobResult struct {
 	Code    string                 `json:"code"`
 	Msg     string                 `json:"msg"`
@@ -1743,7 +1759,10 @@ func GetBenchmarkCountByUserID(userID int64) (int, error) {
 }
 
 func GetWaitingCloudbrainCount(cloudbrainType int, computeResource string, jobTypes ...JobType) (int64, error) {
-	sess := x.In("JobType", jobTypes).And("status=? and type=?", JobWaiting, cloudbrainType)
+	sess := x.Where("status=? and type=?", JobWaiting, cloudbrainType)
+	if len(jobTypes) >= 0 {
+		sess.In("JobType", jobTypes)
+	}
 	if computeResource != "" {
 		sess.And("compute_resource=?", computeResource)
 	}

modules/cloudbrain/resty.go

@@ -30,6 +30,7 @@ const (
 	LogPageSize		  = 500
 	LogPageTokenExpired = "5m"
 	pageSize            = 15
+	QueuesDetailUrl     = "/rest-server/api/v2/queuesdetail"
 )
 
 func getRestyClient() *resty.Client {
@@ -73,6 +74,36 @@ func loginCloudbrain() error {
 	return nil
 }
 
+func GetQueuesDetail() (*map[string]int, error) {
+	checkSetting()
+	client := getRestyClient()
+	var jobResult models.QueueDetailResult
+
+	var result = make(map[string]int, 0)
+
+	res, err := client.R().
+		SetHeader("Content-Type", "application/json").
+		SetAuthToken(TOKEN).
+		SetResult(&jobResult).
+		Get(HOST + QueuesDetailUrl)
+
+	if err != nil {
+		return nil, fmt.Errorf("resty get queues detail failed: %s", err)
+	}
+
+	if jobResult.Code != Success {
+		return nil, fmt.Errorf("jobResult err: %s", res.String())
+	}
+
+	for k, v := range jobResult.Payload {
+
+		result[k] = v.JobScheduleInfo.Pending
+
+	}
+	return &result, nil
+
+}
+
 func CreateJob(jobName string, createJobParams models.CreateJobParams) (*models.CreateJobResult, error) {
 	checkSetting()
 	client := getRestyClient()
@@ -145,16 +176,6 @@ sendjob:
 	return &getJobResult, nil
 }
 
-func GetImages() (*models.GetImagesResult, error) {
-
-	return GetImagesPageable(1, 100, Custom, "")
-
-}
-
-func GetPublicImages() (*models.GetImagesResult, error) {
-	return GetImagesPageable(1, 100, Public, "")
-}
-
 func GetImagesPageable(page int, size int, imageType string, name string) (*models.GetImagesResult, error) {
 	checkSetting()
 	client := getRestyClient()

routers/repo/cloudbrain.go

@@ -112,7 +112,10 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 	ctx.Data["benchmark_categories"] = categories.Category
 
 	ctx.Data["benchmark_types"] = GetBenchmarkTypes(ctx).BenchmarkType
-
+	queuesDetail, _ := cloudbrain.GetQueuesDetail()
+	if queuesDetail != nil {
+		ctx.Data["QueuesDetail"] = queuesDetail
+	}
 	if gpuInfos == nil {
 		json.Unmarshal([]byte(setting.GpuTypes), &gpuInfos)
 	}
@@ -164,8 +167,6 @@ func CloudBrainNew(ctx *context.Context) {
 		ctx.ServerError("get new cloudbrain info failed", err)
 		return
 	}
-	waitCount := cloudbrain.GetWaitingCloudbrainCount(models.TypeCloudBrainOne, "", models.JobTypeDebug)
-	ctx.Data["WaitCount"] = waitCount
 	ctx.HTML(200, tplCloudBrainNew)
 }
 
@@ -1922,8 +1923,6 @@ func CloudBrainBenchmarkNew(ctx *context.Context) {
 		ctx.ServerError("get new cloudbrain info failed", err)
 		return
 	}
-	waitCount := cloudbrain.GetWaitingCloudbrainCount(models.TypeCloudBrainOne, "", models.JobTypeBrainScore, models.JobTypeSnn4imagenet, models.JobTypeBenchmark)
-	ctx.Data["WaitCount"] = waitCount
 	ctx.HTML(200, tplCloudBrainBenchmarkNew)
 }
 
@@ -2367,8 +2366,6 @@ func CloudBrainTrainJobNew(ctx *context.Context) {
 		ctx.ServerError("get new train-job info failed", err)
 		return
 	}
-	waitCount := cloudbrain.GetWaitingCloudbrainCount(models.TypeCloudBrainOne, "", models.JobTypeTrain)
-	ctx.Data["WaitCount"] = waitCount
 	ctx.HTML(http.StatusOK, tplCloudBrainTrainJobNew)
 }
 
@@ -2378,8 +2375,6 @@ func InferenceCloudBrainJobNew(ctx *context.Context) {
 		ctx.ServerError("get new train-job info failed", err)
 		return
 	}
-	waitCount := cloudbrain.GetWaitingCloudbrainCount(models.TypeCloudBrainOne, "", models.JobTypeInference)
-	ctx.Data["WaitCount"] = waitCount
 	ctx.HTML(http.StatusOK, tplCloudBrainInferenceJobNew)
 }