aiforge/routers/repo/ai_model_manage.go

package repo

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"path"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/storage"
	uuid "github.com/satori/go.uuid"
)

const (
	Model_prefix           = "aimodels/"
	tplModelManageIndex    = "repo/modelmanage/index"
	tplModelManageDownload = "repo/modelmanage/download"
	tplModelInfo           = "repo/modelmanage/showinfo"
	MODEL_LATEST           = 1
	MODEL_NOT_LATEST       = 0
)

func saveModelByParameters(jobId string, versionName string, name string, version string, label string, description string, ctx *context.Context) error {
	aiTask, err := models.GetCloudbrainByJobIDAndVersionName(jobId, versionName)
	if err != nil {
		log.Info("query task error." + err.Error())
		return err
	}

	uuid := uuid.NewV4()
	id := uuid.String()
	modelPath := id
	var lastNewModelId string
	var modelSize int64
	cloudType := models.TypeCloudBrainTwo

	log.Info("find task name:" + aiTask.JobName)
	aimodels := models.QueryModelByName(name, aiTask.RepoID)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.Version == version {
				return errors.New(ctx.Tr("repo.model.manage.create_error"))
			}
			if model.New == MODEL_LATEST {
				lastNewModelId = model.ID
			}
		}
	}
	cloudType = aiTask.Type
	//download model zip   //train type
	if cloudType == models.TypeCloudBrainTwo {
		modelPath, modelSize, err = downloadModelFromCloudBrainTwo(id, aiTask.JobName, "", aiTask.TrainUrl)
		if err != nil {
			log.Info("download model from CloudBrainTwo faild." + err.Error())
			return err
		}
	}
	accuracy := make(map[string]string)
	accuracy["F1"] = ""
	accuracy["Recall"] = ""
	accuracy["Accuracy"] = ""
	accuracy["Precision"] = ""
	accuracyJson, _ := json.Marshal(accuracy)
	log.Info("accuracyJson=" + string(accuracyJson))
	aiTaskJson, _ := json.Marshal(aiTask)

	model := &models.AiModelManage{
		ID:            id,
		Version:       version,
		VersionCount:  len(aimodels) + 1,
		Label:         label,
		Name:          name,
		Description:   description,
		New:           MODEL_LATEST,
		Type:          cloudType,
		Path:          modelPath,
		Size:          modelSize,
		AttachmentId:  aiTask.Uuid,
		RepoId:        aiTask.RepoID,
		UserId:        ctx.User.ID,
		CodeBranch:    aiTask.BranchName,
		CodeCommitID:  aiTask.CommitID,
		Engine:        aiTask.EngineID,
		TrainTaskInfo: string(aiTaskJson),
		Accuracy:      string(accuracyJson),
	}

	err = models.SaveModelToDb(model)
	if err != nil {
		return err
	}
	if len(lastNewModelId) > 0 {
		//udpate status and version count
		models.ModifyModelNewProperty(lastNewModelId, MODEL_NOT_LATEST, 0)
	}
	var units []models.RepoUnit
	var deleteUnitTypes []models.UnitType
	units = append(units, models.RepoUnit{
		RepoID: ctx.Repo.Repository.ID,
		Type:   models.UnitTypeModelManage,
		Config: &models.ModelManageConfig{
			EnableModelManage: true,
		},
	})
	deleteUnitTypes = append(deleteUnitTypes, models.UnitTypeModelManage)

	models.UpdateRepositoryUnits(ctx.Repo.Repository, units, deleteUnitTypes)

	log.Info("save model end.")

	return nil
}

func SaveNewNameModel(ctx *context.Context) {
	name := ctx.Query("Name")
	if name == "" {
		ctx.Error(500, fmt.Sprintf("name or version is null."))
		return
	}

	aimodels := models.QueryModelByName(name, ctx.Repo.Repository.ID)
	if len(aimodels) > 0 {
		ctx.Error(500, ctx.Tr("repo.model_rename"))
		return
	}
	SaveModel(ctx)

	log.Info("save model end.")
}

func SaveModel(ctx *context.Context) {
	log.Info("save model start.")
	JobId := ctx.Query("JobId")
	VersionName := ctx.Query("VersionName")
	name := ctx.Query("Name")
	version := ctx.Query("Version")
	label := ctx.Query("Label")
	description := ctx.Query("Description")
	trainTaskCreate := ctx.QueryBool("trainTaskCreate")

	if !trainTaskCreate {
		if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
			ctx.ServerError("No right.", errors.New(ctx.Tr("repo.model_noright")))
			return
		}
	}

	if JobId == "" || VersionName == "" {
		ctx.Error(500, fmt.Sprintf("JobId or VersionName is null."))
		return
	}

	if name == "" || version == "" {
		ctx.Error(500, fmt.Sprintf("name or version is null."))
		return
	}

	err := saveModelByParameters(JobId, VersionName, name, version, label, description, ctx)

	if err != nil {
		log.Info("save model error." + err.Error())
		ctx.Error(500, fmt.Sprintf("save model error. %v", err))
		return
	}

	log.Info("save model end.")
}

