aiforge/models/cloudbrain_spec.go

package models

import (
	"code.gitea.io/gitea/modules/timeutil"
)

type CloudbrainSpec struct {
	CloudbrainID    int64 `xorm:"pk"`
	SpecId          int64 `xorm:"index"`
	SourceSpecId    string
	AccCardsNum     int
	AccCardType     string
	CpuCores        int
	MemGiB          float32
	GPUMemGiB       float32
	ShareMemGiB     float32
	ComputeResource string
	UnitPrice       int
	QueueId         int64
	QueueCode       string
	Cluster         string
	AiCenterCode    string
	AiCenterName    string
	IsExclusive     bool
	ExclusiveOrg    string
	CreatedTime     timeutil.TimeStamp `xorm:"created"`
	UpdatedTime     timeutil.TimeStamp `xorm:"updated"`
}

func (s CloudbrainSpec) ConvertToSpecification() *Specification {
	return &Specification{
		ID:              s.SpecId,
		SourceSpecId:    s.SourceSpecId,
		AccCardsNum:     s.AccCardsNum,
		AccCardType:     s.AccCardType,
		CpuCores:        s.CpuCores,
		MemGiB:          s.MemGiB,
		GPUMemGiB:       s.GPUMemGiB,
		ShareMemGiB:     s.ShareMemGiB,
		ComputeResource: s.ComputeResource,
		UnitPrice:       s.UnitPrice,
		QueueId:         s.QueueId,
		QueueCode:       s.QueueCode,
		Cluster:         s.Cluster,
		AiCenterCode:    s.AiCenterCode,
		AiCenterName:    s.AiCenterName,
		IsExclusive:     s.IsExclusive,
		ExclusiveOrg:    s.ExclusiveOrg,
	}
}

func NewCloudBrainSpec(cloudbrainId int64, s Specification) CloudbrainSpec {
	return CloudbrainSpec{
		CloudbrainID:    cloudbrainId,
		SpecId:          s.ID,
		SourceSpecId:    s.SourceSpecId,
		AccCardsNum:     s.AccCardsNum,
		AccCardType:     s.AccCardType,
		CpuCores:        s.CpuCores,
		MemGiB:          s.MemGiB,
		GPUMemGiB:       s.GPUMemGiB,
		ShareMemGiB:     s.ShareMemGiB,
		ComputeResource: s.ComputeResource,
		UnitPrice:       s.UnitPrice,
		QueueId:         s.QueueId,
		QueueCode:       s.QueueCode,
		Cluster:         s.Cluster,
		AiCenterCode:    s.AiCenterCode,
		AiCenterName:    s.AiCenterName,
		IsExclusive:     s.IsExclusive,
		ExclusiveOrg:    s.ExclusiveOrg,
	}
}

var StatusChangeChan = make(chan *Cloudbrain, 50)

func InsertCloudbrainSpec(c CloudbrainSpec) (int64, error) {
	return x.Insert(&c)
}

func GetCloudbrainSpecByID(cloudbrainId int64) (*CloudbrainSpec, error) {
	r := &CloudbrainSpec{}
	if has, err := x.Where("cloudbrain_id = ?", cloudbrainId).Get(r); err != nil {
		return nil, err
	} else if !has {
		return nil, nil
	}
	return r, nil
}

func FindCloudbrainTask(page, pageSize int) ([]*Cloudbrain, error) {
	r := make([]*Cloudbrain, 0)
	err := x.Unscoped().
		Limit(pageSize, (page-1)*pageSize).
		OrderBy("cloudbrain.id").
		Find(&r)
	if err != nil {
		return nil, err
	}
	return r, nil
}

func CountNoSpecHistoricTask() (int64, error) {
	n, err := x.Unscoped().
		Where(" 1=1 and not exists (select 1 from cloudbrain_spec  where cloudbrain.id = cloudbrain_spec.cloudbrain_id)").
		Count(&Cloudbrain{})
	if err != nil {
		return 0, err
	}
	return n, nil
}

// GetResourceSpecMapByCloudbrainIDs
func GetResourceSpecMapByCloudbrainIDs(ids []int64) (map[int64]*Specification, error) {
	specs := make([]*CloudbrainSpec, 0)
	if err := x.In("cloudbrain_id", ids).Find(&specs); err != nil {
		return nil, err
	}
	r := make(map[int64]*Specification, len(ids))
	for _, s := range specs {
		r[s.CloudbrainID] = s.ConvertToSpecification()
	}
	return r, nil
}

func GetCloudbrainTaskUnitPrice(task Cloudbrain) (int, error) {
	s, err := GetCloudbrainSpecByID(task.ID)
	if err != nil {
		return 0, err
	}
	var n = 1
	if task.WorkServerNumber > 1 {
		n = task.WorkServerNumber
	}
	return s.UnitPrice * n, nil
}