aiforge/models/resource_scene.go

package models

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

const (
	Exclusive = iota + 1
	NotExclusive
)

type ResourceScene struct {
	ID           int64 `xorm:"pk autoincr"`
	SceneName    string
	JobType      string
	IsExclusive  bool
	ExclusiveOrg string
	QueueId      int64
	CreatedTime  timeutil.TimeStamp `xorm:"created"`
	CreatedBy    int64
	UpdatedTime  timeutil.TimeStamp `xorm:"updated"`
	UpdatedBy    int64
	DeleteTime   timeutil.TimeStamp `xorm:"deleted"`
	DeletedBy    int64
}

type ResourceSceneSpec struct {
	ID          int64              `xorm:"pk autoincr"`
	SceneId     int64              `xorm:"unique(idx_scene_spec)"`
	SpecId      int64              `xorm:"unique(idx_scene_spec)"`
	CreatedTime timeutil.TimeStamp `xorm:"created"`
}

type ResourceSceneReq struct {
	ID           int64
	SceneName    string
	JobType      string
	IsExclusive  bool
	ExclusiveOrg string
	QueueId      int64
	CreatorId    int64
	SpecIds      []int64
}

type SearchResourceSceneOptions struct {
	ListOptions
	JobType      string
	IsExclusive  int
	AiCenterCode string
	QueueId      int64
}

type ResourceSceneListRes struct {
	TotalSize int64
	List      []ResourceSceneRes
}

func NewResourceSceneListRes(totalSize int64, list []ResourceSceneRes) *ResourceSceneListRes {
	return &ResourceSceneListRes{
		TotalSize: totalSize,
		List:      list,
	}
}

type ResourceSceneRes struct {
	ID              int64
	SceneName       string
	JobType         JobType
	IsExclusive     bool
	ExclusiveOrg    string
	Cluster         string
	AiCenterCode    string
	QueueCode       string
	QueueId         int64
	ComputeResource string
	AccCardType     string
	Specs           []ResourceSpecWithSceneId
}

func (ResourceSceneRes) TableName() string {
	return "resource_scene"
}

type ResourceSceneBriefRes struct {
	ID        int64
	SceneName string
}

func (ResourceSceneBriefRes) TableName() string {
	return "resource_scene"
}

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

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

func InsertResourceScene(r ResourceSceneReq) error {
	sess := x.NewSession()
	defer sess.Close()

	//check
	specs := make([]ResourceSpecification, 0)
	cond := builder.In("id", r.SpecIds).And(builder.Eq{"queue_id": r.QueueId}).And(builder.Eq{"status": SpecOnShelf})
	if err := sess.Where(cond).Find(&specs); err != nil {
		return err
	}
	if len(specs) < len(r.SpecIds) {
		return errors.New("specIds not correct")
	}

	rs := ResourceScene{
		SceneName:    r.SceneName,
		JobType:      r.JobType,
		IsExclusive:  r.IsExclusive,
		ExclusiveOrg: r.ExclusiveOrg,
		QueueId:      r.QueueId,
		CreatedBy:    r.CreatorId,
		UpdatedBy:    r.CreatorId,
	}
	_, err := sess.InsertOne(&rs)
	if err != nil {
		sess.Rollback()
		return err
	}

	if len(r.SpecIds) == 0 {
		return sess.Commit()
	}
	rss := make([]ResourceSceneSpec, len(r.SpecIds))
	for i, v := range r.SpecIds {
		rss[i] = ResourceSceneSpec{
			SceneId: rs.ID,
			SpecId:  v,
		}
	}

	_, err = sess.Insert(&rss)
	if err != nil {
		sess.Rollback()
		return err
	}

	return sess.Commit()
}

func UpdateResourceScene(r ResourceSceneReq) error {
	sess := x.NewSession()
	var err error
	defer func() {
		if err != nil {
			sess.Rollback()
		}
		sess.Close()
	}()

	// find old scene
	old := ResourceScene{}
	if has, _ := sess.ID(r.ID).Get(&old); !has {
		return errors.New("ResourceScene not exist")
	}
	//check specification
	specs := make([]ResourceSpecification, 0)
	cond := builder.In("id", r.SpecIds).And(builder.Eq{"queue_id": old.QueueId}).And(builder.Eq{"status": SpecOnShelf})
	if err := sess.Where(cond).Find(&specs); err != nil {
		return err
	}
	if len(specs) < len(r.SpecIds) {
		return errors.New("specIds not correct")
	}

	//update scene
	rs := ResourceScene{
		SceneName:    r.SceneName,
		IsExclusive:  r.IsExclusive,
		ExclusiveOrg: r.ExclusiveOrg,
	}
	if _, err = sess.ID(r.ID).UseBool("is_exclusive").Update(&rs); err != nil {
		return err
	}

	//delete scene spec relation
	if _, err = sess.Where("scene_id = ? ", r.ID).Delete(&ResourceSceneSpec{}); err != nil {
		sess.Rollback()
		return err
	}

	if len(r.SpecIds) == 0 {
		return sess.Commit()
	}
	//build new scene spec relation
	rss := make([]ResourceSceneSpec, len(r.SpecIds))
	for i, v := range r.SpecIds {
		rss[i] = ResourceSceneSpec{
			SceneId: r.ID,
			SpecId:  v,
		}
	}
	if _, err = sess.Insert(&rss); err != nil {
		sess.Rollback()
		return err
	}

	return sess.Commit()
}

