aiforge/models/ai_model_manage.go

package models

import (
	"fmt"
	"time"

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

type AiModelManage struct {
	ID                string             `xorm:"pk" json:"id"`
	Name              string             `xorm:"INDEX NOT NULL" json:"name"`
	ModelType         int                `xorm:"NULL" json:"modelType"`
	Version           string             `xorm:"NOT NULL" json:"version"`
	VersionCount      int                `xorm:"NOT NULL DEFAULT 0" json:"versionCount"`
	New               int                `xorm:"NOT NULL" json:"new"`
	Type              int                `xorm:"NOT NULL" json:"type"`
	Size              int64              `xorm:"NOT NULL" json:"size"`
	Description       string             `xorm:"varchar(2000)" json:"description"`
	Label             string             `xorm:"varchar(1000)" json:"label"`
	Path              string             `xorm:"varchar(400) NOT NULL" json:"path"`
	DownloadCount     int                `xorm:"NOT NULL DEFAULT 0" json:"downloadCount"`
	Engine            int64              `xorm:"NOT NULL DEFAULT 0" json:"engine"`
	Status            int                `xorm:"NOT NULL DEFAULT 0" json:"status"`
	StatusDesc        string             `xorm:"varchar(500)" json:"statusDesc"`
	Accuracy          string             `xorm:"varchar(1000)" json:"accuracy"`
	AttachmentId      string             `xorm:"NULL" json:"attachmentId"`
	RepoId            int64              `xorm:"INDEX NULL" json:"repoId"`
	CodeBranch        string             `xorm:"varchar(400) NULL" json:"codeBranch"`
	CodeCommitID      string             `xorm:"NULL" json:"codeCommitID"`
	UserId            int64              `xorm:"NOT NULL" json:"userId"`
	IsPrivate         bool               `xorm:"DEFAULT true" json:"isPrivate"`
	UserName          string             `json:"userName"`
	UserRelAvatarLink string             `json:"userRelAvatarLink"`
	TrainTaskInfo     string             `xorm:"text NULL" json:"trainTaskInfo"`
	CreatedUnix       timeutil.TimeStamp `xorm:"created" json:"createdUnix"`
	UpdatedUnix       timeutil.TimeStamp `xorm:"updated" json:"updatedUnix"`
	IsCanOper         bool               `json:"isCanOper"`
	IsCanDelete       bool               `json:"isCanDelete"`
	IsCanDownload     bool               `json:"isCanDownload"`
}

type AiModelConvert struct {
	ID                 string             `xorm:"pk" json:"id"`
	Name               string             `xorm:"INDEX NOT NULL" json:"name"`
	Status             string             `xorm:"NULL" json:"status"`
	StatusResult       string             `xorm:"NULL" json:"statusResult"`
	SrcEngine          int                `xorm:"NOT NULL DEFAULT 0" json:"srcEngine"`
	RepoId             int64              `xorm:"INDEX NULL" json:"repoId"`
	ModelId            string             `xorm:"NOT NULL" json:"modelId"`
	ModelName          string             `xorm:"NULL" json:"modelName"`
	ModelVersion       string             `xorm:"NOT NULL" json:"modelVersion"`
	ModelPath          string             `xorm:"NULL" json:"modelPath"`
	DestFormat         int                `xorm:"NOT NULL DEFAULT 0" json:"destFormat"`
	NetOutputFormat    int                `xorm:"NULL" json:"netOutputFormat"`
	UserId             int64              `xorm:"NOT NULL" json:"userId"`
	CloudBrainTaskId   string             `xorm:"NULL" json:"cloudBrainTaskId"`
	ModelArtsVersionId string             `xorm:"NULL" json:"modelArtsVersionId"`
	ContainerID        string             `json:"containerID"`
	ContainerIp        string             `json:"containerIp"`
	RunTime            int64              `xorm:"NULL" json:"runTime"`
	TrainJobDuration   string             `json:"trainJobDuration"`
	InputShape         string             `xorm:"varchar(2000)" json:"inputShape"`
	InputDataFormat    string             `xorm:"NOT NULL" json:"inputDataFormat"`
	Description        string             `xorm:"varchar(2000)" json:"description"`
	Path               string             `xorm:"varchar(400) NOT NULL" json:"path"`
	CreatedUnix        timeutil.TimeStamp `xorm:"created" json:"createdUnix"`
	UpdatedUnix        timeutil.TimeStamp `xorm:"updated" json:"updatedUnix"`
	StartTime          timeutil.TimeStamp `json:"startTime"`
	EndTime            timeutil.TimeStamp `json:"endTime"`
	UserName           string             `json:"userName"`
	UserRelAvatarLink  string             `json:"userRelAvatarLink"`
	IsCanOper          bool               `json:"isCanOper"`
	IsCanDelete        bool               `json:"isCanDelete"`
}

type AiModelQueryOptions struct {
	ListOptions
	RepoID   int64 // include all repos if empty
	UserID   int64
	ModelID  string
	SortType string
	New      int
	// JobStatus     CloudbrainStatus
	Type           int
	Status         int
	IsOnlyThisRepo bool
	IsQueryPrivate bool
}

func (a *AiModelConvert) IsGpuTrainTask() bool {
	if a.SrcEngine == 0 || a.SrcEngine == 1 || a.SrcEngine == 4 || a.SrcEngine == 6 {
		return true
	}
	return false
}

func ModelComputeAndSetDuration(task *AiModelConvert, result JobResultPayload) {
	if task.StartTime == 0 {
		task.StartTime = timeutil.TimeStamp(result.JobStatus.CreatedTime / 1000)
	}
	if task.EndTime == 0 {
		if result.JobStatus.CompletedTime > 0 {
			task.EndTime = timeutil.TimeStamp(result.JobStatus.CompletedTime / 1000)
		}
	}
	var d int64
	if task.StartTime == 0 {
		d = 0
	} else if task.EndTime == 0 {
		d = time.Now().Unix() - task.StartTime.AsTime().Unix()
	} else {
		d = task.EndTime.AsTime().Unix() - task.StartTime.AsTime().Unix()
	}

	if d < 0 {
		d = 0
	}
	task.RunTime = d
	task.TrainJobDuration = ConvertDurationToStr(d)
}

func ModelConvertSetDuration(task *AiModelConvert) {
	var d int64
	if task.StartTime == 0 {
		d = 0
	} else if task.EndTime == 0 {
		d = time.Now().Unix() - task.StartTime.AsTime().Unix()
	} else {
		d = task.EndTime.AsTime().Unix() - task.StartTime.AsTime().Unix()
	}

	if d < 0 {
		d = 0
	}
	task.RunTime = d
	task.TrainJobDuration = ConvertDurationToStr(d)
}

func UpdateModelConvertModelArts(id string, CloudBrainTaskId string, VersionId string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("cloud_brain_task_id,model_arts_version_id").Update(&AiModelConvert{
		CloudBrainTaskId:   CloudBrainTaskId,
		ModelArtsVersionId: VersionId,
	})
	if err != nil {
		return err
	}
	log.Info("success to update cloud_brain_task_id from  db.re=" + fmt.Sprint((re)))
	return nil
}

func UpdateModelConvertFailed(id string, status string, statusResult string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("status", "status_result").Update(&AiModelConvert{
		Status:       status,
		StatusResult: statusResult,
	})
	if err != nil {
		return err
	}
	log.Info("success to update cloud_brain_task_id from  db.re=" + fmt.Sprint((re)))
	return nil
}

func UpdateModelConvertCBTI(id string, CloudBrainTaskId string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("cloud_brain_task_id").Update(&AiModelConvert{
		CloudBrainTaskId: CloudBrainTaskId,
	})
	if err != nil {
		return err
	}
	log.Info("success to update cloud_brain_task_id from  db.re=" + fmt.Sprint((re)))
	return nil
}

func UpdateModelConvert(job *AiModelConvert) error {
	return updateModelConvert(x, job)
}

func updateModelConvert(e Engine, job *AiModelConvert) error {
	var sess *xorm.Session
	sess = e.Where("id = ?", job.ID)
	_, err := sess.Cols("status", "train_job_duration", "run_time", "start_time", "end_time", "updated_unix").Update(job)
	return err
}

func SaveModelConvert(modelConvert *AiModelConvert) error {
	sess := x.NewSession()
	defer sess.Close()
	re, err := sess.Insert(modelConvert)
	if err != nil {
		log.Info("insert modelConvert error." + err.Error())
		return err
	}
	log.Info("success to save modelConvert db.re=" + fmt.Sprint((re)))
	return nil
}

func SaveModelToDb(model *AiModelManage) error {
	sess := x.NewSession()
	defer sess.Close()

	re, err := sess.Insert(model)
	if err != nil {
		log.Info("insert error." + err.Error())
		return err
	}
	log.Info("success to save db.re=" + fmt.Sprint((re)))
	return nil
}

func QueryModelConvertByName(name string, repoId int64) ([]*AiModelConvert, error) {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table(new(AiModelConvert)).
		Where("name='" + name + "' and repo_id=" + fmt.Sprint(repoId)).OrderBy("created_unix desc")
	aiModelManageConvertList := make([]*AiModelConvert, 0)
	err := sess.Find(&aiModelManageConvertList)
	if err == nil {
		return aiModelManageConvertList, nil
	}
	return nil, err
}

func QueryModelConvertById(id string) (*AiModelConvert, error) {
	sess := x.NewSession()
	defer sess.Close()
	re := new(AiModelConvert)
	isExist, err := sess.Table(new(AiModelConvert)).ID(id).Get(re)
	if err == nil && isExist {
		return re, nil
	}
	return nil, err
}

func QueryModelById(id string) (*AiModelManage, error) {
	sess := x.NewSession()
	defer sess.Close()
	re := new(AiModelManage)
	isExist, err := sess.Table(new(AiModelManage)).ID(id).Get(re)
	if err == nil && isExist {
		return re, nil
	}
	return nil, err
}

func DeleteModelConvertById(id string) error {
	sess := x.NewSession()
	defer sess.Close()
	re, err := sess.Delete(&AiModelConvert{
		ID: id,
	})
	if err != nil {
		return err
	}
	log.Info("success to delete AiModelManageConvert from  db.re=" + fmt.Sprint((re)))
	return nil
}

func DeleteModelById(id string) error {
	sess := x.NewSession()
	defer sess.Close()
	re, err := sess.Delete(&AiModelManage{
		ID: id,
	})
	if err != nil {
		return err
	}
	log.Info("success to delete from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelDescription(id string, description string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("description").Update(&AiModelManage{
		Description: description,
	})
	if err != nil {
		return err
	}
	log.Info("success to update description from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelPrivate(id string, isPrivate bool) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("is_private").Update(&AiModelManage{
		IsPrivate: isPrivate,
	})
	if err != nil {
		return err
	}
	log.Info("success to update isPrivate from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyLocalModel(id string, name, label, description string, engine int, isPrivate bool) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("name", "label", "description", "engine", "is_private").Update(&AiModelManage{
		Description: description,
		Name:        name,
		Label:       label,
		Engine:      int64(engine),
		IsPrivate:   isPrivate,
	})
	if err != nil {
		return err
	}
	log.Info("success to update description from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelSize(id string, size int64) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("size").Update(&AiModelManage{
		Size: size,
	})
	if err != nil {
		return err
	}
	log.Info("success to update size from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelStatus(id string, modelSize int64, status int, modelPath string, statusDesc string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("size", "status", "path", "status_desc").Update(&AiModelManage{
		Size:       modelSize,
		Status:     status,
		Path:       modelPath,
		StatusDesc: statusDesc,
	})
	if err != nil {
		return err
	}
	log.Info("success to update ModelStatus from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelNewProperty(id string, new int, versioncount int) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("new", "version_count").Update(&AiModelManage{
		New:          new,
		VersionCount: versioncount,
	})
	if err != nil {
		return err
	}
	log.Info("success to update new property from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelDownloadCount(id string) error {
	sess := x.NewSession()
	defer sess.Close()
	if _, err := sess.Exec("UPDATE `ai_model_manage` SET download_count = download_count + 1 WHERE id = ?", id); err != nil {
		return err
	}

	return nil
}

