From 879d29594ce694a628b824583abfdb156c9f36b6 Mon Sep 17 00:00:00 2001 From: chenyifan01 Date: Thu, 11 Aug 2022 12:01:23 +0800 Subject: [PATCH] #2631 add cloudbrain task stop template msg --- modules/auth/wechat/client.go | 2 +- modules/auth/wechat/cloudbrain.go | 155 ++++++++++++++++++++++++++ modules/auth/wechat/template.go | 119 +++----------------- modules/notification/wechat/wechat.go | 15 +-- modules/setting/setting.go | 12 ++ 5 files changed, 187 insertions(+), 116 deletions(-) create mode 100644 modules/auth/wechat/cloudbrain.go diff --git a/modules/auth/wechat/client.go b/modules/auth/wechat/client.go index bb6db09d0..367a80c3a 100644 --- a/modules/auth/wechat/client.go +++ b/modules/auth/wechat/client.go @@ -61,7 +61,7 @@ type TemplateValue struct { Color string `json:"color"` } -type CloudbrainTaskData struct { +type DefaultWechatTemplate struct { First TemplateValue `json:"first"` Keyword1 TemplateValue `json:"keyword1"` Keyword2 TemplateValue `json:"keyword2"` diff --git a/modules/auth/wechat/cloudbrain.go b/modules/auth/wechat/cloudbrain.go new file mode 100644 index 000000000..9af091f15 --- /dev/null +++ b/modules/auth/wechat/cloudbrain.go @@ -0,0 +1,155 @@ +package wechat + +import ( + "code.gitea.io/gitea/models" + "code.gitea.io/gitea/modules/log" + "code.gitea.io/gitea/modules/setting" + "fmt" + "time" +) + +type JobOperateType string + +const ( + JobOperateTypeStart JobOperateType = "start" + JobOperateTypeStop JobOperateType = "stop" +) + +type CloudbrainStartMsg struct { +} + +func (CloudbrainStartMsg) Data(ctx *TemplateContext) *DefaultWechatTemplate { + return &DefaultWechatTemplate{ + First: TemplateValue{Value: setting.CloudbrainStartedTitle}, + Keyword1: TemplateValue{Value: ctx.Cloudbrain.DisplayJobName}, + Keyword2: TemplateValue{Value: getJobTypeDisplayName(ctx.Cloudbrain.JobType)}, + Keyword3: TemplateValue{Value: time.Unix(int64(ctx.Cloudbrain.CreatedUnix), 0).Format("2006-01-02 15:04:05")}, + Remark: TemplateValue{Value: setting.CloudbrainStartedRemark}, + } +} + +func (CloudbrainStartMsg) ShouldSend(ctx *TemplateContext) bool { + if len(setting.CloudbrainStartedNotifyList) == 0 { + return false + } + for _, v := range setting.CloudbrainStartedNotifyList { + if v == ctx.Cloudbrain.JobType { + return true + } + } + return false +} + +func (CloudbrainStartMsg) MsgId(ctx *TemplateContext) string { + return string(JobOperateTypeStart) + "_" + fmt.Sprint(ctx.Cloudbrain.ID) +} + +func (CloudbrainStartMsg) Url(ctx *TemplateContext) string { + repo, err := models.GetRepositoryByID(ctx.Cloudbrain.RepoID) + if err != nil { + log.Error("CloudbrainStartMsg GetRepositoryByID error,%v", err) + return "" + } + return getCloudbrainTemplateUrl(*ctx.Cloudbrain, repo) +} + +type CloudbrainStopMsg struct { +} + +func (CloudbrainStopMsg) Data(ctx *TemplateContext) *DefaultWechatTemplate { + return &DefaultWechatTemplate{ + First: TemplateValue{Value: setting.CloudbrainStoppedTitle}, + Keyword1: TemplateValue{Value: ctx.Cloudbrain.DisplayJobName}, + Keyword2: TemplateValue{Value: getJobTypeDisplayName(ctx.Cloudbrain.JobType)}, + Keyword3: TemplateValue{Value: time.Unix(int64(ctx.Cloudbrain.CreatedUnix), 0).Format("2006-01-02 15:04:05")}, + Remark: TemplateValue{Value: setting.CloudbrainStoppedRemark}, + } +} + +func (CloudbrainStopMsg) ShouldSend(ctx *TemplateContext) bool { + if len(setting.CloudbrainStoppedNotifyList) == 0 { + return false + } + for _, v := range setting.CloudbrainStoppedNotifyList { + if v == ctx.Cloudbrain.JobType { + return true + } + } + return false +} + +func (CloudbrainStopMsg) MsgId(ctx *TemplateContext) string { + return string(JobOperateTypeStop) + "_" + fmt.Sprint(ctx.Cloudbrain.ID) +} + +func (CloudbrainStopMsg) Url(ctx *TemplateContext) string { + repo, err := models.GetRepositoryByID(ctx.Cloudbrain.RepoID) + if err != nil { + log.Error("CloudbrainStopMsg GetRepositoryByID error,%v", err) + return "" + } + return getCloudbrainTemplateUrl(*ctx.Cloudbrain, repo) +} + +var startMsg = &CloudbrainStartMsg{} +var stopMsg = &CloudbrainStopMsg{} + +func GetTemplateFromOperateType(operate JobOperateType) Template { + switch operate { + case JobOperateTypeStart: + return startMsg + case JobOperateTypeStop: + return stopMsg + } + return nil +} + +func GetJobOperateTypeFromCloudbrainStatus(cloudbrain *models.Cloudbrain) JobOperateType { + if cloudbrain.IsTerminal() { + return JobOperateTypeStop + } + if cloudbrain.IsRunning() { + return JobOperateTypeStart + } + return "" +} + +func getCloudbrainTemplateUrl(cloudbrain models.Cloudbrain, repo *models.Repository) string { + url := setting.AppURL + repo.FullName() + + switch cloudbrain.JobType { + case string(models.JobTypeDebug): + if cloudbrain.ComputeResource == "CPU/GPU" { + url += "/cloudbrain/" + fmt.Sprint(cloudbrain.ID) + } else { + url += "/modelarts/notebook/" + fmt.Sprint(cloudbrain.ID) + } + case string(models.JobTypeBenchmark): + url += "/cloudbrain/benchmark/" + fmt.Sprint(cloudbrain.ID) + case string(models.JobTypeTrain): + if cloudbrain.Type == models.TypeCloudBrainOne { + url += "/cloudbrain/train-job/" + fmt.Sprint(cloudbrain.JobID) + } else if cloudbrain.Type == models.TypeCloudBrainTwo { + url += "/modelarts/train-job/" + fmt.Sprint(cloudbrain.JobID) + } else if cloudbrain.Type == models.TypeC2Net { + url += "/grampus/train-job/" + fmt.Sprint(cloudbrain.JobID) + } + case string(models.JobTypeInference): + url += "/modelarts/inference-job/" + fmt.Sprint(cloudbrain.JobID) + } + return url +} + +func getJobTypeDisplayName(jobType string) string { + switch jobType { + case string(models.JobTypeDebug): + return "调试任务" + case string(models.JobTypeBenchmark): + return "评测任务" + case string(models.JobTypeTrain): + return "训练任务" + case string(models.JobTypeInference): + return "推理任务" + } + return "" +} diff --git a/modules/auth/wechat/template.go b/modules/auth/wechat/template.go index 6c19651a6..bc206e3b2 100644 --- a/modules/auth/wechat/template.go +++ b/modules/auth/wechat/template.go @@ -6,140 +6,55 @@ import ( "code.gitea.io/gitea/modules/setting" "errors" "fmt" - "time" ) -type JobOperateType string - -const ( - JobOperateTypeStart JobOperateType = "start" - JobOperateTypeStop JobOperateType = "stop" -) +type Template interface { + ShouldSend(ctx *TemplateContext) bool + Data(ctx *TemplateContext) *DefaultWechatTemplate + MsgId(ctx *TemplateContext) string + Url(ctx *TemplateContext) string +} -func GetJobOperateTypeFromCloudbrainStatus(cloudbrain *models.Cloudbrain) JobOperateType { - if cloudbrain.IsTerminal() { - return JobOperateTypeStop - } - if cloudbrain.IsRunning() { - return JobOperateTypeStart - } - return "" +type TemplateContext struct { + Cloudbrain *models.Cloudbrain } -func SendCloudbrainStartedMsg(operateType JobOperateType, cloudbrain models.Cloudbrain) error { +func SendTemplateMsg(template Template, ctx *TemplateContext, userId int64) error { defer func() { if err := recover(); err != nil { combinedErr := fmt.Errorf("%s\n%s", err, log.Stack(2)) log.Error("PANIC:", combinedErr) } }() - repo, err := models.GetRepositoryByID(cloudbrain.RepoID) - if err != nil { - log.Error("SendCloudbrainStartedMsg GetRepositoryByID error,%v", err) - } - - if setting.CloudbrainStartedTemplateId == "" { + if !template.ShouldSend(ctx) { return nil } - openId := models.GetUserWechatOpenId(cloudbrain.UserID) + openId := models.GetUserWechatOpenId(userId) if openId == "" { return errors.New("Wechat openId not exist") } - data := CloudbrainTaskData{ - First: TemplateValue{Value: getCloudbrainTemplateTitle(operateType)}, - Keyword1: TemplateValue{Value: cloudbrain.DisplayJobName}, - Keyword2: TemplateValue{Value: getJobTypeDisplayName(cloudbrain.JobType)}, - Keyword3: TemplateValue{Value: time.Unix(int64(cloudbrain.CreatedUnix), 0).Format("2006-01-02 15:04:05")}, - Remark: TemplateValue{Value: getCloudbrainTemplateRemark(operateType)}, - } req := TemplateMsgRequest{ ToUser: openId, TemplateId: setting.CloudbrainStartedTemplateId, - Url: getCloudbrainTemplateUrl(cloudbrain, repo), - ClientMsgId: string(operateType) + "_" + fmt.Sprint(cloudbrain.ID), - Data: data, + Url: template.Url(ctx), + ClientMsgId: template.MsgId(ctx), + Data: template.Data(ctx), } err, retryFlag := sendTemplateMsg(req) if retryFlag { - log.Info("retrySendCloudbrainTemplateMsg calling") + log.Info("SendTemplateMsg calling") refreshAccessToken() err, _ = sendTemplateMsg(req) if err != nil { - log.Error("SendCloudbrainStartedMsg err. %v", err) + log.Error("SendTemplateMsg err. %v", err) return err } return nil } if err != nil { - log.Error("SendCloudbrainStartedMsg err. %v", err) + log.Error("SendTemplateMsg err. %v", err) return err } return nil } - -func getCloudbrainTemplateUrl(cloudbrain models.Cloudbrain, repo *models.Repository) string { - url := setting.AppURL + repo.FullName() - - switch cloudbrain.JobType { - case string(models.JobTypeDebug): - if cloudbrain.ComputeResource == "CPU/GPU" { - url += "/cloudbrain/" + fmt.Sprint(cloudbrain.ID) - } else { - url += "/modelarts/notebook/" + fmt.Sprint(cloudbrain.ID) - } - case string(models.JobTypeBenchmark): - url += "/cloudbrain/benchmark/" + fmt.Sprint(cloudbrain.ID) - case string(models.JobTypeTrain): - if cloudbrain.Type == models.TypeCloudBrainOne { - url += "/cloudbrain/train-job/" + fmt.Sprint(cloudbrain.JobID) - } else if cloudbrain.Type == models.TypeCloudBrainTwo { - url += "/modelarts/train-job/" + fmt.Sprint(cloudbrain.JobID) - } else if cloudbrain.Type == models.TypeC2Net { - url += "/grampus/train-job/" + fmt.Sprint(cloudbrain.JobID) - } - case string(models.JobTypeInference): - url += "/modelarts/inference-job/" + fmt.Sprint(cloudbrain.JobID) - } - return url -} - -func getCloudbrainTemplateTitle(operateType JobOperateType) string { - var title string - switch operateType { - case JobOperateTypeStart: - title = "您好，您提交的算力资源申请已通过，任务已启动，请您关注运行情况。" - case JobOperateTypeStop: - title = "您好，您提交的任务已运行结束。" - } - - return title - -} - -func getCloudbrainTemplateRemark(operateType JobOperateType) string { - var remark string - switch operateType { - case JobOperateTypeStart: - remark = "感谢您的耐心等待。" - case JobOperateTypeStop: - remark = "点击可查看运行结果" - } - - return remark - -} - -func getJobTypeDisplayName(jobType string) string { - switch jobType { - case string(models.JobTypeDebug): - return "调试任务" - case string(models.JobTypeBenchmark): - return "评测任务" - case string(models.JobTypeTrain): - return "训练任务" - case string(models.JobTypeInference): - return "推理任务" - } - return "" -} diff --git a/modules/notification/wechat/wechat.go b/modules/notification/wechat/wechat.go index f77bfe741..81c697fc9 100644 --- a/modules/notification/wechat/wechat.go +++ b/modules/notification/wechat/wechat.go @@ -8,7 +8,6 @@ import ( "code.gitea.io/gitea/models" "code.gitea.io/gitea/modules/auth/wechat" "code.gitea.io/gitea/modules/notification/base" - "code.gitea.io/gitea/modules/setting" ) type wechatNotifier struct { @@ -29,16 +28,6 @@ func (*wechatNotifier) NotifyChangeCloudbrainStatus(cloudbrain *models.Cloudbrai if operateType == "" { return } - switch operateType { - case wechat.JobOperateTypeStart: - if len(setting.CloudbrainStartedNotifyList) == 0 { - return - } - for _, v := range setting.CloudbrainStartedNotifyList { - if v == cloudbrain.JobType { - go wechat.SendCloudbrainStartedMsg(operateType, *cloudbrain) - return - } - } - } + template := wechat.GetTemplateFromOperateType(operateType) + go wechat.SendTemplateMsg(template, &wechat.TemplateContext{Cloudbrain: cloudbrain}, cloudbrain.UserID) } diff --git a/modules/setting/setting.go b/modules/setting/setting.go index 8a8a4a052..f26640ca9 100755 --- a/modules/setting/setting.go +++ b/modules/setting/setting.go @@ -586,6 +586,12 @@ var ( //wechat template msg config CloudbrainStartedTemplateId string CloudbrainStartedNotifyList []string + CloudbrainStartedTitle string + CloudbrainStartedRemark string + CloudbrainStoppedTemplateId string + CloudbrainStoppedNotifyList []string + CloudbrainStoppedTitle string + CloudbrainStoppedRemark string //nginx proxy PROXYURL string @@ -1448,6 +1454,12 @@ func NewContext() { TreePathOfSubscribe = sec.Key("SUBSCRIBE_TREE_PATH").MustString("wechat/subscribe_reply.json") CloudbrainStartedTemplateId = sec.Key("CLOUDBRAIN_STARTED_TEMPLATE_ID").MustString("") CloudbrainStartedNotifyList = strings.Split(sec.Key("CLOUDBRAIN_STARTED_NOTIFY_LIST").MustString("DEBUG"), ",") + CloudbrainStartedTitle = sec.Key("CLOUDBRAIN_STARTED_TITLE").MustString("您好，您提交的算力资源申请已通过，任务已启动，请您关注运行情况。") + CloudbrainStartedRemark = sec.Key("CLOUDBRAIN_STARTED_REMARK").MustString("感谢您的耐心等待。") + CloudbrainStoppedTemplateId = sec.Key("CLOUDBRAIN_STOPPED_TEMPLATE_ID").MustString("") + CloudbrainStoppedNotifyList = strings.Split(sec.Key("CLOUDBRAIN_STOPPED_NOTIFY_LIST").MustString("TRAIN"), ",") + CloudbrainStoppedTitle = sec.Key("CLOUDBRAIN_STOPPED_TITLE").MustString("您好，您申请的算力资源已结束使用，任务已完成运行，请您关注运行结果。") + CloudbrainStoppedRemark = sec.Key("CLOUDBRAIN_STOPPED_REMARK").MustString("感谢您的耐心等待。") SetRadarMapConfig()