func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir string, trainUrl string) (string, int64, error) {
	objectkey := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir), "/")
	if trainUrl != "" {
		objectkey = strings.Trim(trainUrl[len(setting.Bucket)+1:], "/")
	}

	modelDbResult, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, objectkey, "")
	log.Info("bucket=" + setting.Bucket + "  objectkey=" + objectkey)
	if err != nil {
		log.Info("get TrainJobListModel failed:", err)
		return "", 0, err
	}
	if len(modelDbResult) == 0 {
		return "", 0, errors.New("cannot create model, as model is empty.")
	}

	prefix := objectkey + "/"
	destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(modelUUID) + "/"

	size, err := storage.ObsCopyManyFile(setting.Bucket, prefix, setting.Bucket, destKeyNamePrefix)

	dataActualPath := setting.Bucket + "/" + destKeyNamePrefix
	return dataActualPath, size, nil
}

func DeleteModel(ctx *context.Context) {
	log.Info("delete model start.")
	id := ctx.Query("ID")
	err := deleteModelByID(ctx, id)
	if err != nil {
		ctx.JSON(500, err.Error())
	} else {
		ctx.JSON(200, map[string]string{
			"result_code": "0",
		})
	}
}

func isCanDelete(ctx *context.Context, model *models.AiModelManage) bool {
	if ctx.User == nil {
		return false
	}
	if ctx.User.IsAdmin || ctx.User.ID == model.UserId {
		return true
	}
	if ctx.Repo.IsOwner() {
		return true
	}
	return false
}

func deleteModelByID(ctx *context.Context, id string) error {
	log.Info("delete model start. id=" + id)
	model, err := models.QueryModelById(id)
	if !isCanDelete(ctx, model) {
		return errors.New(ctx.Tr("repo.model_noright"))
	}
	if err == nil {
		log.Info("bucket=" + setting.Bucket + " path=" + model.Path)
		if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
			err := storage.ObsRemoveObject(setting.Bucket, model.Path[len(setting.Bucket)+1:])
			if err != nil {
				log.Info("Failed to delete model. id=" + id)
				return err
			}
		}
		err = models.DeleteModelById(id)
		if err == nil { //find a model to change new
			aimodels := models.QueryModelByName(model.Name, model.RepoId)
			if model.New == MODEL_LATEST {
				if len(aimodels) > 0 {
					//udpate status and version count
					models.ModifyModelNewProperty(aimodels[0].ID, MODEL_LATEST, len(aimodels))
				}
			} else {
				for _, tmpModel := range aimodels {
					if tmpModel.New == MODEL_LATEST {
						models.ModifyModelNewProperty(tmpModel.ID, MODEL_LATEST, len(aimodels))
						break
					}
				}
			}
		}
	}
	return err
}

func QueryModelByParameters(repoId int64, page int) ([]*models.AiModelManage, int64, error) {

	return models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repoId,
		Type:   -1,
		New:    MODEL_LATEST,
	})
}

func DownloadMultiModelFile(ctx *context.Context) {
	log.Info("DownloadMultiModelFile start.")
	id := ctx.Query("ID")
	log.Info("id=" + id)
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if !isOper(ctx, task.UserId) {
		ctx.ServerError("no right.", errors.New(ctx.Tr("repo.model_noright")))
		return
	}

	path := Model_prefix + models.AttachmentRelativePath(id) + "/"

	allFile, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, path)
	if err == nil {
		//count++
		models.ModifyModelDownloadCount(id)

		returnFileName := task.Name + "_" + task.Version + ".zip"
		ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+returnFileName)
		ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
		w := zip.NewWriter(ctx.Resp)
		defer w.Close()
		for _, oneFile := range allFile {
			if oneFile.IsDir {
				log.Info("zip dir name:" + oneFile.FileName)
			} else {
				log.Info("zip file name:" + oneFile.FileName)
				fDest, err := w.Create(oneFile.FileName)
				if err != nil {
					log.Info("create zip entry error, download file failed: %s\n", err.Error())
					ctx.ServerError("download file failed:", err)
					return
				}
				body, err := storage.ObsDownloadAFile(setting.Bucket, path+oneFile.FileName)
				if err != nil {
					log.Info("download file failed: %s\n", err.Error())
					ctx.ServerError("download file failed:", err)
					return
				} else {
					defer body.Close()
					p := make([]byte, 1024)
					var readErr error
					var readCount int
					// 读取对象内容
					for {
						readCount, readErr = body.Read(p)
						if readCount > 0 {
							fDest.Write(p[:readCount])
						}
						if readErr != nil {
							break
						}
					}
				}
			}
		}
	} else {
		log.Info("error,msg=" + err.Error())
		ctx.ServerError("no file to download.", err)
	}
}

