aiforge/services/cloudbrain/cloudbrainTask/count.go

package cloudbrainTask

import (
	"fmt"
	"strconv"

	"code.gitea.io/gitea/models"
)

type StatusInfo struct {
	CloudBrainTypes  []int
	JobType          []models.JobType
	NotFinalStatuses []string
	ComputeResource  string
}

var cloudbrainOneNotFinalStatuses = []string{string(models.JobWaiting), string(models.JobRunning)}
var cloudbrainTwoNotFinalStatuses = []string{string(models.ModelArtsTrainJobInit), string(models.ModelArtsTrainJobImageCreating), string(models.ModelArtsTrainJobSubmitTrying), string(models.ModelArtsTrainJobWaiting), string(models.ModelArtsTrainJobRunning), string(models.ModelArtsTrainJobScaling), string(models.ModelArtsTrainJobCheckInit), string(models.ModelArtsTrainJobCheckRunning), string(models.ModelArtsTrainJobCheckRunningCompleted)}
var grampusTwoNotFinalStatuses = []string{models.GrampusStatusWaiting, models.GrampusStatusRunning}
var StatusInfoDict = map[string]StatusInfo{string(models.JobTypeDebug) + "-" + strconv.Itoa(models.TypeCloudBrainOne): {
	CloudBrainTypes:  []int{models.TypeCloudBrainOne},
	JobType:          []models.JobType{models.JobTypeDebug},
	NotFinalStatuses: cloudbrainOneNotFinalStatuses,
	ComputeResource:  models.GPUResource,
}, string(models.JobTypeTrain) + "-" + strconv.Itoa(models.TypeCloudBrainOne): {
	CloudBrainTypes:  []int{models.TypeCloudBrainOne},
	JobType:          []models.JobType{models.JobTypeTrain},
	NotFinalStatuses: cloudbrainOneNotFinalStatuses,
	ComputeResource:  models.GPUResource,
}, string(models.JobTypeInference) + "-" + strconv.Itoa(models.TypeCloudBrainOne): {
	CloudBrainTypes:  []int{models.TypeCloudBrainOne},
	JobType:          []models.JobType{models.JobTypeInference},
	NotFinalStatuses: cloudbrainOneNotFinalStatuses,
	ComputeResource:  models.GPUResource,
}, string(models.JobTypeBenchmark) + "-" + strconv.Itoa(models.TypeCloudBrainOne): {
	CloudBrainTypes:  []int{models.TypeCloudBrainOne},
	JobType:          []models.JobType{models.JobTypeBenchmark, models.JobTypeBrainScore, models.JobTypeSnn4imagenet},
	NotFinalStatuses: cloudbrainOneNotFinalStatuses,
	ComputeResource:  models.GPUResource,
}, string(models.JobTypeDebug) + "-" + strconv.Itoa(models.TypeCloudBrainTwo): {
	CloudBrainTypes:  []int{models.TypeCloudBrainTwo, models.TypeCDCenter},
	JobType:          []models.JobType{models.JobTypeDebug},
	NotFinalStatuses: []string{string(models.ModelArtsCreateQueue), string(models.ModelArtsCreating), string(models.ModelArtsStarting), string(models.ModelArtsReadyToStart), string(models.ModelArtsResizing), string(models.ModelArtsStartQueuing), string(models.ModelArtsRunning), string(models.ModelArtsRestarting)},
	ComputeResource:  models.NPUResource,
}, string(models.JobTypeTrain) + "-" + strconv.Itoa(models.TypeCloudBrainTwo): {
	CloudBrainTypes:  []int{models.TypeCloudBrainTwo},
	JobType:          []models.JobType{models.JobTypeTrain},
	NotFinalStatuses: cloudbrainTwoNotFinalStatuses,
	ComputeResource:  models.NPUResource,
}, string(models.JobTypeInference) + "-" + strconv.Itoa(models.TypeCloudBrainTwo): {
	CloudBrainTypes:  []int{models.TypeCloudBrainTwo},
	JobType:          []models.JobType{models.JobTypeInference},
	NotFinalStatuses: cloudbrainTwoNotFinalStatuses,
	ComputeResource:  models.NPUResource,
}, string(models.JobTypeTrain) + "-" + strconv.Itoa(models.TypeC2Net) + "-" + models.GPUResource: {
	CloudBrainTypes:  []int{models.TypeC2Net},
	JobType:          []models.JobType{models.JobTypeTrain},
	NotFinalStatuses: grampusTwoNotFinalStatuses,
	ComputeResource:  models.GPUResource,
}, string(models.JobTypeTrain) + "-" + strconv.Itoa(models.TypeC2Net) + "-" + models.NPUResource: {
	CloudBrainTypes:  []int{models.TypeC2Net},
	JobType:          []models.JobType{models.JobTypeTrain},
	NotFinalStatuses: grampusTwoNotFinalStatuses,
	ComputeResource:  models.NPUResource,
}}

func GetNotFinalStatusTaskCount(uid int64, cloudbrainType int, jobType string, computeResource ...string) (int, error) {
	jobNewType := jobType
	if jobType == string(models.JobTypeSnn4imagenet) || jobType == string(models.JobTypeBrainScore) {
		jobNewType = string(models.JobTypeBenchmark)
	}

	key := jobNewType + "-" + strconv.Itoa(cloudbrainType)
	if len(computeResource) > 0 {
		key = key + "-" + computeResource[0]
	}

	if statusInfo, ok := StatusInfoDict[key]; ok {

		return models.GetNotFinalStatusTaskCount(uid, statusInfo.NotFinalStatuses, statusInfo.JobType, statusInfo.CloudBrainTypes, statusInfo.ComputeResource)

	} else {
		return 0, fmt.Errorf("Can not find the status info.")
	}

}