diff --git a/models/action.go b/models/action.go index 869acb762..247d4df0c 100755 --- a/models/action.go +++ b/models/action.go @@ -67,6 +67,7 @@ const ( ActionChangeUserAvatar //38 ActionCreateGrampusNPUDebugTask //39 ActionCreateGrampusGPUDebugTask //40 + ActionCreateGrampusGCUDebugTask //41 ) // Action represents user operation type and other information to diff --git a/models/task_config.go b/models/task_config.go index f86032fc9..14ca6b223 100644 --- a/models/task_config.go +++ b/models/task_config.go @@ -39,6 +39,7 @@ func GetTaskTypeFromAction(a ActionType) TaskType { ActionCreateGrampusGPUDebugTask, ActionCreateGrampusNPUDebugTask, ActionCreateGrampusNPUTrainTask, + ActionCreateGrampusGCUDebugTask, ActionCreateGrampusGPUTrainTask: return TaskCreateCloudbrainTask case ActionCreateRepo: diff --git a/modules/grampus/grampus.go b/modules/grampus/grampus.go index ff03680da..2c28bda24 100755 --- a/modules/grampus/grampus.go +++ b/modules/grampus/grampus.go @@ -156,7 +156,7 @@ func GenerateNotebookJob(ctx *context.Context, req *GenerateNotebookJobReq) (job var codeGrampus models.GrampusDataset var cpCommand string imageUrl := req.ImageUrl - if ProcessorTypeNPU == req.ProcessType { + if ProcessorTypeNPU == req.ProcessType || ProcessorTypeGCU == req.ProcessType { datasetGrampus = getDatasetGrampus(req.DatasetInfos) if len(req.ModelName) != 0 { datasetGrampus = append(datasetGrampus, models.GrampusDataset{ @@ -225,6 +225,7 @@ func GenerateNotebookJob(ctx *context.Context, req *GenerateNotebookJobReq) (job AutoStopDuration: autoStopDurationMs, Capacity: setting.Capacity, Command: req.Command, + CenterID: req.Spec.GetAvailableCenterIds(ctx.User.ID), }, }, }) @@ -273,6 +274,8 @@ func GenerateNotebookJob(ctx *context.Context, req *GenerateNotebookJobReq) (job actionType = models.ActionCreateGrampusNPUDebugTask } else if req.ComputeResource == models.GPUResource { actionType = models.ActionCreateGrampusGPUDebugTask + } else if req.ComputeResource == models.GCUResource { + actionType = models.ActionCreateGrampusGCUDebugTask } notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, jobID, req.DisplayJobName, actionType) diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go index 115842c30..eb25bd8c2 100644 --- a/routers/repo/ai_model_manage.go +++ b/routers/repo/ai_model_manage.go @@ -74,7 +74,7 @@ func saveModelByParameters(jobId string, versionName string, name string, versio cloudType := aiTask.Type modelSelectedFile := ctx.Query("modelSelectedFile") //download model zip //train type - if aiTask.ComputeResource == models.NPUResource { + if aiTask.ComputeResource == models.NPUResource || aiTask.ComputeResource == models.GCUResource { cloudType = models.TypeCloudBrainTwo } else if aiTask.ComputeResource == models.GPUResource { cloudType = models.TypeCloudBrainOne diff --git a/routers/repo/grampus.go b/routers/repo/grampus.go index 4655a5dc6..643886a8e 100755 --- a/routers/repo/grampus.go +++ b/routers/repo/grampus.go @@ -288,7 +288,6 @@ func GrampusNotebookCreate(ctx *context.Context, form auth.CreateGrampusNotebook req.ModelVersion = form.ModelVersion req.PreTrainModelUrl = form.PreTrainModelUrl req.PreTrainModelPath = getPreTrainModelPath(form.PreTrainModelUrl, form.CkptName) - } _, err = grampus.GenerateNotebookJob(ctx, req) @@ -1660,7 +1659,11 @@ func GrampusNotebookRestart(ctx *context.Context) { if task.ComputeResource == models.NPUResource { computeSourceSimple = models.NPU action = models.ActionCreateGrampusNPUDebugTask + } else if task.ComputeResource == models.GCUResource { + computeSourceSimple = models.GCU + action = models.ActionCreateGrampusGCUDebugTask } + spec, err = resource.GetAndCheckSpec(ctx.User.ID, oldSpec.ID, models.FindSpecsOptions{ JobType: models.JobType(task.JobType), ComputeResource: computeSourceSimple, diff --git a/services/cloudbrain/cloudbrainTask/count.go b/services/cloudbrain/cloudbrainTask/count.go index 372afc7d9..5ebbba798 100644 --- a/services/cloudbrain/cloudbrainTask/count.go +++ b/services/cloudbrain/cloudbrainTask/count.go @@ -72,6 +72,11 @@ var StatusInfoDict = map[string]StatusInfo{string(models.JobTypeDebug) + "-" + s JobType: []models.JobType{models.JobTypeDebug}, NotFinalStatuses: GrampusNotFinalStatuses, ComputeResource: models.NPUResource, +}, string(models.JobTypeDebug) + "-" + strconv.Itoa(models.TypeC2Net) + "-" + models.GCUResource: { + CloudBrainTypes: []int{models.TypeC2Net}, + JobType: []models.JobType{models.JobTypeDebug}, + NotFinalStatuses: GrampusNotFinalStatuses, + ComputeResource: models.GCUResource, }} func GetNotFinalStatusTaskCount(uid int64, cloudbrainType int, jobType string, computeResource ...string) (int, error) {