From f9f9286d09b2ced5faf03dff6bedfe1b9189e55d Mon Sep 17 00:00:00 2001 From: zouap Date: Thu, 4 Nov 2021 09:56:16 +0800 Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E4=BB=A3=E7=A0=81=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: zouap --- models/ai_model_manage.go | 35 ++++++- routers/repo/ai_model_manage.go | 161 ++++++++++++++++++++++++++++++-- routers/repo/cloudbrain.go | 12 +-- 3 files changed, 189 insertions(+), 19 deletions(-) diff --git a/models/ai_model_manage.go b/models/ai_model_manage.go index 60c45fe8c..b6169b806 100644 --- a/models/ai_model_manage.go +++ b/models/ai_model_manage.go @@ -11,7 +11,7 @@ type AiModelManage struct { ID string `xorm:"pk"` Name string `xorm:"NOT NULL"` Version string `xorm:"NOT NULL"` - Parent int64 `xorm:"NOT NULL"` + Parent string `xorm:"NOT NULL"` Type int `xorm:"NOT NULL"` Size int64 `xorm:"NOT NULL"` Description string `xorm:"varchar(2000)"` @@ -22,10 +22,14 @@ type AiModelManage struct { Engine int `xorm:"NOT NULL DEFAULT 0"` Status int `xorm:"NOT NULL DEFAULT 0"` Accuracy string `xorm:"varchar(1000)"` - DatasetId int64 `xorm:"NULL"` + AttachmentId string `xorm:"NULL"` RepoId int64 `xorm:"NULL"` - CodePath string `xorm:"varchar(400) NULL"` + CodeBranch string `xorm:"varchar(400) NULL"` + CodeCommitID string `xorm:"NULL"` + UserId int64 `xorm:"NOT NULL"` + TrainTaskInfo string `xorm:"text NULL"` CreatedUnix timeutil.TimeStamp `xorm:"created"` + UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"` } func SaveModelToDb(model *AiModelManage) error { @@ -39,3 +43,28 @@ func SaveModelToDb(model *AiModelManage) error { log.Info("success to save 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 QueryModelByName(name string, uid int64) []*AiModelManage { + sess := x.NewSession() + defer sess.Close() + sess.Select("*").Table("ai_model_manage"). + Where("name=" + name + " and user_id=" + fmt.Sprint(uid)) + aiModelManageList := make([]*AiModelManage, 0) + sess.Find(&aiModelManageList) + return aiModelManageList +} diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go index 98455e92f..2e75c5634 100644 --- a/routers/repo/ai_model_manage.go +++ b/routers/repo/ai_model_manage.go @@ -1,36 +1,92 @@ package repo import ( + "archive/zip" + "errors" + "fmt" + "io" + "io/ioutil" + "os" + "path/filepath" + "code.gitea.io/gitea/models" "code.gitea.io/gitea/modules/context" "code.gitea.io/gitea/modules/log" + "code.gitea.io/gitea/modules/setting" uuid "github.com/satori/go.uuid" ) func SaveModel(ctx *context.Context) { log.Info("save model start.") - jobId := ctx.QueryInt64("JobId") + trainTaskId := ctx.QueryInt64("TrainTask") name := ctx.Query("Name") + version := ctx.Query("Version") + label := ctx.Query("Label") + description := ctx.Query("Description") - aiTasks, count, err := models.Cloudbrains(&models.CloudbrainsOptions{ - JobID: jobId, + aiTasks, _, err := models.Cloudbrains(&models.CloudbrainsOptions{ + JobID: trainTaskId, }) if err != nil { - log.Info("query task error.") + log.Info("query task error." + err.Error()) + ctx.Error(500, fmt.Sprintf("query cloud brain train task error. %v", err)) return } - if count > 0 { - for _, task := range aiTasks { - log.Info("find task name:" + task.JobName) + uuid := uuid.NewV4() + id := uuid.String() + modelPath := id + parent := id + var modelSize int64 + cloudType := models.TypeCloudBrainTwo + if len(aiTasks) != 1 { + log.Info("query task error. len=" + fmt.Sprint(len(aiTasks))) + ctx.Error(500, fmt.Sprintf("query cloud brain train task error. %v", err)) + return + } + aiTask := aiTasks[0] + log.Info("find task name:" + aiTask.JobName) + aimodels := models.QueryModelByName(name, ctx.User.ID) + if len(aimodels) > 0 { + for _, model := range aimodels { + if model.ID == model.Parent { + parent = model.ID + } } } + cloudType = aiTask.Cloudbrain.Type + //download model zip + if cloudType == models.TypeCloudBrainOne { + modelPath, modelSize, err = downloadModelFromCloudBrainOne(id, aiTask.JobName, "") + if err != nil { + log.Info("download model from CloudBrainOne faild." + err.Error()) + ctx.Error(500, fmt.Sprintf("%v", err)) + return + } + } else if cloudType == models.TypeCloudBrainTwo { + modelPath, err = downloadModelFromCloudBrainTwo(id) + if err == nil { - id := uuid.NewV4() + } else { + log.Info("download model from CloudBrainTwo faild." + err.Error()) + ctx.Error(500, fmt.Sprintf("%v", err)) + return + } + } model := &models.AiModelManage{ - ID: id.String(), - Name: name, + ID: id, + Version: version, + Label: label, + Name: name, + Description: description, + Parent: parent, + Type: cloudType, + Path: modelPath, + Size: modelSize, + AttachmentId: aiTask.Uuid, + RepoId: aiTask.RepoID, + UserId: ctx.User.ID, } models.SaveModelToDb(model) @@ -38,8 +94,93 @@ func SaveModel(ctx *context.Context) { log.Info("save model end.") } +func downloadModelFromCloudBrainOne(modelUUID string, jobName string, parentDir string) (string, int64, error) { + + modelActualPath := setting.Attachment.Minio.RealPath + + setting.Attachment.Minio.Bucket + "/" + + "aimodels/" + + models.AttachmentRelativePath(modelUUID) + + "/" + os.MkdirAll(modelActualPath, 0755) + zipFile := modelActualPath + "model.zip" + + modelDir := setting.JobPath + jobName + "/model/" + + dir, _ := ioutil.ReadDir(modelDir) + if len(dir) == 0 { + return "", 0, errors.New("cannot create model, as model is empty.") + } + + err := zipDir(modelDir, zipFile) + if err != nil { + return "", 0, err + } + + fi, err := os.Stat(zipFile) + if err == nil { + return modelActualPath, fi.Size(), nil + } else { + return "", 0, err + } +} + +func zipDir(dir, zipFile string) error { + fz, err := os.Create(zipFile) + if err != nil { + log.Info("Create zip file failed: %s\n", err.Error()) + return err + } + defer fz.Close() + + w := zip.NewWriter(fz) + defer w.Close() + + err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { + if !info.IsDir() { + fDest, err := w.Create(path[len(dir)+1:]) + if err != nil { + log.Info("Create failed: %s\n", err.Error()) + return err + } + fSrc, err := os.Open(path) + if err != nil { + log.Info("Open failed: %s\n", err.Error()) + return err + } + defer fSrc.Close() + _, err = io.Copy(fDest, fSrc) + if err != nil { + log.Info("Copy failed: %s\n", err.Error()) + return err + } + } + return nil + }) + if err != nil { + return err + } + return nil +} + +func downloadModelFromCloudBrainTwo(modelUUID string) (string, error) { + dataActualPath := setting.Bucket + "/" + + "aimodels/" + + models.AttachmentRelativePath(modelUUID) + + "/" + return dataActualPath, nil +} + func DeleteModel(ctx *context.Context) { log.Info("delete model start.") + id := ctx.Query("ID") + err := models.DeleteModelById(id) + if err != nil { + ctx.JSON(500, err.Error()) + } else { + ctx.JSON(200, map[string]string{ + "result_code": "0", + }) + } } func DownloadModel(ctx *context.Context) { diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index 03fba6cd1..b70ec5fda 100755 --- a/routers/repo/cloudbrain.go +++ b/routers/repo/cloudbrain.go @@ -323,7 +323,7 @@ func CloudBrainDebug(ctx *context.Context) { var jobID = ctx.Params(":jobid") if !ctx.IsSigned { log.Error("the user has not signed in") - ctx.Error(http.StatusForbidden, "","the user has not signed in") + ctx.Error(http.StatusForbidden, "", "the user has not signed in") return } task, err := models.GetCloudbrainByJobID(jobID) @@ -340,7 +340,7 @@ func CloudBrainCommitImage(ctx *context.Context, form auth.CommitImageCloudBrain var jobID = ctx.Params(":jobid") if !ctx.IsSigned { log.Error("the user has not signed in") - ctx.Error(http.StatusForbidden, "","the user has not signed in") + ctx.Error(http.StatusForbidden, "", "the user has not signed in") return } task, err := models.GetCloudbrainByJobID(jobID) @@ -513,10 +513,10 @@ func CloudBrainShowModels(ctx *context.Context) { } //get dirs - dirs, err := getModelDirs(task.JobName, parentDir) + dirs, err := GetModelDirs(task.JobName, parentDir) if err != nil { - log.Error("getModelDirs failed:%v", err.Error(), ctx.Data["msgID"]) - ctx.ServerError("getModelDirs failed:", err) + log.Error("GetModelDirs failed:%v", err.Error(), ctx.Data["msgID"]) + ctx.ServerError("GetModelDirs failed:", err) return } @@ -567,7 +567,7 @@ func getImages(ctx *context.Context, imageType string) { log.Info("Get images end") } -func getModelDirs(jobName string, parentDir string) (string, error) { +func GetModelDirs(jobName string, parentDir string) (string, error) { var req string modelActualPath := setting.JobPath + jobName + "/model/" if parentDir == "" {