aiforge/models/resource_specification.go

package models

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

const (
	SpecNotVerified int = iota + 1
	SpecOnShelf
	SpecOffShelf
)

type ResourceSpecification struct {
	ID              int64  `xorm:"pk autoincr"`
	QueueId         int64  `xorm:"INDEX"`
	SourceSpecId    string `xorm:"INDEX"`
	AccCardsNum     int
	CpuCores        int
	MemGiB          float32
	GPUMemGiB       float32
	ShareMemGiB     float32
	UnitPrice       int
	Status          int
	IsAutomaticSync bool
	CreatedTime     timeutil.TimeStamp `xorm:"created"`
	CreatedBy       int64
	UpdatedTime     timeutil.TimeStamp `xorm:"updated"`
	UpdatedBy       int64
}

func (r ResourceSpecification) ConvertToRes() *ResourceSpecificationRes {
	return &ResourceSpecificationRes{
		ID:           r.ID,
		SourceSpecId: r.SourceSpecId,
		AccCardsNum:  r.AccCardsNum,
		CpuCores:     r.CpuCores,
		MemGiB:       r.MemGiB,
		ShareMemGiB:  r.ShareMemGiB,
		GPUMemGiB:    r.GPUMemGiB,
		UnitPrice:    r.UnitPrice,
		Status:       r.Status,
		UpdatedTime:  r.UpdatedTime,
	}
}

type ResourceSpecificationReq struct {
	QueueId         int64 `binding:"Required"`
	SourceSpecId    string
	AccCardsNum     int
	CpuCores        int
	MemGiB          float32
	GPUMemGiB       float32
	ShareMemGiB     float32
	UnitPrice       int
	Status          int
	IsAutomaticSync bool
	CreatorId       int64
}

func (r ResourceSpecificationReq) ToDTO() ResourceSpecification {
	return ResourceSpecification{
		QueueId:         r.QueueId,
		SourceSpecId:    r.SourceSpecId,
		AccCardsNum:     r.AccCardsNum,
		CpuCores:        r.CpuCores,
		MemGiB:          r.MemGiB,
		GPUMemGiB:       r.GPUMemGiB,
		ShareMemGiB:     r.ShareMemGiB,
		UnitPrice:       r.UnitPrice,
		Status:          r.Status,
		IsAutomaticSync: r.IsAutomaticSync,
		CreatedBy:       r.CreatorId,
		UpdatedBy:       r.CreatorId,
	}
}

type SearchResourceSpecificationOptions struct {
	ListOptions
	QueueId int64
	Status  int
	Cluster string
}

type SearchResourceBriefSpecificationOptions struct {
	QueueId int64
	Cluster string
}

type ResourceSpecAndQueueListRes struct {
	TotalSize int64
	List      []*ResourceSpecAndQueueRes
}

func NewResourceSpecAndQueueListRes(totalSize int64, list []ResourceSpecAndQueue) *ResourceSpecAndQueueListRes {
	resList := make([]*ResourceSpecAndQueueRes, len(list))
	for i, v := range list {
		resList[i] = v.ConvertToRes()
	}
	return &ResourceSpecAndQueueListRes{
		TotalSize: totalSize,
		List:      resList,
	}
}

type ResourceSpecificationRes struct {
	ID           int64
	SourceSpecId string
	AccCardsNum  int
	CpuCores     int
	MemGiB       float32
	GPUMemGiB    float32
	ShareMemGiB  float32
	UnitPrice    int
	Status       int
	UpdatedTime  timeutil.TimeStamp
}

func (ResourceSpecificationRes) TableName() string {
	return "resource_specification"
}

type ResourceSpecAndQueueRes struct {
	Spec  *ResourceSpecificationRes
	Queue *ResourceQueueRes
}

type ResourceSpecAndQueue struct {
	ResourceSpecification `xorm:"extends"`
	ResourceQueue         `xorm:"extends"`
}

func (*ResourceSpecAndQueue) TableName() string {
	return "resource_specification"
}

func (r ResourceSpecAndQueue) ConvertToRes() *ResourceSpecAndQueueRes {
	return &ResourceSpecAndQueueRes{
		Spec:  r.ResourceSpecification.ConvertToRes(),
		Queue: r.ResourceQueue.ConvertToRes(),
	}
}

func InsertResourceSpecification(r ResourceSpecification) (int64, error) {
	return x.Insert(&r)
}

func UpdateResourceSpecificationById(queueId int64, spec ResourceSpecification) (int64, error) {
	return x.ID(queueId).Update(&spec)
}
func UpdateSpecUnitPriceById(id int64, unitPrice int) error {
	_, err := x.Exec("update resource_specification set unit_price = ? ,updated_time = ? where id = ?", unitPrice, timeutil.TimeStampNow(), id)
	return err
}

