aiforge/routers/repo/ai_model_convert.go

package repo

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

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

const (
	tplModelManageConvertIndex = "repo/modelmanage/convertIndex"
	tplModelConvertInfo        = "repo/modelmanage/convertshowinfo"
	PYTORCH_ENGINE             = 0
	TENSORFLOW_ENGINE          = 1
	MINDSPORE_ENGIN            = 2
	ModelMountPath             = "/model"
	CodeMountPath              = "/code"
	DataSetMountPath           = "/dataset"
	LogFile                    = "log.txt"
	DefaultBranchName          = "master"
	SubTaskName                = "task1"
	GpuQueue                   = "openidgx"
	Success                    = "S000"
	GPU_PYTORCH_IMAGE          = "dockerhub.pcl.ac.cn:5000/user-images/openi:tensorRT_7_zouap"
	GPU_TENSORFLOW_IMAGE       = "dockerhub.pcl.ac.cn:5000/user-images/openi:tf2onnx"

	PytorchBootFile       = "convert_pytorch.py"
	MindsporeBootFile     = "convert_mindspore.py"
	TensorFlowNpuBootFile = "convert_tensorflow.py"
	TensorFlowGpuBootFile = "convert_tensorflow_gpu.py"

	ConvertRepoPath = "https://git.openi.org.cn/zouap/npu_test"

	REPO_ID = 33267

	NPU_MINDSPORE_IMAGE_ID  = 122
	NPU_TENSORFLOW_IMAGE_ID = 121

	NPU_FlavorCode = "modelarts.bm.910.arm.public.1"
	NPU_PoolID     = "pool7908321a"
)

var (
	TrainResourceSpecs *models.ResourceSpecs
)

func SaveModelConvert(ctx *context.Context) {
	log.Info("save model convert start.")
	if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
		ctx.JSON(403, ctx.Tr("repo.model_noright"))
		return
	}
	name := ctx.Query("name")
	desc := ctx.Query("desc")
	modelId := ctx.Query("modelId")
	modelPath := ctx.Query("ModelFile")
	SrcEngine := ctx.QueryInt("SrcEngine")
	InputShape := ctx.Query("inputshape")
	InputDataFormat := ctx.Query("inputdataformat")
	DestFormat := ctx.QueryInt("DestFormat")
	NetOutputFormat := ctx.QueryInt("NetOutputFormat")

	task, err := models.QueryModelById(modelId)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}

	uuid := uuid.NewV4()
	id := uuid.String()
	modelConvert := &models.AiModelConvert{
		ID:              id,
		Name:            name,
		Description:     desc,
		Status:          string(models.JobWaiting),
		SrcEngine:       SrcEngine,
		RepoId:          ctx.Repo.Repository.ID,
		ModelName:       task.Name,
		ModelVersion:    task.Version,
		ModelId:         modelId,
		ModelPath:       modelPath,
		DestFormat:      DestFormat,
		NetOutputFormat: NetOutputFormat,
		InputShape:      InputShape,
		InputDataFormat: InputDataFormat,
		UserId:          ctx.User.ID,
	}
	models.SaveModelConvert(modelConvert)
	go goCreateTask(modelConvert, ctx, task)

	ctx.JSON(200, map[string]string{
		"result_code": "0",
	})
}

func goCreateTask(modelConvert *models.AiModelConvert, ctx *context.Context, task *models.AiModelManage) error {
	if modelConvert.IsGpuTrainTask() {
		log.Info("create gpu train job.")
		return createGpuTrainJob(modelConvert, ctx, task)
	} else {
		//create npu job
		log.Info("create npu train job.")
		return createNpuTrainJob(modelConvert, ctx, task.Path)
	}
}