func QueryTrainJobVersionList(ctx *context.Context) {
	log.Info("query train job version list. start.")
	JobID := ctx.Query("JobID")

	VersionListTasks, count, err := models.QueryModelTrainJobVersionList(JobID)

	log.Info("query return count=" + fmt.Sprint(count))

	if err != nil {
		ctx.ServerError("QueryTrainJobList:", err)
	} else {
		ctx.JSON(200, VersionListTasks)
	}
}

func QueryTrainJobList(ctx *context.Context) {
	log.Info("query train job list. start.")
	repoId := ctx.QueryInt64("repoId")

	VersionListTasks, count, err := models.QueryModelTrainJobList(repoId)
	log.Info("query return count=" + fmt.Sprint(count))

	if err != nil {
		ctx.ServerError("QueryTrainJobList:", err)
	} else {
		ctx.JSON(200, VersionListTasks)
	}

}

func DownloadSingleModelFile(ctx *context.Context) {
	log.Info("DownloadSingleModelFile start.")
	id := ctx.Params(":ID")
	parentDir := ctx.Query("parentDir")
	fileName := ctx.Query("fileName")
	path := Model_prefix + models.AttachmentRelativePath(id) + "/" + parentDir + fileName
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if !isOper(ctx, task.UserId) {
		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
		return
	}
	if setting.PROXYURL != "" {
		body, err := storage.ObsDownloadAFile(setting.Bucket, path)
		if err != nil {
			log.Info("download error.")
		} else {
			//count++
			models.ModifyModelDownloadCount(id)
			defer body.Close()
			ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+fileName)
			ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
			p := make([]byte, 1024)
			var readErr error
			var readCount int
			// 读取对象内容
			for {
				readCount, readErr = body.Read(p)
				if readCount > 0 {
					ctx.Resp.Write(p[:readCount])
					//fmt.Printf("%s", p[:readCount])
				}
				if readErr != nil {
					break
				}
			}
		}
	} else {
		url, err := storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
		if err != nil {
			log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
			ctx.ServerError("GetObsCreateSignedUrl", err)
			return
		}
		//count++
		models.ModifyModelDownloadCount(id)
		http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusMovedPermanently)
	}
}

func ShowModelInfo(ctx *context.Context) {
	ctx.Data["ID"] = ctx.Query("ID")
	ctx.Data["name"] = ctx.Query("name")
	ctx.Data["isModelManage"] = true
	ctx.HTML(200, tplModelInfo)
}

func ShowSingleModel(ctx *context.Context) {
	name := ctx.Query("name")

	log.Info("Show single ModelInfo start.name=" + name)
	models := models.QueryModelByName(name, ctx.Repo.Repository.ID)

	userIds := make([]int64, len(models))
	for i, model := range models {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model)
		userIds[i] = model.UserId
	}
	userNameMap := queryUserName(userIds)

	for _, model := range models {
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}
	ctx.JSON(http.StatusOK, models)
}

func queryUserName(intSlice []int64) map[int64]*models.User {
	keys := make(map[int64]string)
	uniqueElements := []int64{}
	for _, entry := range intSlice {
		if _, value := keys[entry]; !value {
			keys[entry] = ""
			uniqueElements = append(uniqueElements, entry)
		}
	}
	result := make(map[int64]*models.User)
	userLists, err := models.GetUsersByIDs(uniqueElements)
	if err == nil {
		for _, user := range userLists {
			result[user.ID] = user
		}
	}

	return result
}

func ShowOneVersionOtherModel(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	name := ctx.Query("name")
	aimodels := models.QueryModelByName(name, repoId)

	userIds := make([]int64, len(aimodels))
	for i, model := range aimodels {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model)
		userIds[i] = model.UserId
	}
	userNameMap := queryUserName(userIds)

	for _, model := range aimodels {
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}

	if len(aimodels) > 0 {
		ctx.JSON(200, aimodels[1:])
	} else {
		ctx.JSON(200, aimodels)
	}
}

