aiforge/models/cloudbrain.go

package models

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"

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

type CloudbrainStatus string

const (
	JobWaiting   CloudbrainStatus = "WAITING"
	JobStopped   CloudbrainStatus = "STOPPED"
	JobSucceeded CloudbrainStatus = "SUCCEEDED"
	JobFailed    CloudbrainStatus = "FAILED"
)

type Cloudbrain struct {
	ID          int64  `xorm:"pk autoincr"`
	JobID       string `xorm:"INDEX NOT NULL"`
	JobName     string
	Status      string             `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
	JobID         int64
	SortType      string
	CloudbrainIDs []int64
	// JobStatus     CloudbrainStatus
}
type TaskPod struct {
	TaskRoleStatus struct {
		Name string `json:"name"`
	} `json:"taskRoleStatus"`
	TaskStatuses []struct {
		TaskIndex       int       `json:"taskIndex"`
		PodUID          string    `json:"podUid"`
		PodIP           string    `json:"podIp"`
		PodName         string    `json:"podName"`
		ContainerID     string    `json:"containerId"`
		ContainerIP     string    `json:"containerIp"`
		ContainerGpus   string    `json:"containerGpus"`
		State           string    `json:"state"`
		StartAt         time.Time `json:"startAt"`
		FinishedAt      time.Time `json:"finishedAt"`
		ExitCode        int       `json:"exitCode"`
		ExitDiagnostics string    `json:"exitDiagnostics"`
		RetriedCount    int       `json:"retriedCount"`
	} `json:"taskStatuses"`
}

func ConvertToTaskPod(input map[string]interface{}) (TaskPod, error) {
	data, _ := json.Marshal(input)
	var taskPod TaskPod
	err := json.Unmarshal(data, &taskPod)
	return taskPod, err
}

type JobResultPayload struct {
	ID        string `json:"id"`
	Name      string `json:"name"`
	Platform  string `json:"platform"`
	JobStatus struct {
		Username           string      `json:"username"`
		State              string      `json:"state"`
		SubState           string      `json:"subState"`
		ExecutionType      string      `json:"executionType"`
		Retries            int         `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   interface{} `json:"appCompletedTime"`
		AppExitCode        int         `json:"appExitCode"`
		AppExitDiagnostics string      `json:"appExitDiagnostics"`
		AppExitType        interface{} `json:"appExitType"`
		VirtualCluster     string      `json:"virtualCluster"`
	} `json:"jobStatus"`
	TaskRoles map[string]interface{} `json:"taskRoles"`
	Resource  struct {
		CPU          int    `json:"cpu"`
		Memory       string `json:"memory"`
		NvidiaComGpu int    `json:"nvidia.com/gpu"`
	} `json:"resource"`
	Config struct {
		Image     string `json:"image"`
		JobID     string `json:"jobId"`
		GpuType   string `json:"gpuType"`
		JobName   string `json:"jobName"`
		JobType   string `json:"jobType"`
		TaskRoles []struct {
			Name                  string `json:"name"`
			ShmMB                 int    `json:"shmMB"`
			Command               string `json:"command"`
			MemoryMB              int    `json:"memoryMB"`
			CPUNumber             int    `json:"cpuNumber"`
			GpuNumber             int    `json:"gpuNumber"`
			IsMainRole            bool   `json:"isMainRole"`
			TaskNumber            int    `json:"taskNumber"`
			NeedIBDevice          bool   `json:"needIBDevice"`
			MinFailedTaskCount    int    `json:"minFailedTaskCount"`
			MinSucceededTaskCount int    `json:"minSucceededTaskCount"`
		} `json:"taskRoles"`
		RetryCount int `json:"retryCount"`
	} `json:"config"`
	Userinfo struct {
		User  string `json:"user"`
		OrgID string `json:"org_id"`
	} `json:"userinfo"`
}

func ConvertToJobResultPayload(input map[string]interface{}) (JobResultPayload, error) {
	data, _ := json.Marshal(input)
	var jobResultPayload JobResultPayload
	err := json.Unmarshal(data, &jobResultPayload)
	return jobResultPayload, err
}

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},
		)
	}

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

	// 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
}

func getRepoCloudBrain(cb *Cloudbrain) (*Cloudbrain, error) {
	has, err := x.Get(cb)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, errors.New("cloudbrain task is not found")
	}
	return cb, nil
}

func GetRepoCloudBrainByJobID(repoID int64, jobID string) (*Cloudbrain, error) {
	cb := &Cloudbrain{JobID: jobID, RepoID: repoID}
	return getRepoCloudBrain(cb)
}

func GetCloudbrainByJobID(jobID string) (*Cloudbrain, error) {
	cb := &Cloudbrain{JobID: jobID}
	return getRepoCloudBrain(cb)
}

func SetCloudbrainStatusByJobID(jobID string, status CloudbrainStatus) (err error) {
	cb := &Cloudbrain{JobID: jobID, Status: string(status)}
	_, err = x.Cols("status").Where("cloudbrain.job_id=?", jobID).Update(cb)
	return
}