func QueryModelByName(name string, repoId int64) []*AiModelManage {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table("ai_model_manage").
		Where("name='" + name + "' and repo_id=" + fmt.Sprint(repoId)).OrderBy("created_unix desc")
	aiModelManageList := make([]*AiModelManage, 0)
	sess.Find(&aiModelManageList)
	return aiModelManageList
}

func QueryModelByPath(path string) (*AiModelManage, error) {
	modelManage := new(AiModelManage)
	has, err := x.Where("path=?", path).Get(modelManage)
	if err != nil {
		return nil, err
	}
	if !has {
		return nil, ErrNotExist{}
	}
	return modelManage, nil
}

func QueryModel(opts *AiModelQueryOptions) ([]*AiModelManage, int64, error) {
	sess := x.NewSession()
	defer sess.Close()

	var cond = builder.NewCond()
	if opts.RepoID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.repo_id": opts.RepoID},
		)
	}

	if opts.UserID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.user_id": opts.UserID},
		)
	}

	if opts.New >= 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.new": opts.New},
		)
	}

	if len(opts.ModelID) > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.id": opts.ModelID},
		)
	}

	if (opts.Type) >= 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.type": opts.Type},
		)
	}

	if (opts.Status) >= 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.status": opts.Status},
		)
	}
	if !opts.IsQueryPrivate {
		cond = cond.And(
			builder.Eq{"ai_model_manage.is_private": false},
		)
	}
	count, err := sess.Where(cond).Count(new(AiModelManage))
	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("ai_model_manage.created_unix DESC")
	aiModelManages := make([]*AiModelManage, 0, setting.UI.IssuePagingNum)
	if err := sess.Table("ai_model_manage").Where(cond).
		Find(&aiModelManages); err != nil {
		return nil, 0, fmt.Errorf("Find: %v", err)
	}

	return aiModelManages, count, nil
}

