aiforge/routers/api/v1/repo/modelarts.go

// Copyright 2016 The Gogs Authors. All rights reserved.
// Copyright 2018 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package repo

import (
	"net/http"
	"path"
	"strconv"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/modelarts"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/storage"
)

func GetModelArtsNotebook(ctx *context.APIContext) {
	var (
		err error
	)

	jobID := ctx.Params(":jobid")
	repoID := ctx.Repo.Repository.ID
	job, err := models.GetRepoCloudBrainByJobID(repoID, jobID)
	if err != nil {
		ctx.NotFound(err)
		return
	}
	result, err := modelarts.GetJob(jobID)
	if err != nil {
		ctx.NotFound(err)
		return
	}

	job.Status = result.Status
	err = models.UpdateJob(job)
	if err != nil {
		log.Error("UpdateJob failed:", err)
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":     jobID,
		"JobStatus": result.Status,
	})

}

func GetModelArtsTrainJob(ctx *context.APIContext) {
	var (
		err error
	)

	jobID := ctx.Params(":jobid")
	repoID := ctx.Repo.Repository.ID
	job, err := models.GetRepoCloudBrainByJobID(repoID, jobID)
	if err != nil {
		ctx.NotFound(err)
		return
	}
	result, err := modelarts.GetTrainJob(jobID, strconv.FormatInt(job.VersionID, 10))
	if err != nil {
		ctx.NotFound(err)
		return
	}

	job.Status = modelarts.TransTrainJobStatus(result.IntStatus)
	job.Duration = result.Duration
	job.TrainJobDuration = result.TrainJobDuration
	err = models.UpdateJob(job)
	if err != nil {
		log.Error("UpdateJob failed:", err)
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"JobStatus":   job.Status,
		"JobDuration": job.Duration,
	})

}

func GetModelArtsTrainJobVersion(ctx *context.APIContext) {
	var (
		err error
	)

	jobID := ctx.Params(":jobid")
	versionName := ctx.Query("version_name")
	job, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		ctx.NotFound(err)
		return
	}
	result, err := modelarts.GetTrainJob(jobID, strconv.FormatInt(job.VersionID, 10))
	if err != nil {
		ctx.NotFound(err)
		return
	}

	job.Status = modelarts.TransTrainJobStatus(result.IntStatus)
	job.Duration = result.Duration
	job.TrainJobDuration = result.TrainJobDuration

	if result.Duration != 0 {
		job.TrainJobDuration = addZero(result.Duration/3600000) + ":" + addZero(result.Duration%3600000/60000) + ":" + addZero(result.Duration%60000/1000)

	} else {
		job.TrainJobDuration = "00:00:00"
	}

	err = models.UpdateTrainJobVersion(job)
	if err != nil {
		log.Error("UpdateJob failed:", err)
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"JobStatus":   job.Status,
		"JobDuration": job.TrainJobDuration,
	})

}

func addZero(t int64) (m string) {
	if t < 10 {
		m = "0" + strconv.FormatInt(t, 10)
		return m
	} else {
		return strconv.FormatInt(t, 10)
	}
}

func TrainJobGetLog(ctx *context.APIContext) {
	var (
		err error
	)

	var jobID = ctx.Params(":jobid")
	var versionName = ctx.Query("version_name")
	// var logFileName = ctx.Query("file_name")
	var baseLine = ctx.Query("base_line")
	var order = ctx.Query("order")
	var lines = ctx.Query("lines")
	lines_int, err := strconv.Atoi(lines)
	if err != nil {
		log.Error("change lines(%d) string to int failed", lines_int)
	}

	if order != modelarts.OrderDesc && order != modelarts.OrderAsc {
		log.Error("order(%s) check failed", order)
		ctx.JSON(http.StatusBadRequest, map[string]interface{}{
			"err_msg": "order check failed",
		})
		return
	}

	resultLogFile, result, err := trainJobGetLogContent(jobID, versionName, baseLine, order, lines_int)
	if err != nil {
		log.Error("trainJobGetLog(%s) failed:%v", jobID, err.Error())
		// ctx.RenderWithErr(err.Error(), tplModelArtsTrainJobShow, nil)
		return
	}

	ctx.Data["log_file_name"] = resultLogFile.LogFileList[0]

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"LogFileName": resultLogFile.LogFileList[0],
		"StartLine":   result.StartLine,
		"EndLine":     result.EndLine,
		"Content":     result.Content,
		"Lines":       result.Lines,
	})
}

