aiforge/models/cloudbrain.go

package models

import (
	"fmt"

	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/timeutil"
	"xorm.io/builder"
)

type CloudbrainStatus int8

const (
	JobWaiting CloudbrainStatus = iota
	JobStopped
	JobSucceeded
	JobFailed
)

type Cloudbrain struct {
	ID          int64  `xorm:"pk autoincr"`
	JobID       string `xorm:"INDEX NOT NULL"`
	JobName     string
	Status      int32              `xorm:"INDEX"`
	UserID      int64              `xorm:"INDEX"`
	RepoID      int64              `xorm:"INDEX"`
	CreatedUnix timeutil.TimeStamp `xorm:"INDEX created"`
	UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"`

	User *User       `xorm:"-"`
	Repo *Repository `xorm:"-"`
}

type CloudBrainLoginResult struct {
	Code    string
	Msg     string
	Payload map[string]interface{}
}

type TaskRole struct {
	Name                  string `json:"name"`
	TaskNumber            int8   `json:"taskNumber"`
	MinSucceededTaskCount int8   `json:"minSucceededTaskCount"`
	MinFailedTaskCount    int8   `json:"minFailedTaskCount"`
	CPUNumber             int8   `json:"cpuNumber"`
	GPUNumber             int8   `json:"gpuNumber"`
	MemoryMB              int    `json:"memoryMB"`
	ShmMB                 int    `json:"shmMB"`
	Command               string `json:"command"`
	NeedIBDevice          bool   `json:"needIBDevice"`
	IsMainRole            bool   `json:"isMainRole"`
}

type CreateJobParams struct {
	JobName    string     `json:"jobName"`
	RetryCount int8       `json:"retryCount"`
	GpuType    string     `json:"gpuType"`
	Image      string     `json:"image"`
	TaskRoles  []TaskRole `json:"taskRoles"`
}

type CreateJobResult struct {
	Code    string
	Msg     string
	Payload map[string]interface{}
}

type GetJobResult struct {
	Code    string                 `json:"code"`
	Msg     string                 `json:"msg"`
	Payload map[string]interface{} `json:"payload"`
}

type CloudbrainsOptions struct {
	ListOptions
	RepoID        int64 // include all repos if empty
	UserID        int64
	JobStatus     CloudbrainStatus
	SortType      string
	CloudbrainIDs []int64
}

func Cloudbrains(opts *CloudbrainsOptions) ([]*Cloudbrain, int64, error) {
	sess := x.NewSession()
	defer sess.Close()

	var cond = builder.NewCond()
	if opts.RepoID > 0 {
		cond.And(
			builder.Eq{"cloudbrain.repo_id": opts.RepoID},
		)
	}

	if opts.UserID > 0 {
		cond.And(
			builder.Eq{"cloudbrain.user_id": opts.UserID},
		)
	}

	switch opts.JobStatus {
	case JobWaiting:
		cond.And(builder.Eq{"cloudbrain.status": int(JobWaiting)})
	case JobFailed:
		cond.And(builder.Eq{"cloudbrain.status": int(JobFailed)})
	case JobStopped:
		cond.And(builder.Eq{"cloudbrain.status": int(JobStopped)})
	case JobSucceeded:
		cond.And(builder.Eq{"cloudbrain.status": int(JobSucceeded)})
	}

	if len(opts.CloudbrainIDs) > 0 {
		cond.And(builder.In("cloudbrain.id", opts.CloudbrainIDs))
	}

	count, err := sess.Where(cond).Count(new(Cloudbrain))
	if err != nil {
		return nil, 0, fmt.Errorf("Count: %v", err)
	}

	if opts.Page >= 0 && opts.PageSize > 0 {
		var start int
		if opts.Page == 0 {
			start = 0
		} else {
			start = (opts.Page - 1) * opts.PageSize
		}
		sess.Limit(opts.PageSize, start)
	}

	sess.OrderBy("cloudbrain.created_unix DESC")
	cloudbrains := make([]*Cloudbrain, 0, setting.UI.IssuePagingNum)
	if err := sess.Find(&cloudbrains); err != nil {
		return nil, 0, fmt.Errorf("Find: %v", err)
	}
	sess.Close()

	return cloudbrains, count, nil
}

func CreateCloudbrain(cloudbrain *Cloudbrain) (err error) {
	if _, err = x.Insert(cloudbrain); err != nil {
		return err
	}

	return nil
}