func createNpuTrainJob(modelConvert *models.AiModelConvert, ctx *context.Context, modelRelativePath string) error {
	VersionOutputPath := "V0001"
	codeLocalPath := setting.JobPath + modelConvert.ID + modelarts.CodePath
	codeObsPath := "/" + setting.Bucket + modelarts.JobPath + modelConvert.ID + modelarts.CodePath
	outputObsPath := "/" + setting.Bucket + modelarts.JobPath + modelConvert.ID + modelarts.OutputPath + VersionOutputPath + "/"
	logObsPath := "/" + setting.Bucket + modelarts.JobPath + modelConvert.ID + modelarts.LogPath + VersionOutputPath + "/"
	dataPath := "/" + modelRelativePath

	_, err := ioutil.ReadDir(codeLocalPath)
	if err == nil {
		os.RemoveAll(codeLocalPath)
	}
	if err := downloadConvertCode(ConvertRepoPath, codeLocalPath, DefaultBranchName); err != nil {
		log.Error("downloadCode failed, server timed out: %s (%v)", ConvertRepoPath, err)
		return err
	}
	if err := obsMkdir(setting.CodePathPrefix + modelConvert.ID + modelarts.OutputPath + VersionOutputPath + "/"); err != nil {
		log.Error("Failed to obsMkdir_output: %s (%v)", modelConvert.ID+modelarts.OutputPath, err)
		return err
	}
	if err := obsMkdir(setting.CodePathPrefix + modelConvert.ID + modelarts.LogPath + VersionOutputPath + "/"); err != nil {
		log.Error("Failed to obsMkdir_log: %s (%v)", modelConvert.ID+modelarts.LogPath, err)
		return err
	}
	if err := uploadCodeToObs(codeLocalPath, modelConvert.ID, ""); err != nil {
		log.Error("Failed to uploadCodeToObs: %s (%v)", modelConvert.ID, err)
		return err
	}
	intputshape := strings.Split(modelConvert.InputShape, ",")
	n := "256"
	c := "1"
	h := "28"
	w := "28"
	if len(intputshape) == 4 {
		n = intputshape[0]
		c = intputshape[1]
		h = intputshape[2]
		w = intputshape[3]
	}
	param := make([]models.Parameter, 0)
	modelPara := models.Parameter{
		Label: "model",
		Value: modelConvert.ModelPath,
	}
	param = append(param, modelPara)

	batchSizePara := models.Parameter{
		Label: "n",
		Value: fmt.Sprint(n),
	}
	param = append(param, batchSizePara)
	channelSizePara := models.Parameter{
		Label: "c",
		Value: fmt.Sprint(c),
	}
	param = append(param, channelSizePara)
	heightPara := models.Parameter{
		Label: "h",
		Value: fmt.Sprint(h),
	}
	param = append(param, heightPara)
	widthPara := models.Parameter{
		Label: "w",
		Value: fmt.Sprint(w),
	}
	param = append(param, widthPara)

	var engineId int64
	engineId = int64(NPU_MINDSPORE_IMAGE_ID)
	bootfile := MindsporeBootFile
	if modelConvert.SrcEngine == TENSORFLOW_ENGINE {
		engineId = int64(NPU_TENSORFLOW_IMAGE_ID)
		bootfile = TensorFlowNpuBootFile
	}
	req := &modelarts.GenerateTrainJobReq{
		JobName:          modelConvert.ID,
		DisplayJobName:   modelConvert.Name,
		DataUrl:          dataPath,
		Description:      modelConvert.Description,
		CodeObsPath:      codeObsPath,
		BootFileUrl:      codeObsPath + bootfile,
		BootFile:         bootfile,
		TrainUrl:         outputObsPath,
		FlavorCode:       NPU_FlavorCode,
		WorkServerNumber: 1,
		IsLatestVersion:  modelarts.IsLatestVersion,
		EngineID:         engineId,
		LogUrl:           logObsPath,
		PoolID:           NPU_PoolID,
		Parameters:       param,
		BranchName:       DefaultBranchName,
	}
	result, err := modelarts.GenerateModelConvertTrainJob(req)
	if err == nil {
		log.Info("jobId=" + fmt.Sprint(result.JobID) + " versionid=" + fmt.Sprint(result.VersionID))
		models.UpdateModelConvertModelArts(modelConvert.ID, fmt.Sprint(result.JobID), fmt.Sprint(result.VersionID))
	}
	return err
}

func downloadConvertCode(repopath string, codePath, branchName string) error {
	//add "file:///" prefix to make the depth valid
	if err := git.Clone(repopath, codePath, git.CloneRepoOptions{Branch: branchName, Depth: 1}); err != nil {
		log.Error("Failed to clone repository: %s (%v)", repopath, err)
		return err
	}
	log.Info("srcPath=" + repopath + " codePath=" + codePath)
	configFile, err := os.OpenFile(codePath+"/.git/config", os.O_RDWR, 0666)
	if err != nil {
		log.Error("open file(%s) failed:%v", codePath+"/,git/config", err)
		return err
	}

	defer configFile.Close()

	pos := int64(0)
	reader := bufio.NewReader(configFile)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				log.Error("not find the remote-url")
				return nil
			} else {
				log.Error("read error: %v", err)
				return err
			}
		}

		if strings.Contains(line, "url") && strings.Contains(line, ".git") {
			originUrl := "\turl = " + repopath + "\n"
			if len(line) > len(originUrl) {
				originUrl += strings.Repeat(" ", len(line)-len(originUrl))
			}
			bytes := []byte(originUrl)
			_, err := configFile.WriteAt(bytes, pos)
			if err != nil {
				log.Error("WriteAt failed:%v", err)
				return err
			}
			break
		}

		pos += int64(len(line))
	}

	return nil
}