func trainJobGetLogContent(jobID string, versionName string, baseLine string, order string, lines int) (*models.GetTrainJobLogFileNamesResult, *models.GetTrainJobLogResult, error) {
	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", jobID, err.Error())
		return nil, nil, err
	}

	resultLogFile, err := modelarts.GetTrainJobLogFileNames(jobID, strconv.FormatInt(task.VersionID, 10))
	if err != nil {
		log.Error("GetTrainJobLogFileNames(%s) failed:%v", jobID, err.Error())
		return nil, nil, err
	}

	result, err := modelarts.GetTrainJobLog(jobID, strconv.FormatInt(task.VersionID, 10), baseLine, resultLogFile.LogFileList[0], order, lines)
	if err != nil {
		log.Error("GetTrainJobLog(%s) failed:%v", jobID, err.Error())
		return nil, nil, err
	}

	return resultLogFile, result, err
}

func DelTrainJobVersion(ctx *context.APIContext) {
	var (
		err error
	)

	var jobID = ctx.Params(":jobid")
	var versionName = ctx.Query("version_name")
	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
		ctx.NotFound(err)
		return
	}

	//删除modelarts上的记录
	_, err = modelarts.DelTrainJobVersion(jobID, strconv.FormatInt(task.VersionID, 10))
	if err != nil {
		log.Error("DelTrainJobVersion(%s) failed:%v", task.JobName, err.Error())
		ctx.NotFound(err)
		return
	}

	//删除数据库记录
	err = models.DeleteJob(task)
	if err != nil {
		ctx.ServerError("DeleteJob failed", err)
		ctx.NotFound(err)
		return
	}

	//获取删除后的版本数量
	repo := ctx.Repo.Repository
	VersionTaskList, VersionListCount, err := models.CloudbrainsVersionList(&models.CloudbrainsOptions{
		RepoID:  repo.ID,
		Type:    models.TypeCloudBrainTwo,
		JobType: string(models.JobTypeTrain),
		JobID:   jobID,
	})
	if err != nil {
		ctx.ServerError("get VersionListCount faild", err)
		return
	}

	// 判断当前删掉的任务是否是最新版本，若是，将排序后的TotalVersionCount置为删掉的最新版本的TotalVersionCount，若不是，按时间排序后的版本列表的第一个版本设置为最新版本，TotalVersionCount不变
	if task.IsLatestVersion == modelarts.IsLatestVersion {
		err = models.SetVersionCountAndLatestVersion(jobID, VersionTaskList[0].Cloudbrain.VersionName, VersionListCount, modelarts.IsLatestVersion, task.TotalVersionCount)
		if err != nil {
			ctx.ServerError("UpdateJobVersionCount failed", err)
			return
		}
	} else {
		err = models.SetVersionCountAndLatestVersion(jobID, VersionTaskList[0].VersionName, VersionListCount, modelarts.IsLatestVersion, VersionTaskList[0].Cloudbrain.TotalVersionCount)
		if err != nil {
			ctx.ServerError("UpdateJobVersionCount failed", err)
			return
		}
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"VersionName": versionName,
		"StatusOK":    0,
	})
}

func StopTrainJobVersion(ctx *context.APIContext) {
	var (
		err error
	)
	var jobID = ctx.Params(":jobid")
	var versionName = ctx.Query("version_name")
	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
		return
	}

	_, err = modelarts.StopTrainJob(jobID, strconv.FormatInt(task.VersionID, 10))
	if err != nil {
		log.Error("StopTrainJob(%s) failed:%v", task.JobName, err.Error())
		return
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"VersionName": versionName,
		"StatusOK":    0,
	})
}

func ModelList(ctx *context.APIContext) {
	var (
		err error
	)

	var jobID = ctx.Params(":jobid")
	var versionName = ctx.Query("version_name")
	parentDir := ctx.Query("parentDir")
	dirArray := strings.Split(parentDir, "/")
	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
		return
	}
	models, err := storage.GetObsListObject(task.JobName, parentDir, versionName)
	if err != nil {
		log.Info("get TrainJobListModel failed:", err)
		ctx.ServerError("GetObsListObject:", err)
		return
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":            jobID,
		"VersionName":      versionName,
		"StatusOK":         0,
		"Path":             dirArray,
		"Dirs":             models,
		"task":             task,
		"PageIsCloudBrain": true,
	})
}

func ModelDownload(ctx *context.Context) {
	var (
		err error
	)

	var jobID = ctx.Params(":jobid")
	versionName := ctx.Query("version_name")
	// versionName := "V0001"
	parentDir := ctx.Query("parent_dir")
	fileName := ctx.Query("file_name")
	log.Info("DownloadSingleModelFile start.")
	// id := ctx.Params(":ID")
	// path := Model_prefix + models.AttachmentRelativePath(id) + "/" + parentDir + fileName
	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
		return
	}

	path := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, task.JobName, setting.OutPutPath, versionName, parentDir, fileName), "/")
	log.Info("Download path is:%s", path)
	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)
	}
}