func SetModelCount(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	Type := -1
	_, count, _ := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     1,
			PageSize: 2,
		},
		RepoID: repoId,
		Type:   Type,
		New:    MODEL_LATEST,
	})
	ctx.Data["MODEL_COUNT"] = count
}

func ShowModelTemplate(ctx *context.Context) {
	ctx.Data["isModelManage"] = true
	repoId := ctx.Repo.Repository.ID
	SetModelCount(ctx)

	_, trainCount, _ := models.QueryModelTrainJobList(repoId)
	log.Info("query train count=" + fmt.Sprint(trainCount))
	ctx.Data["TRAIN_COUNT"] = trainCount
	ctx.HTML(200, tplModelManageIndex)
}

func isQueryRight(ctx *context.Context) bool {
	if ctx.Repo.Repository.IsPrivate {
		if ctx.Repo.CanRead(models.UnitTypeModelManage) || ctx.User.IsAdmin || ctx.Repo.IsAdmin() || ctx.Repo.IsOwner() {
			return true
		}
		return false
	} else {
		return true
	}
}

func isOper(ctx *context.Context, modelUserId int64) bool {
	if ctx.User == nil {
		return false
	}
	if ctx.User.IsAdmin || ctx.User.ID == modelUserId {
		return true
	}
	return false
}

func ShowModelPageInfo(ctx *context.Context) {
	log.Info("ShowModelInfo start.")
	if !isQueryRight(ctx) {
		ctx.ServerError("no right.", errors.New(ctx.Tr("repo.model_noright")))
		return
	}
	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	pageSize := ctx.QueryInt("pageSize")
	if pageSize <= 0 {
		pageSize = setting.UI.IssuePagingNum
	}
	repoId := ctx.Repo.Repository.ID
	Type := -1
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: pageSize,
		},
		RepoID: repoId,
		Type:   Type,
		New:    MODEL_LATEST,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}

	userIds := make([]int64, len(modelResult))
	for i, model := range modelResult {
		model.IsCanOper = isOper(ctx, model.UserId)
		userIds[i] = model.UserId
	}

	userNameMap := queryUserName(userIds)

	for _, model := range modelResult {
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}

	mapInterface := make(map[string]interface{})
	mapInterface["data"] = modelResult
	mapInterface["count"] = count
	ctx.JSON(http.StatusOK, mapInterface)
}

func ModifyModel(id string, description string) error {
	err := models.ModifyModelDescription(id, description)
	if err == nil {
		log.Info("modify success.")
	} else {
		log.Info("Failed to modify.id=" + id + "  desc=" + description + " error:" + err.Error())
	}
	return err
}

func ModifyModelInfo(ctx *context.Context) {
	log.Info("modify model start.")
	id := ctx.Query("ID")
	description := ctx.Query("Description")

	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if !isOper(ctx, task.UserId) {
		ctx.ServerError("no right.", errors.New(ctx.Tr("repo.model_noright")))
		return
	}

	err = ModifyModel(id, description)

	if err != nil {
		log.Info("modify error," + err.Error())
		ctx.ServerError("error.", err)
	} else {
		ctx.JSON(200, "success")
	}

}

func QueryModelListForPredict(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     -1,
			PageSize: -1,
		},
		RepoID: repoId,
		Type:   -1,
		New:    -1,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}
	log.Info("query return count=" + fmt.Sprint(count))

	nameList := make([]string, 0)

	nameMap := make(map[string][]*models.AiModelManage)
	for _, model := range modelResult {
		if _, value := nameMap[model.Name]; !value {
			models := make([]*models.AiModelManage, 0)
			models = append(models, model)
			nameMap[model.Name] = models
			nameList = append(nameList, model.Name)
		} else {
			nameMap[model.Name] = append(nameMap[model.Name], model)
		}
	}

	mapInterface := make(map[string]interface{})
	mapInterface["nameList"] = nameList
	mapInterface["nameMap"] = nameMap
	ctx.JSON(http.StatusOK, mapInterface)
}

func QueryModelFileForPredict(ctx *context.Context) {
	id := ctx.Query("ID")
	model, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	prefix := model.Path[len(setting.Bucket)+1:]
	fileinfos, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, prefix)
	ctx.JSON(http.StatusOK, fileinfos)
}

func QueryOneLevelModelFile(ctx *context.Context) {
	id := ctx.Query("ID")
	parentDir := ctx.Query("parentDir")
	model, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	prefix := model.Path[len(setting.Bucket)+1:]
	fileinfos, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, prefix, parentDir)
	ctx.JSON(http.StatusOK, fileinfos)
}