func downloadFromObsToLocal(task *models.AiModelManage, localPath string) error {
	path := Model_prefix + models.AttachmentRelativePath(task.ID) + "/"
	allFile, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, path)
	if err == nil {
		_, errState := os.Stat(localPath)
		if errState != nil {
			if err = os.MkdirAll(localPath, os.ModePerm); err != nil {
				return err
			}
		}
		for _, oneFile := range allFile {
			if oneFile.IsDir {
				log.Info(" dir name:" + oneFile.FileName)
			} else {
				allFileName := localPath + "/" + oneFile.FileName
				index := strings.LastIndex(allFileName, "/")
				if index != -1 {
					parentDir := allFileName[0:index]
					if err = os.MkdirAll(parentDir, os.ModePerm); err != nil {
						log.Info("make dir may be error," + err.Error())
					}
				}
				fDest, err := os.Create(allFileName)
				if err != nil {
					log.Info("create file error, download file failed: %s\n", err.Error())
					return err
				}
				body, err := storage.ObsDownloadAFile(setting.Bucket, path+oneFile.FileName)
				if err != nil {
					log.Info("download file failed: %s\n", err.Error())
					return err
				} 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())
		return err
	}
	return nil
}

func createGpuTrainJob(modelConvert *models.AiModelConvert, ctx *context.Context, model *models.AiModelManage) error {
	modelRelativePath := model.Path
	command := ""
	IMAGE_URL := GPU_PYTORCH_IMAGE
	dataActualPath := setting.Attachment.Minio.RealPath + modelRelativePath

	if modelConvert.SrcEngine == PYTORCH_ENGINE {
		command = getGpuModelConvertCommand(modelConvert.ID, modelConvert.ModelPath, modelConvert, PytorchBootFile)
	} else if modelConvert.SrcEngine == TENSORFLOW_ENGINE {
		IMAGE_URL = GPU_TENSORFLOW_IMAGE
		command = getGpuModelConvertCommand(modelConvert.ID, modelConvert.ModelPath, modelConvert, TensorFlowGpuBootFile)
		//如果模型在OBS上，需要下载到本地，并上传到minio中
		if model.Type == models.TypeCloudBrainTwo {
			relatetiveModelPath := setting.JobPath + modelConvert.ID + "/dataset"
			log.Info("local dataset path:" + relatetiveModelPath)
			downloadFromObsToLocal(model, relatetiveModelPath)
			uploadCodeToMinio(relatetiveModelPath+"/", modelConvert.ID, "/dataset/")
			dataActualPath = setting.Attachment.Minio.RealPath + setting.Attachment.Minio.Bucket + "/" + setting.CBCodePathPrefix + modelConvert.ID + "/dataset"
		}

	}
	log.Info("dataActualPath=" + dataActualPath)

	log.Info("command=" + command)

	codePath := setting.JobPath + modelConvert.ID + CodeMountPath
	downloadConvertCode(ConvertRepoPath, codePath, DefaultBranchName)

	uploadCodeToMinio(codePath+"/", modelConvert.ID, CodeMountPath+"/")

	minioCodePath := setting.Attachment.Minio.RealPath + setting.Attachment.Minio.Bucket + "/" + setting.CBCodePathPrefix + modelConvert.ID + "/code"
	log.Info("minio codePath=" + minioCodePath)

	modelPath := setting.JobPath + modelConvert.ID + ModelMountPath + "/"
	log.Info("local modelPath=" + modelPath)
	mkModelPath(modelPath)

	uploadCodeToMinio(modelPath, modelConvert.ID, ModelMountPath+"/")
	minioModelPath := setting.Attachment.Minio.RealPath + setting.Attachment.Minio.Bucket + "/" + setting.CBCodePathPrefix + modelConvert.ID + "/model"
	log.Info("minio model path=" + minioModelPath)

	if TrainResourceSpecs == nil {
		json.Unmarshal([]byte(setting.TrainResourceSpecs), &TrainResourceSpecs)
	}
	resourceSpec := TrainResourceSpecs.ResourceSpec[1]
	jobResult, err := cloudbrain.CreateJob(modelConvert.ID, models.CreateJobParams{
		JobName:    modelConvert.ID,
		RetryCount: 1,
		GpuType:    GpuQueue,
		Image:      IMAGE_URL,
		TaskRoles: []models.TaskRole{
			{
				Name:                  SubTaskName,
				TaskNumber:            1,
				MinSucceededTaskCount: 1,
				MinFailedTaskCount:    1,
				CPUNumber:             resourceSpec.CpuNum,
				GPUNumber:             resourceSpec.GpuNum,
				MemoryMB:              resourceSpec.MemMiB,
				ShmMB:                 resourceSpec.ShareMemMiB,
				Command:               command,
				NeedIBDevice:          false,
				IsMainRole:            false,
				UseNNI:                false,
			},
		},
		Volumes: []models.Volume{
			{
				HostPath: models.StHostPath{
					Path:      minioCodePath,
					MountPath: CodeMountPath,
					ReadOnly:  false,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      dataActualPath,
					MountPath: DataSetMountPath,
					ReadOnly:  true,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      minioModelPath,
					MountPath: ModelMountPath,
					ReadOnly:  false,
				},
			},
		},
	})
	if err != nil {
		log.Error("CreateJob failed:", err.Error(), ctx.Data["MsgID"])
		return err
	}
	if jobResult.Code != Success {
		log.Error("CreateJob(%s) failed:%s", modelConvert.ID, jobResult.Msg, ctx.Data["MsgID"])
		return errors.New(jobResult.Msg)
	}

	var jobID = jobResult.Payload["jobId"].(string)
	log.Info("jobId=" + jobID)
	models.UpdateModelConvertCBTI(modelConvert.ID, jobID)

	return nil
}

