Merge pull request '训练任务停止时发送微信通知' (#2695) from fix-2631 into V20220815

Reviewed-on: https://git.openi.org.cn/OpenI/aiforge/pulls/2695
Reviewed-by: lewis <747342561@qq.com>

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"`

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 ""
+}

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 ""
-}

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)
 }

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()