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

type FindSpecsOptions struct {
	JobType         JobType
	ComputeResource string
	Cluster         string
	AiCenterCode    string
	SpecId          int64
	QueueCode       string
	SourceSpecId    string
	AccCardsNum     int
	UseAccCardsNum  bool
	AccCardType     string
	CpuCores        int
	UseCpuCores     bool
	MemGiB          float32
	UseMemGiB       bool
	GPUMemGiB       float32
	UseGPUMemGiB    bool
	ShareMemGiB     float32
	UseShareMemGiB  bool
	//if true,find specs no matter used or not used in scene. if false,only find specs used in scene
	RequestAll bool
}

type Specification struct {
	ID              int64
	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
}

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

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

//FindSpecs
func FindSpecs(opts FindSpecsOptions) ([]*Specification, error) {
	var cond = builder.NewCond()
	if !opts.RequestAll && opts.JobType != "" {
		cond = cond.And(builder.Eq{"resource_scene.job_type": opts.JobType})
	}
	if opts.ComputeResource != "" {
		cond = cond.And(builder.Eq{"resource_queue.compute_resource": opts.ComputeResource})
	}
	if opts.ComputeResource != "" {
		cond = cond.And(builder.Eq{"resource_queue.cluster": opts.Cluster})
	}
	if opts.AiCenterCode != "" {
		cond = cond.And(builder.Eq{"resource_queue.ai_center_code": opts.AiCenterCode})
	}
	if opts.SpecId > 0 {
		cond = cond.And(builder.Eq{"resource_specification.id": opts.SpecId})
	}
	if opts.QueueCode != "" {
		cond = cond.And(builder.Eq{"resource_queue.queue_code": opts.QueueCode})
	}
	if opts.SourceSpecId != "" {
		cond = cond.And(builder.Eq{"resource_specification.source_spec_id": opts.SourceSpecId})
	}
	if opts.UseAccCardsNum {
		cond = cond.And(builder.Eq{"resource_specification.acc_cards_num": opts.AccCardsNum})
	}
	if opts.AccCardType != "" {
		cond = cond.And(builder.Eq{"resource_queue.acc_card_type": opts.AccCardType})
	}
	if opts.UseCpuCores {
		cond = cond.And(builder.Eq{"resource_specification.cpu_cores": opts.CpuCores})
	}
	if opts.UseMemGiB {
		cond = cond.And(builder.Eq{"resource_specification.mem_gi_b": opts.MemGiB})
	}
	if opts.UseGPUMemGiB {
		cond = cond.And(builder.Eq{"resource_specification.gpu_mem_gi_b": opts.GPUMemGiB})
	}
	if opts.UseShareMemGiB {
		cond = cond.And(builder.Eq{"resource_specification.share_mem_gi_b": opts.ShareMemGiB})
	}
	r := make([]*Specification, 0)
	s := x.Where(cond).
		Join("INNER", "resource_queue", "resource_queue.id = resource_specification.queue_id")

	if !opts.RequestAll {
		s = s.Join("INNER", "resource_scene_spec", "resource_scene_spec.spec_id = resource_specification.id").
			Join("INNER", "resource_scene", "resource_scene_spec.scene_id = resource_scene.id")
	}
	err := s.OrderBy("resource_queue.compute_resource asc,resource_queue.acc_card_type asc,resource_specification.acc_cards_num asc").
		Unscoped().Find(&r)
	if err != nil {
		return nil, err
	}
	return r, nil
}

func InitQueueAndSpec(queue ResourceQueue, spec ResourceSpecification) (*Specification, error) {
	sess := x.NewSession()
	defer sess.Close()

	sess.Begin()
	param := ResourceQueue{
		QueueCode:       queue.QueueCode,
		Cluster:         queue.Cluster,
		AiCenterCode:    queue.AiCenterCode,
		ComputeResource: queue.ComputeResource,
		AccCardType:     queue.AccCardType,
	}
	_, err := sess.Get(&param)
	if err != nil {
		sess.Rollback()
		return nil, err
	}
	if param.ID == 0 {
		_, err = sess.InsertOne(&queue)
		if err != nil {
			sess.Rollback()
			return nil, err
		}
	} else {
		queue = param
	}

	spec.QueueId = queue.ID
	_, err = sess.InsertOne(&spec)
	if err != nil {
		sess.Rollback()
		return nil, err
	}
	sess.Commit()
	return &Specification{
		ID:              spec.ID,
		SourceSpecId:    spec.SourceSpecId,
		AccCardsNum:     spec.AccCardsNum,
		AccCardType:     queue.AccCardType,
		CpuCores:        spec.CpuCores,
		MemGiB:          spec.MemGiB,
		GPUMemGiB:       spec.GPUMemGiB,
		ShareMemGiB:     spec.ShareMemGiB,
		ComputeResource: queue.ComputeResource,
		UnitPrice:       spec.UnitPrice,
		QueueId:         queue.ID,
		QueueCode:       queue.QueueCode,
		Cluster:         queue.Cluster,
		AiCenterCode:    queue.AiCenterCode,
		AiCenterName:    queue.AiCenterName,
	}, nil
}

func GetCloudbrainOneAccCardType(queueCode string) string {
	switch queueCode {
	case "a100":
		return "A100"
	case "openidebug":
		return "T4"
	case "openidgx":
		return "V100"

	}
	return ""
}