aiforge/routers/admin/cloudbrains.go

package admin

import (
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/360EntSecGroup-Skylar/excelize/v2"

	"code.gitea.io/gitea/modules/modelarts"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

const (
	tplCloudBrains     base.TplName = "admin/cloudbrain/list"
	EXCEL_DATE_FORMAT               = "20060102150405"
	CREATE_TIME_FORMAT              = "2006/01/02 15:04:05.00"
)

func CloudBrains(ctx *context.Context) {
	ctx.Data["Title"] = ctx.Tr("admin.cloudBrains")
	ctx.Data["PageIsAdmin"] = true
	ctx.Data["PageIsAdminCloudBrains"] = true

	listType := ctx.Query("listType")
	jobType := ctx.Query("jobType")
	jobStatus := ctx.Query("jobStatus")

	ctx.Data["ListType"] = listType
	ctx.Data["JobType"] = jobType
	ctx.Data["JobStatus"] = jobStatus

	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	debugType := modelarts.DebugType
	if listType == models.GPUResource {
		debugType = models.TypeCloudBrainOne
	} else if listType == models.NPUResource {
		debugType = models.TypeCloudBrainTwo
	}

	var jobTypes []string
	jobTypeNot := false
	if jobType == string(models.JobTypeDebug) {
		jobTypes = append(jobTypes, string(models.JobTypeSnn4imagenet), string(models.JobTypeBrainScore), string(models.JobTypeDebug))
	} else if jobType != "all" && jobType != "" {
		jobTypes = append(jobTypes, jobType)
	}

	var jobStatuses []string
	jobStatusNot := false
	if jobStatus == "other" {
		jobStatusNot = true
		jobStatuses = append(jobStatuses, string(models.ModelArtsTrainJobWaiting), string(models.ModelArtsTrainJobFailed), string(models.ModelArtsRunning), string(models.ModelArtsTrainJobCompleted),
			string(models.ModelArtsStarting), string(models.ModelArtsRestarting), string(models.ModelArtsStartFailed),
			string(models.ModelArtsStopping), string(models.ModelArtsStopped))
	} else if jobStatus != "all" && jobStatus != "" {
		jobStatuses = append(jobStatuses, jobStatus)
	}

	keyword := strings.Trim(ctx.Query("q"), " ")

	ciTasks, count, err := models.Cloudbrains(&models.CloudbrainsOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		Keyword:      keyword,
		Type:         debugType,
		JobTypeNot:   jobTypeNot,
		JobStatusNot: jobStatusNot,
		JobStatus:    jobStatuses,
		JobTypes:     jobTypes,
		NeedRepoInfo: true,
	})
	if err != nil {
		ctx.ServerError("Get job failed:", err)
		return
	}

	for i, task := range ciTasks {
		ciTasks[i].CanDebug = true
		ciTasks[i].CanDel = true
		ciTasks[i].Cloudbrain.ComputeResource = task.ComputeResource
	}

	pager := context.NewPagination(int(count), setting.UI.IssuePagingNum, page, getTotalPage(count, setting.UI.IssuePagingNum))
	pager.SetDefaultParams(ctx)
	pager.AddParam(ctx, "listType", "ListType")
	ctx.Data["Page"] = pager
	ctx.Data["PageIsCloudBrain"] = true
	ctx.Data["Tasks"] = ciTasks
	ctx.Data["CanCreate"] = true
	ctx.Data["Keyword"] = keyword

	ctx.HTML(200, tplCloudBrains)

}

func DownloadCloudBrains(ctx *context.Context) {

	page := 1

	pageSize := 300

	var cloudBrain = ctx.Tr("repo.cloudbrain")
	fileName := getFileName(cloudBrain)

	_, total, err := models.Cloudbrains(&models.CloudbrainsOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: 1,
		},
		Type:         modelarts.DebugType,
		NeedRepoInfo: false,
	})

	if err != nil {
		log.Warn("Can not get cloud brain info", err)
		ctx.Error(http.StatusBadRequest, ctx.Tr("repo.cloudbrain_query_fail"))
		return
	}

	totalPage := getTotalPage(total, pageSize)

	f := excelize.NewFile()

	index := f.NewSheet(cloudBrain)
	f.DeleteSheet("Sheet1")

	for k, v := range allHeader(ctx) {
		f.SetCellValue(cloudBrain, k, v)
	}

	var row = 2
	for i := 0; i < totalPage; i++ {

		pageRecords, _, err := models.Cloudbrains(&models.CloudbrainsOptions{
			ListOptions: models.ListOptions{
				Page:     page,
				PageSize: pageSize,
			},
			Type:         modelarts.DebugType,
			NeedRepoInfo: true,
		})
		if err != nil {
			log.Warn("Can not get cloud brain info", err)
			continue
		}
		for _, record := range pageRecords {

			for k, v := range allValues(row, record, ctx) {
				f.SetCellValue(cloudBrain, k, v)
			}
			row++

		}

		page++
	}
	f.SetActiveSheet(index)

	ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+url.QueryEscape(fileName))
	ctx.Resp.Header().Set("Content-Type", "application/octet-stream")

	f.WriteTo(ctx.Resp)
}

func allValues(row int, rs *models.CloudbrainInfo, ctx *context.Context) map[string]string {
	return map[string]string{getCellName("A", row): rs.JobName, getCellName("B", row): rs.Status, getCellName("C", row): rs.JobType, getCellName("D", row): time.Unix(int64(rs.Cloudbrain.CreatedUnix), 0).Format(CREATE_TIME_FORMAT), getCellName("E", row): getDurationTime(rs),
		getCellName("F", row): rs.ComputeResource, getCellName("G", row): rs.Name, getCellName("H", row): rs.Repo.OwnerName + "/" + rs.Repo.Alias, getCellName("I", row): rs.JobName,
	}
}

func getDurationTime(rs *models.CloudbrainInfo) string {
	if rs.JobType == "TRAIN" || rs.JobType == "INFERENCE" {
		return rs.TrainJobDuration
	} else {
		return "-"
	}
}

func getFileName(baseName string) string {
	return baseName + "_" + time.Now().Format(EXCEL_DATE_FORMAT) + ".xlsx"

}

func getTotalPage(total int64, pageSize int) int {

	another := 0
	if int(total)%pageSize != 0 {
		another = 1
	}
	return int(total)/pageSize + another

}

func allHeader(ctx *context.Context) map[string]string {

	return map[string]string{"A1": ctx.Tr("repo.cloudbrain_task"), "B1": ctx.Tr("repo.modelarts.status"), "C1": ctx.Tr("repo.cloudbrain_task_type"), "D1": ctx.Tr("repo.modelarts.createtime"), "E1": ctx.Tr("repo.modelarts.train_job.dura_time"), "F1": ctx.Tr("repo.modelarts.computing_resources"), "G1": ctx.Tr("repo.cloudbrain_creator"), "H1": ctx.Tr("repo.repo_name"), "I1": ctx.Tr("repo.cloudbrain_task_name")}

}

func getCellName(col string, row int) string {
	return col + strconv.Itoa(row)
}