提交代码

modules/modelarts/modelarts.go

@@ -7,11 +7,12 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
-	"os"
 	"path"
 	"strconv"
 	"strings"
 
+	"github.com/go-resty/resty/v2"
+
 	"code.gitea.io/gitea/modules/cloudbrain"
 
 	"code.gitea.io/gitea/modules/modelarts_cd"
@@ -932,20 +933,12 @@ func HandleNotebookInfo(task *models.Cloudbrain) error {
 func uploadNoteBookFile(task *models.Cloudbrain, result *models.GetNotebook2Result) {
 	jupyterUrl := result.Url + "?token=" + result.Token
 	client := getRestyClient()
-	res, err := client.R().Get(jupyterUrl)
-	codePath := setting.JobPath + task.JobName + cloudbrain.CodeMountPath
-	if err != nil {
-		log.Error("browser jupyterUrl failed:%v", task.DisplayJobName, err)
-
+	cookies, xsrf := getCookiesAndCsrf(client, jupyterUrl)
+	if xsrf == "" {
+		log.Error("browser jupyterUrl failed:%v", task.DisplayJobName)
 	} else {
-		cookies := res.Cookies()
-		xsrf := ""
-		for _, cookie := range cookies {
-			if cookie.Name == "_xsrf" {
-				xsrf = cookie.Value
-			}
-		}
 
+		codePath := setting.JobPath + task.JobName + cloudbrain.CodeMountPath
 		fileContents, err := ioutil.ReadFile(codePath + "/" + task.BootFile)
 		if err != nil {
 			log.Error("read jupyter file failed:%v", task.DisplayJobName, err)
@@ -953,8 +946,8 @@ func uploadNoteBookFile(task *models.Cloudbrain, result *models.GetNotebook2Resu
 
 		base64Content := base64.StdEncoding.EncodeToString(fileContents)
 
-		uploadUrl := result.Url + "/api/contents/" + path.Base(task.BootFile)
-		res, err = client.R().
+		uploadUrl := getJupyterBaseUrl(result.Url) + "api/contents/" + path.Base(task.BootFile)
+		res, err := client.R().
 			SetCookies(cookies).
 			SetHeader("X-XSRFToken", xsrf).
 			SetBody(map[string]interface{}{
@@ -971,7 +964,40 @@ func uploadNoteBookFile(task *models.Cloudbrain, result *models.GetNotebook2Resu
 		}
 
 	}
-	go os.RemoveAll(codePath)
+
+}
+
+func getJupyterBaseUrl(url string) string {
+	jupyterUrlLength := len(url)
+	baseUrl := url[0 : jupyterUrlLength-len(path.Base(url))]
+	return baseUrl
+}
+
+func getCookiesAndCsrf(client *resty.Client, jupyterUrl string) ([]*http.Cookie, string) {
+
+	var cookies []*http.Cookie
+	for i := 0; i < 4; i++ {
+		res, err := client.R().Get(jupyterUrl)
+		if err != nil {
+			log.Error("browser jupyterUrl failed.")
+			return cookies, ""
+		} else {
+			cookies = res.Cookies()
+			xsrf := ""
+			for _, cookie := range cookies {
+				if cookie.Name == "_xsrf" {
+					xsrf = cookie.Value
+				}
+
+			}
+			if xsrf != "" {
+				return cookies, xsrf
+			}
+
+		}
+	}
+	return cookies, ""
+
 }
 
 func SyncTempStatusJob() {

modules/setting/setting.go

@@ -1587,6 +1587,15 @@ func NewContext() {
 	Course.OrgName = sec.Key("org_name").MustString("")
 	Course.TeamName = sec.Key("team_name").MustString("")
 
+	sec = Cfg.Section("file_notebook")
+	FileNoteBook.ProjectName = sec.Key("project_name").MustString("openi-notebook")
+	FileNoteBook.ImageIdNPU = sec.Key("imageid_npu").MustString("")
+	FileNoteBook.ImageNPU = sec.Key("image_npu").MustString("")
+	FileNoteBook.ImageGPU = sec.Key("image_gpu").MustString("")
+	FileNoteBook.SpecIdCPU = sec.Key("specid_cpu").MustInt64(-1)
+	FileNoteBook.SpecIdGPU = sec.Key("specid_gpu").MustInt64(-1)
+	FileNoteBook.SpecIdNPU = sec.Key("specid_npu").MustInt64(-1)
+
 	getGrampusConfig()
 	getModelartsCDConfig()
 	getModelConvertConfig()

routers/api/v1/repo/modelarts.go

@@ -39,8 +39,11 @@ func GetModelArtsNotebook2(ctx *context.APIContext) {
 	ID := ctx.Params(":id")
 	job, err := models.GetCloudbrainByID(ID)
 	if err != nil {
-		ctx.NotFound(err)
-		return
+		job, err = models.GetCloudbrainByJobID(ID)
+		if err != nil {
+			ctx.NotFound(err)
+			return
+		}
 	}
 	err = modelarts.HandleNotebookInfo(job)
 	if err != nil {

services/cloudbrain/cloudbrainTask/notebook.go

@@ -3,7 +3,7 @@ package cloudbrainTask
 import (
 	"fmt"
 	"net/http"
-	"os"
+	"path"
 
 	"code.gitea.io/gitea/modules/modelarts"
 	"code.gitea.io/gitea/modules/modelarts_cd"
@@ -28,11 +28,25 @@ import (
 	"code.gitea.io/gitea/modules/util"
 )
 
+const NoteBookExtension = ".ipynb"
+
 func FileNotebookCreate(ctx *context.Context, option api.CreateFileNotebookJobOption) {
 
 	if ctx.Written() {
 		return
 	}
+
+	if path.Ext(option.File) != NoteBookExtension {
+		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_select_wrong")))
+		return
+	}
+
+	isNotebookFileExist, _ := isNoteBookFileExist(ctx, option)
+	if !isNotebookFileExist {
+		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_file_not_exist")))
+		return
+	}
+
 	//create repo if not exist
 	repo, err := models.GetRepositoryByName(ctx.User.ID, setting.FileNoteBook.ProjectName)
 	if repo == nil {
@@ -51,6 +65,12 @@ func FileNotebookCreate(ctx *context.Context, option api.CreateFileNotebookJobOp
 	}
 	if err != nil {
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi("repo.failed_to_create_repo"))
+		return
+	}
+	if option.Type <= 1 {
+		cloudBrainFileNoteBookCreate(ctx, option, repo)
+	} else {
+		modelartsFileNoteBookCreate(ctx, option, repo)
 	}
 
 }
@@ -99,40 +119,19 @@ func cloudBrainFileNoteBookCreate(ctx *context.Context, option api.CreateFileNot
 			return
 		}
 	}
-	repoPath := models.RepoPath(repo.OwnerName, repo.Name)
-	gitRepo, err := git.OpenRepository(repoPath)
-	if err != nil {
-		ctx.Error(500, "RepoRef Invalid repo "+repoPath, err.Error())
-		return
-	}
-	// We opened it, we should close it
-	defer func() {
-		// If it's been set to nil then assume someone else has closed it.
-		if ctx.Repo.GitRepo != nil {
-			ctx.Repo.GitRepo.Close()
-		}
-	}()
 
-	ctx.Repo = &context.Repository{
-		Repository: repo,
-		GitRepo:    gitRepo,
-	}
-
-	fileExist, err := ctx.Repo.FileExists(option.File, option.BranchName)
-	if err != nil || !fileExist {
-		log.Error("Get file error:", err, ctx.Data["MsgID"])
+	sourceRepo, err := models.GetRepositoryByOwnerAndName(option.OwnerName, option.ProjectName)
+	if err != nil {
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_file_not_exist")))
 		return
 	}
 
-	command := cloudbrain.GetCloudbrainDebugCommand()
-
-	errStr := uploadCodeFile(repo, codePath, option.BranchName, option.File, jobName)
+	errStr := uploadCodeFile(sourceRepo, codePath, option.BranchName, option.File, jobName)
 	if errStr != "" {
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_file_not_exist")))
 		return
 	}
-
+	command := cloudbrain.GetCloudbrainDebugCommand()
 	commitID, _ := ctx.Repo.GitRepo.GetBranchCommitID(option.BranchName)
 	specId := setting.FileNoteBook.SpecIdGPU
 	if option.Type == 0 {
@@ -242,33 +241,14 @@ func modelartsFileNoteBookCreate(ctx *context.Context, option api.CreateFileNote
 			return
 		}
 	}
-	repoPath := models.RepoPath(repo.OwnerName, repo.Name)
-	gitRepo, err := git.OpenRepository(repoPath)
-	if err != nil {
-		ctx.Error(500, "RepoRef Invalid repo "+repoPath, err.Error())
-		return
-	}
-	// We opened it, we should close it
-	defer func() {
-		// If it's been set to nil then assume someone else has closed it.
-		if ctx.Repo.GitRepo != nil {
-			ctx.Repo.GitRepo.Close()
-		}
-	}()
 
-	ctx.Repo = &context.Repository{
-		Repository: repo,
-		GitRepo:    gitRepo,
-	}
-
-	fileExist, err := ctx.Repo.FileExists(option.File, option.BranchName)
-	if err != nil || !fileExist {
-		log.Error("Get file error:", err, ctx.Data["MsgID"])
+	sourceRepo, err := models.GetRepositoryByOwnerAndName(option.OwnerName, option.ProjectName)
+	if err != nil {
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_file_not_exist")))
 		return
 	}
 
-	err = downloadCode(repo, getCodePath(jobName), option.BranchName)
+	err = downloadCode(sourceRepo, getCodePath(jobName), option.BranchName)
 	if err != nil {
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("cloudbrain.load_code_failed")))
 		return
@@ -292,6 +272,9 @@ func modelartsFileNoteBookCreate(ctx *context.Context, option api.CreateFileNote
 		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("points.insufficient_points_balance")))
 		return
 	}
+	ctx.Repo = &context.Repository{
+		Repository: repo,
+	}
 
 	var jobId string
 	if setting.ModelartsCD.Enabled {
@@ -315,6 +298,30 @@ func modelartsFileNoteBookCreate(ctx *context.Context, option api.CreateFileNote
 
 }
 
+func isNoteBookFileExist(ctx *context.Context, option api.CreateFileNotebookJobOption) (bool, error) {
+	repoPathOfNoteBook := models.RepoPath(option.OwnerName, option.ProjectName)
+
+	gitRepoOfNoteBook, err := git.OpenRepository(repoPathOfNoteBook)
+	if err != nil {
+		log.Error("RepoRef Invalid repo "+repoPathOfNoteBook, err.Error())
+		return false, err
+	}
+	// We opened it, we should close it
+	defer func() {
+		// If it's been set to nil then assume someone else has closed it.
+		if gitRepoOfNoteBook != nil {
+			gitRepoOfNoteBook.Close()
+		}
+	}()
+	fileExist, err := fileExists(gitRepoOfNoteBook, option.File, option.BranchName)
+	if err != nil || !fileExist {
+		log.Error("Get file error:", err, ctx.Data["MsgID"])
+
+		return false, err
+	}
+	return true, nil
+}
+
 func uploadCodeFile(repo *models.Repository, codePath string, branchName string, filePath string, jobName string) string {
 	err := downloadCode(repo, codePath, branchName)
 	if err != nil {
@@ -325,6 +332,17 @@ func uploadCodeFile(repo *models.Repository, codePath string, branchName string,
 	if err != nil {
 		return "cloudbrain.load_code_failed"
 	}
-	go os.RemoveAll(codePath)
 	return ""
 }
+
+func fileExists(gitRepo *git.Repository, path string, branch string) (bool, error) {
+
+	commit, err := gitRepo.GetBranchCommit(branch)
+	if err != nil {
+		return false, err
+	}
+	if _, err := commit.GetTreeEntryByPath(path); err != nil {
+		return false, err
+	}
+	return true, nil
+}