调试支持模型

models/ai_model_manage.go

@@ -371,6 +371,18 @@ func QueryModelByName(name string, repoId int64) []*AiModelManage {
 	return aiModelManageList
 }
 
+func QueryModelByPath(path string) (*AiModelManage, error) {
+	modelManage := new(AiModelManage)
+	has, err := x.Where("path=?", path).Get(modelManage)
+	if err != nil {
+		return nil, err
+	}
+	if !has {
+		return nil, ErrNotExist{}
+	}
+	return modelManage, nil
+}
+
 func QueryModel(opts *AiModelQueryOptions) ([]*AiModelManage, int64, error) {
 	sess := x.NewSession()
 	defer sess.Close()

models/cloudbrain.go

@@ -298,6 +298,12 @@ func (task *Cloudbrain) IsUserHasRight(user *User) bool {
 	}
 	return user.IsAdmin || user.ID == task.UserID
 }
+func (task *Cloudbrain) IsGPUTask() bool {
+	return task.ComputeResource == GPUResource
+}
+func (task *Cloudbrain) IsNPUTask() bool {
+	return task.ComputeResource == NPUResource
+}
 
 func ConvertDurationToStr(duration int64) string {
 	if duration <= 0 {
@@ -1200,6 +1206,13 @@ type DatasetDownload struct {
 	IsDelete            bool   `json:"is_delete"`
 }
 
+type ModelDownload struct {
+	Name           string `json:"name"`
+	DownloadLink   string `json:"download_link"`
+	RepositoryLink string `json:"repository_link"`
+	IsDelete       bool   `json:"is_delete"`
+}
+
 type DataSource struct {
 	DatasetID      string `json:"dataset_id"`
 	DatasetVersion string `json:"dataset_version"`

modules/cloudbrain/cloudbrain.go

@@ -492,14 +492,13 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 	}
 
 	if task.PreTrainModelUrl != "" { //预训练
-		realPath := setting.Attachment.Minio.RealPath + task.PreTrainModelUrl
-		_, err := os.Stat(realPath)
+		_, err := models.QueryModelByPath(task.PreTrainModelUrl)
 		if err != nil {
 			log.Warn("The model may be deleted", err)
 		} else {
 			volumes = append(volumes, models.Volume{
 				HostPath: models.StHostPath{
-					Path:      realPath,
+					Path:      setting.Attachment.Minio.RealPath + task.PreTrainModelUrl,
 					MountPath: PretrainModelMountPath,
 					ReadOnly:  true,
 				},
@@ -557,10 +556,16 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 		GpuQueue:        task.GpuQueue,
 		ResourceSpecId:  task.ResourceSpecId,
 		ComputeResource: task.ComputeResource,
-		CreatedUnix:     createTime,
-		UpdatedUnix:     createTime,
-		BranchName:      task.BranchName,
-		Spec:            spec,
+
+		CreatedUnix:      createTime,
+		UpdatedUnix:      createTime,
+		BranchName:       task.BranchName,
+		Spec:             spec,
+		ModelName:        task.ModelName,
+		ModelVersion:     task.ModelVersion,
+		LabelName:        task.LabelName,
+		PreTrainModelUrl: task.PreTrainModelUrl,
+		CkptName:         task.CkptName,
 	}
 
 	err = models.RestartCloudbrain(task, newTask)
@@ -693,3 +698,22 @@ func GetCloudBrainByIdOrJobId(id string, initialQuery string) (*models.Cloudbrai
 	}
 	return job, err
 }
+
+type GenerateModelArtsNotebookReq struct {
+	JobName        string
+	DisplayJobName string
+	Uuid           string
+	Description    string
+
+	BootFile string
+
+	ImageId            string
+	AutoStopDurationMs int64
+
+	Spec             *models.Specification
+	ModelName        string
+	LabelName        string
+	CkptName         string
+	ModelVersion     string
+	PreTrainModelUrl string
+}

modules/modelarts/modelarts.go

@@ -180,24 +180,23 @@ type Parameters struct {
 	} `json:"parameter"`
 }
 
-
-func GenerateNotebook2(ctx *context.Context, displayJobName, jobName, uuid, description, imageId string, spec *models.Specification, bootFile string,autoStopDurationInMs int64) (string, error) {
+func GenerateNotebook2(ctx *context.Context, req cloudbrain.GenerateModelArtsNotebookReq) (string, error) {
 	if poolInfos == nil {
 		json.Unmarshal([]byte(setting.PoolInfos), &poolInfos)
 	}
 
-	imageName, err := GetNotebookImageName(imageId)
+	imageName, err := GetNotebookImageName(req.ImageId)
 	if err != nil {
 		log.Error("GetNotebookImageName failed: %v", err.Error())
 		return "", err
 	}
 	createTime := timeutil.TimeStampNow()
 	jobResult, err := createNotebook2(models.CreateNotebook2Params{
-		JobName:     jobName,
-		Description: description,
-		Flavor:      spec.SourceSpecId,
-		Duration:    autoStopDurationInMs,
-		ImageID:     imageId,
+		JobName:     req.JobName,
+		Description: req.Description,
+		Flavor:      req.Spec.SourceSpecId,
+		Duration:    req.AutoStopDurationMs,
+		ImageID:     req.ImageId,
 		PoolID:      poolInfos.PoolInfo[0].PoolId,
 		Feature:     models.NotebookFeature,
 		Volume: models.VolumeReq{
@@ -210,13 +209,13 @@ func GenerateNotebook2(ctx *context.Context, displayJobName, jobName, uuid, desc
 	if err != nil {
 		log.Error("createNotebook2 failed: %v", err.Error())
 		if strings.HasPrefix(err.Error(), UnknownErrorPrefix) {
-			log.Info("(%s)unknown error, set temp status", displayJobName)
+			log.Info("(%s)unknown error, set temp status", req.DisplayJobName)
 			errTemp := models.InsertCloudbrainTemp(&models.CloudbrainTemp{
 				JobID:     models.TempJobId,
 				VersionID: models.TempVersionId,
 				Status:    models.TempJobStatus,
 				Type:      models.TypeCloudBrainTwo,
-				JobName:   jobName,
+				JobName:   req.JobName,
 				JobType:   string(models.JobTypeDebug),
 			})
 			if errTemp != nil {
@@ -227,23 +226,28 @@ func GenerateNotebook2(ctx *context.Context, displayJobName, jobName, uuid, desc
 		return "", err
 	}
 	task := &models.Cloudbrain{
-		Status:          jobResult.Status,
-		UserID:          ctx.User.ID,
-		RepoID:          ctx.Repo.Repository.ID,
-		JobID:           jobResult.ID,
-		JobName:         jobName,
-		FlavorCode:      spec.SourceSpecId,
-		DisplayJobName:  displayJobName,
-		JobType:         string(models.JobTypeDebug),
-		Type:            models.TypeCloudBrainTwo,
-		Uuid:            uuid,
-		ComputeResource: models.NPUResource,
-		Image:           imageName,
-		BootFile:        bootFile,
-		Description:     description,
-		CreatedUnix:     createTime,
-		UpdatedUnix:     createTime,
-		Spec:            spec,
+		Status:           jobResult.Status,
+		UserID:           ctx.User.ID,
+		RepoID:           ctx.Repo.Repository.ID,
+		JobID:            jobResult.ID,
+		JobName:          req.JobName,
+		FlavorCode:       req.Spec.SourceSpecId,
+		DisplayJobName:   req.DisplayJobName,
+		JobType:          string(models.JobTypeDebug),
+		Type:             models.TypeCloudBrainTwo,
+		Uuid:             req.Uuid,
+		ComputeResource:  models.NPUResource,
+		Image:            imageName,
+		BootFile:         req.BootFile,
+		Description:      req.Description,
+		CreatedUnix:      createTime,
+		UpdatedUnix:      createTime,
+		Spec:             req.Spec,
+		ModelName:        req.ModelName,
+		ModelVersion:     req.ModelVersion,
+		LabelName:        req.LabelName,
+		PreTrainModelUrl: req.PreTrainModelUrl,
+		CkptName:         req.CkptName,
 	}
 
 	err = models.CreateCloudbrain(task)
@@ -252,7 +256,7 @@ func GenerateNotebook2(ctx *context.Context, displayJobName, jobName, uuid, desc
 	}
 
 	stringId := strconv.FormatInt(task.ID, 10)
-	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, stringId, displayJobName, models.ActionCreateDebugNPUTask)
+	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, stringId, req.DisplayJobName, models.ActionCreateDebugNPUTask)
 	return jobResult.ID, nil
 }
 
@@ -872,14 +876,14 @@ func getJupyterBaseUrl(url string) string {
 }
 
 func getCookiesAndCsrf(jupyterUrl string) ([]*http.Cookie, string) {
-    log.Info("jupyter url:"+jupyterUrl)
+	log.Info("jupyter url:" + jupyterUrl)
 	var cookies []*http.Cookie
 	const retryTimes = 10
 	for i := 0; i < retryTimes; i++ {
 		res, err := http.Get(jupyterUrl)
 		if err != nil {
-			log.Error("browser jupyterUrl failed.",err)
-			if i==retryTimes-1{
+			log.Error("browser jupyterUrl failed.", err)
+			if i == retryTimes-1 {
 				return cookies, ""
 			}
 

modules/modelarts_cd/modelarts.go

@@ -5,6 +5,8 @@ import (
 	"strconv"
 	"strings"
 
+	"code.gitea.io/gitea/modules/cloudbrain"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/context"
 	"code.gitea.io/gitea/modules/log"
@@ -88,19 +90,19 @@ type Parameters struct {
 	} `json:"parameter"`
 }
 
-func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, description, imageId string, spec *models.Specification, bootFile string,autoStopDurationInMs int64) (string, error) {
-	imageName, err := GetNotebookImageName(imageId)
+func GenerateNotebook(ctx *context.Context, req cloudbrain.GenerateModelArtsNotebookReq) (string, error) {
+	imageName, err := GetNotebookImageName(req.ImageId)
 	if err != nil {
 		log.Error("GetNotebookImageName failed: %v", err.Error())
 		return "", err
 	}
 	createTime := timeutil.TimeStampNow()
 	jobResult, err := createNotebook(models.CreateNotebookWithoutPoolParams{
-		JobName:     jobName,
-		Description: description,
-		Flavor:      spec.SourceSpecId,
-		Duration:    autoStopDurationInMs,
-		ImageID:     imageId,
+		JobName:     req.JobName,
+		Description: req.Description,
+		Flavor:      req.Spec.SourceSpecId,
+		Duration:    req.AutoStopDurationMs,
+		ImageID:     req.ImageId,
 		Feature:     models.NotebookFeature,
 		Volume: models.VolumeReq{
 			Capacity:  setting.Capacity,
@@ -112,13 +114,13 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr
 	if err != nil {
 		log.Error("createNotebook failed: %v", err.Error())
 		if strings.HasPrefix(err.Error(), UnknownErrorPrefix) {
-			log.Info("(%s)unknown error, set temp status", displayJobName)
+			log.Info("(%s)unknown error, set temp status", req.DisplayJobName)
 			errTemp := models.InsertCloudbrainTemp(&models.CloudbrainTemp{
 				JobID:     models.TempJobId,
 				VersionID: models.TempVersionId,
 				Status:    models.TempJobStatus,
 				Type:      models.TypeCDCenter,
-				JobName:   jobName,
+				JobName:   req.JobName,
 				JobType:   string(models.JobTypeDebug),
 			})
 			if errTemp != nil {
@@ -129,23 +131,28 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr
 		return "", err
 	}
 	task := &models.Cloudbrain{
-		Status:          jobResult.Status,
-		UserID:          ctx.User.ID,
-		RepoID:          ctx.Repo.Repository.ID,
-		JobID:           jobResult.ID,
-		JobName:         jobName,
-		FlavorCode:      spec.SourceSpecId,
-		DisplayJobName:  displayJobName,
-		JobType:         string(models.JobTypeDebug),
-		Type:            models.TypeCDCenter,
-		Uuid:            uuid,
-		ComputeResource: models.NPUResource,
-		Image:           imageName,
-		Description:     description,
-		CreatedUnix:     createTime,
-		UpdatedUnix:     createTime,
-		Spec:            spec,
-		BootFile:        bootFile,
+		Status:           jobResult.Status,
+		UserID:           ctx.User.ID,
+		RepoID:           ctx.Repo.Repository.ID,
+		JobID:            jobResult.ID,
+		JobName:          req.JobName,
+		FlavorCode:       req.Spec.SourceSpecId,
+		DisplayJobName:   req.DisplayJobName,
+		JobType:          string(models.JobTypeDebug),
+		Type:             models.TypeCDCenter,
+		Uuid:             req.Uuid,
+		ComputeResource:  models.NPUResource,
+		Image:            imageName,
+		Description:      req.Description,
+		CreatedUnix:      createTime,
+		UpdatedUnix:      createTime,
+		Spec:             req.Spec,
+		BootFile:         req.BootFile,
+		ModelName:        req.ModelName,
+		ModelVersion:     req.ModelVersion,
+		LabelName:        req.LabelName,
+		PreTrainModelUrl: req.PreTrainModelUrl,
+		CkptName:         req.CkptName,
 	}
 
 	err = models.CreateCloudbrain(task)
@@ -154,7 +161,7 @@ func GenerateNotebook(ctx *context.Context, displayJobName, jobName, uuid, descr
 	}
 
 	stringId := strconv.FormatInt(task.ID, 10)
-	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, stringId, displayJobName, models.ActionCreateDebugNPUTask)
+	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, stringId, req.DisplayJobName, models.ActionCreateDebugNPUTask)
 	return jobResult.ID, nil
 }
 

routers/repo/cloudbrain.go

@@ -373,6 +373,13 @@ func cloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	}
 
 	if form.ModelName != "" { //使用预训练模型训练
+		_, err := models.QueryModelByPath(form.PreTrainModelUrl)
+		if err != nil {
+			log.Error("Can not find model", err)
+			cloudBrainNewDataPrepare(ctx, jobType)
+			ctx.RenderWithErr(ctx.Tr("repo.modelconvert.manage.model_not_exist"), tpl, &form)
+			return
+		}
 		req.ModelName = form.ModelName
 		req.LabelName = form.LabelName
 		req.CkptName = form.CkptName
@@ -684,6 +691,15 @@ func CloudBrainRestart(ctx *context.Context) {
 				break
 			}
 		}
+		if task.PreTrainModelUrl != "" {
+			_, err := models.QueryModelByPath(task.PreTrainModelUrl)
+			if err != nil {
+				log.Error("Can not find model", err)
+				resultCode = "-1"
+				errorMsg = ctx.Tr("repo.modelconvert.manage.model_not_exist")
+				break
+			}
+		}
 
 		err = cloudbrain.RestartTask(ctx, task, &ID)
 		if err != nil {
@@ -842,10 +858,10 @@ func cloudBrainShow(ctx *context.Context, tpName base.TplName, jobType models.Jo
 func CloudBrainDebug(ctx *context.Context) {
 	task := ctx.Cloudbrain
 	debugUrl := setting.DebugServerHost + "jpylab_" + task.JobID + "_" + task.SubTaskName
-	if task.BootFile!=""{
-		ctx.Redirect(getFileUrl(debugUrl,task.BootFile))
+	if task.BootFile != "" {
+		ctx.Redirect(getFileUrl(debugUrl, task.BootFile))
 
-	}else{
+	} else {
 		ctx.Redirect(debugUrl)
 	}
 

routers/repo/grampus.go

@@ -253,6 +253,14 @@ func GrampusNotebookCreate(ctx *context.Context, form auth.CreateGrampusNotebook
 	}
 
 	if form.ModelName != "" { //使用预训练模型训练
+
+		_, err := models.QueryModelByPath(form.PreTrainModelUrl)
+		if err != nil {
+			log.Error("Can not find model", err)
+			grampusNotebookNewDataPrepare(ctx, processType)
+			ctx.RenderWithErr(ctx.Tr("repo.modelconvert.manage.model_not_exist"), tpl, &form)
+			return
+		}
 		req.ModelName = form.ModelName
 		req.LabelName = form.LabelName
 		req.CkptName = form.CkptName
@@ -1148,6 +1156,7 @@ func GrampusNotebookShow(ctx *context.Context) {
 
 	ctx.Data["task"] = task
 	ctx.Data["datasetDownload"] = getDatasetDownloadInfo(ctx, task)
+	ctx.Data["modelDownload"] = getModelDownloadInfo(ctx, task)
 	ctx.Data["canDownload"] = cloudbrain.CanModifyJob(ctx, task)
 	ctx.Data["ai_center"] = cloudbrainService.GetAiCenterShow(task.AiCenter, ctx)
 	ctx.HTML(http.StatusOK, tplGrampusNotebookShow)
@@ -1157,28 +1166,62 @@ func getDatasetDownloadInfo(ctx *context.Context, task *models.Cloudbrain) []mod
 	datasetDownload := make([]models.DatasetDownload, 0)
 	if ctx.IsSigned {
 		if task.Uuid != "" && task.UserID == ctx.User.ID {
-			datasetDownload = GetCloudBrainDataSetInfo(task.Uuid, task.DatasetName, true)
+			if task.IsGPUTask() {
+				return GetCloudBrainDataSetInfo(task.Uuid, task.DatasetName, false)
+			} else {
+				datasetDownload = GetCloudBrainDataSetInfo(task.Uuid, task.DatasetName, true)
+				datasetObsUrlList := make([]NotebookDataset, 0)
+				_ = json.Unmarshal([]byte(task.DataUrl), &datasetObsUrlList)
+
+				for _, datasetInfo := range datasetDownload {
+					match := false
+					for _, datasetObs := range datasetObsUrlList {
+						if strings.Contains(datasetObs.DatasetUrl, datasetInfo.DatasetName) {
+							datasetInfo.DatasetDownloadLink = datasetObs.DatasetUrl
+							match = true
+							break
+						}
+					}
+					if !match {
+						datasetInfo.DatasetDownloadLink = ""
+					}
+
+				}
+
+			}
+
 		}
 	}
-	datasetObsUrlList := make([]NotebookDataset, 0)
-	_ = json.Unmarshal([]byte(task.DataUrl), &datasetObsUrlList)
 
-	for _, datasetInfo := range datasetDownload {
-		match := false
-		for _, datasetObs := range datasetObsUrlList {
+	return datasetDownload
+}
+
+func getModelDownloadInfo(ctx *context.Context, task *models.Cloudbrain) *models.ModelDownload {
+	var modelDownload models.ModelDownload
+	if ctx.IsSigned {
+		if task.ModelName != "" && task.UserID == ctx.User.ID {
+			if task.IsNPUTask() {
+				modelDownload = models.ModelDownload{
+					Name:         task.CkptName,
+					DownloadLink: "",
+					IsDelete:     false,
+				}
+				datasetObsUrlList := make([]NotebookDataset, 0)
+				_ = json.Unmarshal([]byte(task.DataUrl), &datasetObsUrlList)
+				for _, datasetObs := range datasetObsUrlList {
+					if strings.Contains(datasetObs.DatasetUrl, task.CkptName) {
+						modelDownload.DownloadLink = datasetObs.DatasetUrl
+						break
+					}
+				}
 
-			if strings.Contains(datasetObs.DatasetUrl, datasetInfo.DatasetName) {
-				datasetInfo.DatasetDownloadLink = datasetObs.DatasetUrl
-				match = true
-				break
 			}
-		}
-		if !match {
-			datasetInfo.DatasetDownloadLink = ""
+
 		}
 
 	}
-	return datasetDownload
+
+	return &modelDownload
 }
 
 func GrampusTrainJobShow(ctx *context.Context) {
@@ -1605,6 +1648,16 @@ func GrampusNotebookRestart(ctx *context.Context) {
 			errorMsg = ctx.Tr("points.insufficient_points_balance")
 			break
 		}
+
+		if task.ModelName != "" { //使用预训练模型训练
+			_, err := models.QueryModelByPath(task.PreTrainModelUrl)
+			if err != nil {
+				log.Error("Can not find model", err)
+				errorMsg = ctx.Tr("repo.modelconvert.manage.model_not_exist")
+				break
+			}
+		}
+
 		createTime := timeutil.TimeStampNow()
 
 		res, err := grampus.RestartNotebookJob(task.JobID)
@@ -1615,21 +1668,26 @@ func GrampusNotebookRestart(ctx *context.Context) {
 		}
 
 		newTask := &models.Cloudbrain{
-			Status:          res.Status,
-			UserID:          task.UserID,
-			RepoID:          task.RepoID,
-			JobID:           res.NewId,
-			JobName:         task.JobName,
-			DisplayJobName:  task.DisplayJobName,
-			JobType:         task.JobType,
-			Type:            task.Type,
-			Uuid:            task.Uuid,
-			Image:           task.Image,
-			ComputeResource: task.ComputeResource,
-			Description:     task.Description,
-			CreatedUnix:     createTime,
-			UpdatedUnix:     createTime,
-			Spec:            spec,
+			Status:           res.Status,
+			UserID:           task.UserID,
+			RepoID:           task.RepoID,
+			JobID:            res.NewId,
+			JobName:          task.JobName,
+			DisplayJobName:   task.DisplayJobName,
+			JobType:          task.JobType,
+			Type:             task.Type,
+			Uuid:             task.Uuid,
+			Image:            task.Image,
+			ComputeResource:  task.ComputeResource,
+			Description:      task.Description,
+			CreatedUnix:      createTime,
+			UpdatedUnix:      createTime,
+			Spec:             spec,
+			ModelName:        task.ModelName,
+			ModelVersion:     task.ModelVersion,
+			LabelName:        task.LabelName,
+			PreTrainModelUrl: task.PreTrainModelUrl,
+			CkptName:         task.CkptName,
 		}
 
 		err = models.RestartCloudbrain(task, newTask)

routers/repo/modelarts.go

@@ -239,10 +239,37 @@ func Notebook2Create(ctx *context.Context, form auth.CreateModelArtsNotebookForm
 		return
 	}
 
+	req := cloudbrain.GenerateModelArtsNotebookReq{
+		DisplayJobName:     displayJobName,
+		JobName:            jobName,
+		Description:        description,
+		Uuid:               uuid,
+		ImageId:            imageId,
+		Spec:               spec,
+		BootFile:           "",
+		AutoStopDurationMs: modelarts.AutoStopDurationMs,
+	}
+
+	if form.ModelName != "" { //使用预训练模型训练
+		_, err := models.QueryModelByPath(form.PreTrainModelUrl)
+		if err != nil {
+			log.Error("Can not find model", err)
+			notebookNewDataPrepare(ctx)
+			ctx.RenderWithErr(ctx.Tr("repo.modelconvert.manage.model_not_exist"), tplModelArtsNotebookNew, &form)
+			return
+		}
+		req.ModelName = form.ModelName
+		req.LabelName = form.LabelName
+		req.CkptName = form.CkptName
+		req.ModelVersion = form.ModelVersion
+		req.PreTrainModelUrl = form.PreTrainModelUrl
+
+	}
+
 	if setting.ModelartsCD.Enabled {
-		_, err = modelarts_cd.GenerateNotebook(ctx, displayJobName, jobName, uuid, description, imageId, spec, "", modelarts.AutoStopDurationMs)
+		_, err = modelarts_cd.GenerateNotebook(ctx, req)
 	} else {
-		_, err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, uuid, description, imageId, spec, "", modelarts.AutoStopDurationMs)
+		_, err = modelarts.GenerateNotebook2(ctx, req)
 	}
 
 	if err != nil {
@@ -280,10 +307,16 @@ func NotebookShow(ctx *context.Context) {
 	}
 
 	datasetDownload := make([]models.DatasetDownload, 0)
+	var modelDownload models.ModelDownload
 	if ctx.IsSigned {
 		if task.Uuid != "" && task.UserID == ctx.User.ID {
 			datasetDownload = GetCloudBrainDataSetInfo(task.Uuid, task.DatasetName, true)
 		}
+		if task.ModelName != "" && task.UserID == ctx.User.ID {
+			modelDownload = GetModelDownload(task)
+
+		}
+
 	}
 	user, err := models.GetUserByID(task.UserID)
 	if err == nil {
@@ -304,6 +337,7 @@ func NotebookShow(ctx *context.Context) {
 	}
 	ctx.Data["duration"] = task.TrainJobDuration
 	ctx.Data["datasetDownload"] = datasetDownload
+	ctx.Data["modelDownload"] = modelDownload
 	ctx.Data["task"] = task
 	ctx.Data["ID"] = ID
 	ctx.Data["jobName"] = task.JobName
@@ -311,6 +345,22 @@ func NotebookShow(ctx *context.Context) {
 	ctx.HTML(200, tplModelArtsNotebookShow)
 }
 
+func GetModelDownload(task *models.Cloudbrain) models.ModelDownload {
+	modelDownload := models.ModelDownload{
+		Name:         task.CkptName,
+		DownloadLink: "s3://" + task.PreTrainModelUrl + task.CkptName,
+		IsDelete:     false,
+	}
+
+	_, err := models.QueryModelByPath(task.PreTrainModelUrl)
+
+	if err != nil {
+		log.Warn("Can not get model by path:"+task.PreTrainModelUrl, err)
+		modelDownload.IsDelete = true
+	}
+	return modelDownload
+}
+
 func GetCloudBrainDataSetInfo(uuid string, datasetname string, isNeedDown bool) []models.DatasetDownload {
 	datasetDownload := make([]models.DatasetDownload, 0)
 	if len(uuid) == 0 {
@@ -476,6 +526,15 @@ func NotebookRestart(ctx *context.Context) {
 			errorMsg = ctx.Tr("points.insufficient_points_balance")
 			break
 		}
+		if task.ModelName != "" { //使用预训练模型训练
+			_, err := models.QueryModelByPath(task.PreTrainModelUrl)
+			if err != nil {
+				log.Error("Can not find model", err)
+				errorMsg = ctx.Tr("repo.modelconvert.manage.model_not_exist")
+				break
+			}
+		}
+
 		createTime := timeutil.TimeStampNow()
 		param := models.NotebookAction{
 			Action: models.ActionStart,
@@ -511,21 +570,26 @@ func NotebookRestart(ctx *context.Context) {
 		}
 
 		newTask := &models.Cloudbrain{
-			Status:          res.Status,
-			UserID:          task.UserID,
-			RepoID:          task.RepoID,
-			JobID:           task.JobID,
-			JobName:         task.JobName,
-			DisplayJobName:  task.DisplayJobName,
-			JobType:         task.JobType,
-			Type:            task.Type,
-			Uuid:            task.Uuid,
-			Image:           task.Image,
-			ComputeResource: task.ComputeResource,
-			Description:     task.Description,
-			CreatedUnix:     createTime,
-			UpdatedUnix:     createTime,
-			Spec:            spec,
+			Status:           res.Status,
+			UserID:           task.UserID,
+			RepoID:           task.RepoID,
+			JobID:            task.JobID,
+			JobName:          task.JobName,
+			DisplayJobName:   task.DisplayJobName,
+			JobType:          task.JobType,
+			Type:             task.Type,
+			Uuid:             task.Uuid,
+			Image:            task.Image,
+			ComputeResource:  task.ComputeResource,
+			Description:      task.Description,
+			CreatedUnix:      createTime,
+			UpdatedUnix:      createTime,
+			Spec:             spec,
+			ModelName:        task.ModelName,
+			ModelVersion:     task.ModelVersion,
+			LabelName:        task.LabelName,
+			PreTrainModelUrl: task.PreTrainModelUrl,
+			CkptName:         task.CkptName,
 		}
 
 		err = models.RestartCloudbrain(task, newTask)

services/cloudbrain/cloudbrainTask/notebook.go

@@ -82,7 +82,7 @@ func FileNotebookCreate(ctx *context.Context, option api.CreateFileNotebookJobOp
 		})
 	}
 	if err != nil {
-		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.failed_to_create_notebook_repo",setting.FileNoteBook.ProjectName)))
+		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.failed_to_create_notebook_repo", setting.FileNoteBook.ProjectName)))
 		return
 	}
 	if option.Type <= 1 {
@@ -291,10 +291,21 @@ func modelartsFileNoteBookCreate(ctx *context.Context, option api.CreateFileNote
 	}
 
 	var jobId string
+	req := cloudbrain.GenerateModelArtsNotebookReq{
+		DisplayJobName:     displayJobName,
+		JobName:            jobName,
+		Description:        getDescription(option),
+		ImageId:            setting.FileNoteBook.ImageIdNPU,
+		Spec:               spec,
+		BootFile:           "",
+		AutoStopDurationMs: modelarts.AutoStopDurationMs / 4,
+	}
+
 	if setting.ModelartsCD.Enabled {
-		jobId, err = modelarts_cd.GenerateNotebook(ctx, displayJobName, jobName, "", getDescription(option), setting.FileNoteBook.ImageIdNPUCD, spec, option.File,modelarts.AutoStopDurationMs/4)
+		req.ImageId = setting.FileNoteBook.ImageIdNPUCD
+		jobId, err = modelarts_cd.GenerateNotebook(ctx, req)
 	} else {
-		jobId, err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, "", getDescription(option), setting.FileNoteBook.ImageIdNPU, spec, option.File,modelarts.AutoStopDurationMs/4)
+		jobId, err = modelarts.GenerateNotebook2(ctx, req)
 	}
 
 	if err != nil {