func SearchResourceSpecification(opts SearchResourceSpecificationOptions) (int64, []ResourceSpecAndQueue, error) {
	var cond = builder.NewCond()
	if opts.Page <= 0 {
		opts.Page = 1
	}
	if opts.QueueId > 0 {
		cond = cond.And(builder.Eq{"resource_specification.queue_id": opts.QueueId})
	}
	if opts.Status > 0 {
		cond = cond.And(builder.Eq{"resource_specification.status": opts.Status})
	}
	if opts.Cluster != "" {
		cond = cond.And(builder.Eq{"resource_queue.cluster": opts.Cluster})
	}
	//cond = cond.And(builder.Or(builder.Eq{"resource_queue.deleted_time": 0}).Or(builder.IsNull{"resource_queue.deleted_time"}))
	n, err := x.Where(cond).Join("INNER", "resource_queue", "resource_queue.ID = resource_specification.queue_id").
		Unscoped().Count(&ResourceSpecAndQueue{})
	if err != nil {
		return 0, nil, err
	}

	r := make([]ResourceSpecAndQueue, 0)
	err = x.Where(cond).
		Join("INNER", "resource_queue", "resource_queue.ID = resource_specification.queue_id").
		Desc("resource_specification.id").
		Limit(opts.PageSize, (opts.Page-1)*opts.PageSize).
		Unscoped().Find(&r)
	if err != nil {
		return 0, nil, err
	}
	return n, r, nil
}

func GetSpecScenes(specId int64) ([]ResourceSceneBriefRes, error) {
	r := make([]ResourceSceneBriefRes, 0)
	err := x.Where("resource_scene_spec.spec_id = ?", specId).
		Join("INNER", "resource_scene_spec", "resource_scene_spec.scene_id = resource_scene.id").
		Find(&r)
	if err != nil {
		return nil, err
	}
	return r, nil
}

func ResourceSpecOnShelf(id int64, unitPrice int) error {
	_, err := x.Exec("update resource_specification set unit_price = ?,updated_time = ?,status = ? where id = ?", unitPrice, timeutil.TimeStampNow(), SpecOnShelf, id)
	return err
}

func ResourceSpecOffShelf(id int64) (int64, error) {
	sess := x.NewSession()
	var err error
	defer func() {
		if err != nil {
			sess.Rollback()
		}
		sess.Close()
	}()
	//delete scene spec relation
	if _, err = sess.Where("spec_id = ?", id).Delete(&ResourceSceneSpec{}); err != nil {
		return 0, err
	}

	param := ResourceSpecification{
		Status: SpecOffShelf,
	}
	n, err := sess.Where("id = ? and status = ?", id, SpecOnShelf).Update(&param)
	if err != nil {
		return 0, err
	}
	sess.Commit()
	return n, err
}

func GetResourceSpecification(r *ResourceSpecification) (*ResourceSpecification, error) {
	has, err := x.Get(r)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, nil
	}
	return r, nil
}

func SyncGrampusSpecs(updateList []ResourceSpecification, insertList []ResourceSpecification, existIds []int64) error {
	sess := x.NewSession()
	var err error
	defer func() {
		if err != nil {
			sess.Rollback()
		}
		sess.Close()
	}()
	//delete specs and scene that no longer exists
	deleteIds := make([]int64, 0)
	cond := builder.NewCond()
	cond = cond.And(builder.NotIn("resource_specification.id", existIds)).And(builder.Eq{"resource_queue.cluster": C2NetCluster})
	if err := sess.Cols("resource_specification.id").Table("resource_specification").
		Where(cond).Join("INNER", "resource_queue", "resource_queue.id = resource_specification.queue_id").
		Find(&deleteIds); err != nil {
		return err
	}
	if len(deleteIds) > 0 {
		if _, err = sess.In("id", deleteIds).Update(&ResourceSpecification{Status: SpecOffShelf}); err != nil {
			return err
		}
		if _, err = sess.In("spec_id", deleteIds).Delete(&ResourceSceneSpec{}); err != nil {
			return err
		}
	}

	//update exists specs
	if len(updateList) > 0 {
		for _, v := range updateList {
			if _, err = sess.ID(v.ID).Update(&v); err != nil {
				return err
			}
		}

	}

	//insert new specs
	if len(insertList) > 0 {
		if _, err = sess.Insert(insertList); err != nil {
			return err
		}
	}

	return sess.Commit()
}