func getGpuModelConvertCommand(name string, modelFile string, modelConvert *models.AiModelConvert, bootfile string) string {
	var command string

	intputshape := strings.Split(modelConvert.InputShape, ",")
	n := "256"
	c := "1"
	h := "28"
	w := "28"
	if len(intputshape) == 4 {
		n = intputshape[0]
		c = intputshape[1]
		h = intputshape[2]
		w = intputshape[3]
	}
	command += "python3 /code/" + bootfile + " --model " + modelFile + " --n " + n + " --c " + c + " --h " + h + " --w " + w + " > " + ModelMountPath + "/" + name + "-" + LogFile
	return command
}

func DeleteModelConvert(ctx *context.Context) {
	log.Info("delete model convert start.")
	id := ctx.Params(":id")
	err := models.DeleteModelConvertById(id)
	if err != nil {
		ctx.JSON(500, err.Error())
	} else {
		ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/modelmanage/convert_model")
	}
}

func StopModelConvert(ctx *context.Context) {
	id := ctx.Params(":id")
	log.Info("stop model convert start.id=" + id)
	job, err := models.QueryModelConvertById(id)
	if err != nil {
		ctx.ServerError("Not found task.", err)
		return
	}
	if job.IsGpuTrainTask() {
		err = cloudbrain.StopJob(job.CloudBrainTaskId)
		if err != nil {
			log.Error("Stop cloudbrain Job(%s) failed:%v", job.CloudBrainTaskId, err)
		}
	} else {
		_, err = modelarts.StopTrainJob(job.CloudBrainTaskId, job.ModelArtsVersionId)
		if err != nil {
			log.Error("Stop modelarts Job(%s) failed:%v", job.CloudBrainTaskId, err)
		}
	}
	job.Status = string(models.JobStopped)
	if job.EndTime == 0 {
		job.EndTime = timeutil.TimeStampNow()
	}
	models.ModelConvertSetDuration(job)
	err = models.UpdateModelConvert(job)
	if err != nil {
		log.Error("UpdateModelConvert failed:", err)
	}
	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/modelmanage/convert_model")
}

