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"`
	Name              string `xorm:"INDEX NOT NULL"`
	Version           string `xorm:"NOT NULL"`
	VersionCount      int    `xorm:"NOT NULL DEFAULT 0"`
	New               int    `xorm:"NOT NULL"`
	Type              int    `xorm:"NOT NULL"`
	Size              int64  `xorm:"NOT NULL"`
	Description       string `xorm:"varchar(2000)"`
	Label             string `xorm:"varchar(1000)"`
	Path              string `xorm:"varchar(400) NOT NULL"`
	DownloadCount     int    `xorm:"NOT NULL DEFAULT 0"`
	Engine            int64  `xorm:"NOT NULL DEFAULT 0"`
	Status            int    `xorm:"NOT NULL DEFAULT 0"`
	Accuracy          string `xorm:"varchar(1000)"`
	AttachmentId      string `xorm:"NULL"`
	RepoId            int64  `xorm:"INDEX NULL"`
	CodeBranch        string `xorm:"varchar(400) NULL"`
	CodeCommitID      string `xorm:"NULL"`
	UserId            int64  `xorm:"NOT NULL"`
	UserName          string
	UserRelAvatarLink string
	TrainTaskInfo     string             `xorm:"text NULL"`
	CreatedUnix       timeutil.TimeStamp `xorm:"created"`
	UpdatedUnix       timeutil.TimeStamp `xorm:"updated"`
	IsCanOper         bool
	IsCanDelete       bool
}

type AiModelConvert struct {
	ID                 string `xorm:"pk"`
	Name               string `xorm:"INDEX NOT NULL"`
	Status             string `xorm:"NULL"`
	StatusResult       string `xorm:"NULL"`
	SrcEngine          int    `xorm:"NOT NULL DEFAULT 0"`
	RepoId             int64  `xorm:"INDEX NULL"`
	ModelId            string `xorm:"NOT NULL"`
	ModelName          string `xorm:"NULL"`
	ModelVersion       string `xorm:"NOT NULL"`
	ModelPath          string `xorm:"NULL"`
	DestFormat         int    `xorm:"NOT NULL DEFAULT 0"`
	NetOutputFormat    int    `xorm:"NULL"`
	UserId             int64  `xorm:"NOT NULL"`
	CloudBrainTaskId   string `xorm:"NULL"`
	ModelArtsVersionId string `xorm:"NULL"`
	ContainerID        string
	ContainerIp        string
	RunTime            int64 `xorm:"NULL"`
	TrainJobDuration   string
	InputShape         string             `xorm:"varchar(2000)"`
	InputDataFormat    string             `xorm:"NOT NULL"`
	Description        string             `xorm:"varchar(2000)"`
	Path               string             `xorm:"varchar(400) NOT NULL"`
	CreatedUnix        timeutil.TimeStamp `xorm:"created"`
	UpdatedUnix        timeutil.TimeStamp `xorm:"updated"`
	StartTime          timeutil.TimeStamp
	EndTime            timeutil.TimeStamp
	UserName           string
	UserRelAvatarLink  string
	IsCanOper          bool
	IsCanDelete        bool
}

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
}

func (a *AiModelConvert) IsGpuTrainTask() bool {
	if a.SrcEngine == 0 || a.SrcEngine == 1 {
		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 QueryModelConvertById(id string) (*AiModelConvert, error) {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table(new(AiModelConvert)).Where("id='" + id + "'")
	aiModelManageConvertList := make([]*AiModelConvert, 0)
	err := sess.Find(&aiModelManageConvertList)
	if err == nil {
		if len(aiModelManageConvertList) == 1 {
			return aiModelManageConvertList[0], nil
		}
	}
	return nil, err
}

func QueryModelById(id string) (*AiModelManage, error) {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table("ai_model_manage").
		Where("id='" + id + "'")
	aiModelManageList := make([]*AiModelManage, 0)
	err := sess.Find(&aiModelManageList)
	if err == nil {
		if len(aiModelManageList) == 1 {
			return aiModelManageList[0], nil
		}
	} else {
		log.Info("error=" + err.Error())
	}
	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 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 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},
		)
	}

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