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
	// Title         string `xorm:"INDEX NOT NULL"`
	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 struct {
		UserID   string `json:"userId"`
		RealName string `json:"realName"`
		Token    string `json:"token"`
		Admin    bool   `json:"admin"`
	}
}

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 struct {
		JobID string `json:"jobId"`
	}
}

type GetJobResult struct {
	Code    string
	Msg     string
	Payload struct {
		ID        string `json:"Id"`
		Name      string
		Platform  string
		JobStatus struct {
			Username           string
			State              string
			SubState           string `json:"subState"`
			ExecutionType      string `json:"executionType"`
			Retries            int8   `json:"retries"`
			CreatedTime        int64  `json:"createdTime"`
			CompletedTime      int64  `json:"completedTime"`
			AppID              string `json:"appId"`
			AppProgress        string `json:"appProgress"`
			AppTrackingURL     string `json:"appTrackingUrl"`
			AppLaunchedTime    int64  `json:"appLaunchedTime"`
			AppCompletedTime   int64  `json:"appCompletedTime"`
			AppExitCode        int8   `json:"appExitCode"`
			AppExitDiagnostics string `json:"appExitDiagnostics"`
			AppExitType        string `json:"appExitType"`
			VirtualCluster     string `json:"virtualCluster"`
		} `json:"jobStatus"`

		TaskRoles string `json:"taskRoles"`

		Resource struct {
			CPU    int8 `json:"cpu"`
			Memory string
			GPU    string `json:"nvidia.com/gpu"`
		} `json:"resource"`

		Config struct {
			Image      string
			JobID      string `json:"jobId"`
			GpuType    string `json:"gpuType"`
			JobName    string `json:"jobName"`
			JobType    string `json:"jobType"`
			RetryCount int8   `json:"retryCount"`
			TaskRoles  []struct {
				Name                  string `json:"name"`
				ShmMB                 int32  `json:"shmMB"`
				Command               string `json:"command"`
				MemoryMB              int64  `json:"memoryMB"`
				CPUNumber             int8   `json:"cpuNumber"`
				GPUNumber             int8   `json:"gpuNumber"`
				IsMainRole            bool   `json:"isMainRole"`
				TaskNumber            int32  `json:"taskNumber"`
				NeedIBDevice          bool   `json:"needIBDevice"`
				MinFailedTaskCount    int8   `json:"minFailedTaskCount"`
				MinSucceededTaskCount int8   `json:"minSucceededTaskCount"`
			} `json:"taskRoles"`
		}

		Userinfo struct {
			User  string
			OrgID string `json:"org_id"`
		}
	}
}

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
}