func ShowModelConvertInfo(ctx *context.Context) {
	ctx.Data["ID"] = ctx.Query("ID")
	ctx.Data["isModelManage"] = true
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)

	job, err := models.QueryModelConvertById(ctx.Query("ID"))
	if err == nil {
		ctx.Data["task"] = job
	} else {
		ctx.ServerError("Not found task.", err)
		return
	}
	ctx.Data["canDownload"] = isOper(ctx, job.UserId)
	user, err := models.GetUserByID(job.UserId)
	if err == nil {
		job.UserName = user.Name
		job.UserRelAvatarLink = user.RelAvatarLink()
	}

	if job.IsGpuTrainTask() {
		ctx.Data["npu_display"] = "none"
		ctx.Data["gpu_display"] = "block"
		result, err := cloudbrain.GetJob(job.CloudBrainTaskId)
		if err != nil {
			log.Info("error:" + err.Error())
			ctx.Data["error"] = err.Error()
			return
		}
		if result != nil {
			jobRes, _ := models.ConvertToJobResultPayload(result.Payload)
			ctx.Data["result"] = jobRes
			taskRoles := jobRes.TaskRoles
			taskRes, _ := models.ConvertToTaskPod(taskRoles[cloudbrain.SubTaskName].(map[string]interface{}))
			ctx.Data["taskRes"] = taskRes
			ctx.Data["ExitDiagnostics"] = taskRes.TaskStatuses[0].ExitDiagnostics
			ctx.Data["AppExitDiagnostics"] = jobRes.JobStatus.AppExitDiagnostics

			job.Status = jobRes.JobStatus.State

			if jobRes.JobStatus.State != string(models.JobWaiting) && jobRes.JobStatus.State != string(models.JobFailed) {
				job.ContainerIp = taskRes.TaskStatuses[0].ContainerIP
				job.ContainerID = taskRes.TaskStatuses[0].ContainerID
				job.Status = taskRes.TaskStatuses[0].State
			}
			if jobRes.JobStatus.State != string(models.JobWaiting) {
				models.ModelComputeAndSetDuration(job, jobRes)
				err = models.UpdateModelConvert(job)
				if err != nil {
					log.Error("UpdateModelConvert failed:", err)
				}
			}
		}
	} else {
		ctx.Data["npu_display"] = "block"
		ctx.Data["gpu_display"] = "none"
		ctx.Data["ExitDiagnostics"] = ""
		ctx.Data["AppExitDiagnostics"] = ""

	}

	ctx.HTML(200, tplModelConvertInfo)
}

func ConvertModelTemplate(ctx *context.Context) {
	ctx.Data["isModelManage"] = true
	ctx.Data["MODEL_COUNT"] = 0
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)
	ctx.Data["TRAIN_COUNT"] = 0
	ShowModelConvertPageInfo(ctx)
	ctx.HTML(200, tplModelManageConvertIndex)
}

func ShowModelConvertPageInfo(ctx *context.Context) {
	log.Info("ShowModelConvertInfo start.")
	if !isQueryRight(ctx) {
		log.Info("no right.")
		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
		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
	modelResult, count, err := models.QueryModelConvert(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: pageSize,
		},
		RepoID: repoId,
	})
	if err != nil {
		log.Info("query db error." + err.Error())
		ctx.ServerError("Cloudbrain", err)
		return
	}
	userIds := make([]int64, len(modelResult))
	for i, model := range modelResult {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(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()
		}
	}
	pager := context.NewPagination(int(count), page, pageSize, 5)
	ctx.Data["Page"] = pager
	ctx.Data["Tasks"] = modelResult

}

func ModelConvertDownloadModel(ctx *context.Context) {
	log.Info("enter here......")
	id := ctx.Params(":id")
	job, err := models.QueryModelConvertById(id)
	if err != nil {
		ctx.ServerError("Not found task.", err)
		return
	}
	AllDownload := ctx.QueryBool("AllDownload")
	if AllDownload {
		if job.IsGpuTrainTask() {
			path := setting.CBCodePathPrefix + job.ID + "/model/"
			allFile, err := storage.GetAllObjectByBucketAndPrefixMinio(setting.Attachment.Minio.Bucket, path)
			if err == nil {
				returnFileName := job.Name + ".zip"
				MinioDownloadManyFile(path, ctx, returnFileName, allFile)
			} else {
				log.Info("error,msg=" + err.Error())
				ctx.ServerError("no file to download.", err)
			}
		} else {

		}
	} else {
		if job.IsGpuTrainTask() {
			parentDir := ctx.Query("parentDir")
			fileName := ctx.Query("fileName")
			jobName := ctx.Query("jobName")
			filePath := "jobs/" + jobName + "/model/" + parentDir
			url, err := storage.Attachments.PresignedGetURL(filePath, fileName)
			if err != nil {
				log.Error("PresignedGetURL failed: %v", err.Error(), ctx.Data["msgID"])
				ctx.ServerError("PresignedGetURL", err)
				return
			}
			//ctx.JSON(200, url)
			http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusTemporaryRedirect)
		} else {

		}
	}

}