func QueryModelConvertCountByRepoID(repoId int64) int64 {
	convert := new(AiModelConvert)
	total, _ := x.Where("repo_id =?", repoId).Count(convert)
	return total
}

func QueryModelConvertByRepoID(repoId int64) ([]*AiModelConvert, error) {
	sess := x.NewSession()
	defer sess.Close()
	var cond = builder.NewCond()
	cond = cond.And(
		builder.Eq{"ai_model_convert.repo_id": repoId},
	)
	sess.OrderBy("ai_model_convert.created_unix DESC")
	aiModelManageConvert := make([]*AiModelConvert, 0)
	if err := sess.Table(new(AiModelConvert)).Where(cond).
		Find(&aiModelManageConvert); err != nil {
		return nil, fmt.Errorf("Find: %v", err)
	}
	return aiModelManageConvert, nil
}

func QueryModelConvertByUserID(userID int64) ([]*AiModelConvert, error) {
	sess := x.NewSession()
	defer sess.Close()
	var cond = builder.NewCond()
	cond = cond.And(
		builder.Eq{"ai_model_convert.user_id": userID},
	)
	sess.OrderBy("ai_model_convert.created_unix DESC")
	aiModelManageConvert := make([]*AiModelConvert, 0)
	if err := sess.Table(new(AiModelConvert)).Where(cond).
		Find(&aiModelManageConvert); err != nil {
		return nil, fmt.Errorf("Find: %v", err)
	}
	return aiModelManageConvert, nil
}

func QueryModelConvert(opts *AiModelQueryOptions) ([]*AiModelConvert, int64, error) {
	sess := x.NewSession()
	defer sess.Close()
	var cond = builder.NewCond()
	if opts.RepoID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_convert.repo_id": opts.RepoID},
		)
	}
	if opts.UserID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_convert.user_id": opts.UserID},
		)
	}
	count, err := sess.Where(cond).Count(new(AiModelConvert))
	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("ai_model_convert.created_unix DESC")
	aiModelManageConvert := make([]*AiModelConvert, 0, setting.UI.IssuePagingNum)
	if err := sess.Table(new(AiModelConvert)).Where(cond).
		Find(&aiModelManageConvert); err != nil {
		return nil, 0, fmt.Errorf("Find: %v", err)
	}

	return aiModelManageConvert, count, nil
}