diff --git a/models/cloudbrain.go b/models/cloudbrain.go index 600c8ffc4..b866e4c1a 100755 --- a/models/cloudbrain.go +++ b/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) } diff --git a/modules/cloudbrain/resty.go b/modules/cloudbrain/resty.go index de8d3779c..7b714c4b5 100755 --- a/modules/cloudbrain/resty.go +++ b/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() diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index 0b3613bee..1573df3c8 100755 --- a/routers/repo/cloudbrain.go +++ b/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) }