func DeleteResourceScene(sceneId int64) error {
	sess := x.NewSession()
	var err error
	defer func() {
		if err != nil {
			sess.Rollback()
		}
		sess.Close()
	}()

	if _, err = sess.ID(sceneId).Delete(&ResourceScene{}); err != nil {
		return err
	}
	if _, err = sess.Where("scene_id = ? ", sceneId).Delete(&ResourceSceneSpec{}); err != nil {
		return err
	}
	return sess.Commit()
}

func SearchResourceScene(opts SearchResourceSceneOptions) (int64, []ResourceSceneRes, error) {
	var cond = builder.NewCond()
	if opts.Page <= 0 {
		opts.Page = 1
	}
	if opts.JobType != "" {
		cond = cond.And(builder.Eq{"resource_scene.job_type": opts.JobType})
	}
	if opts.IsExclusive == Exclusive {
		cond = cond.And(builder.Eq{"resource_scene.is_exclusive": 1})
	} else if opts.IsExclusive == NotExclusive {
		cond = cond.And(builder.Eq{"resource_scene.is_exclusive": 0})
	}
	if opts.AiCenterCode != "" {
		cond = cond.And(builder.Eq{"resource_queue.ai_center_code": opts.AiCenterCode})
	}
	if opts.QueueId > 0 {
		cond = cond.And(builder.Eq{"resource_scene.queue_id": opts.QueueId})
	}
	cond = cond.And(builder.NewCond().Or(builder.Eq{"resource_scene.delete_time": 0}).Or(builder.IsNull{"resource_scene.delete_time"}))
	count, err := x.Where(cond).Join("INNER", "resource_queue", "resource_queue.id = resource_scene.queue_id").Count(&ResourceSceneRes{})
	if err != nil {
		return 0, nil, err
	}
	cols := []string{"resource_scene.id", "resource_scene.scene_name", "resource_scene.job_type", "resource_scene.is_exclusive",
		"resource_scene.exclusive_org", "resource_queue.cluster", "resource_queue.ai_center_code", "resource_queue.acc_card_type",
		"resource_queue.id as queue_id", "resource_queue.compute_resource", "resource_queue.queue_code"}
	r := make([]ResourceSceneRes, 0)
	if err = x.Where(cond).Cols(cols...).
		Join("INNER", "resource_queue", "resource_queue.id = resource_scene.queue_id").
		Desc("resource_scene.id").
		Limit(opts.PageSize, (opts.Page-1)*opts.PageSize).
		Find(&r); err != nil {
		return 0, nil, err
	}

	if len(r) == 0 {
		return 0, r, err
	}
	//find related specs
	sceneIds := make([]int64, 0, len(r))
	for _, v := range r {
		sceneIds = append(sceneIds, v.ID)
	}

	specs := make([]ResourceSpecWithSceneId, 0)

	if err := x.Cols("resource_specification.id",
		"resource_specification.source_spec_id",
		"resource_specification.acc_cards_num",
		"resource_specification.cpu_cores",
		"resource_specification.mem_gi_b",
		"resource_specification.gpu_mem_gi_b",
		"resource_specification.share_mem_gi_b",
		"resource_specification.unit_price",
		"resource_specification.status",
		"resource_specification.updated_time",
		"resource_scene_spec.scene_id",
	).In("resource_scene_spec.scene_id", sceneIds).
		Join("INNER", "resource_scene_spec", "resource_scene_spec.spec_id = resource_specification.id").
		OrderBy("resource_specification.acc_cards_num").
		Find(&specs); err != nil {
		return 0, nil, err
	}

	specsMap := make(map[int64][]ResourceSpecWithSceneId, 0)
	for _, v := range specs {
		if _, ok := specsMap[v.SceneId]; !ok {
			specsMap[v.SceneId] = []ResourceSpecWithSceneId{v}
		} else {
			specsMap[v.SceneId] = append(specsMap[v.SceneId], v)
		}
	}

	for i, v := range r {
		s := specsMap[v.ID]
		if s == nil {
			s = make([]ResourceSpecWithSceneId, 0)
		}
		r[i].Specs = s
	}

	return count, r, nil
}