aiforge/modules/auth/wechat/cloudbrain.go

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

func (CloudbrainStartMsg) TemplateId(ctx *TemplateContext) string {
	return setting.CloudbrainStartedTemplateId
}

type CloudbrainStopMsg struct {
}

func (CloudbrainStopMsg) Data(ctx *TemplateContext) *DefaultWechatTemplate {
	return &DefaultWechatTemplate{
		First:    TemplateValue{Value: fmt.Sprintf(setting.CloudbrainStoppedTitle, ctx.Cloudbrain.Status)},
		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")},
		Keyword4: TemplateValue{Value: time.Unix(int64(ctx.Cloudbrain.EndTime), 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 {
			if ctx.Cloudbrain.Duration > 0 && ctx.Cloudbrain.EndTime > 0 {
				return true
			}
			break
		}
	}
	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)
}

func (CloudbrainStopMsg) TemplateId(ctx *TemplateContext) string {
	return setting.CloudbrainStoppedTemplateId
}

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