xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-08 18:23:45 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

auto restore user's power for failure tasks

Browse Source
This commit is contained in:
RockYang
2024-11-11 18:12:35 +08:00
parent ce8a2d0222
commit 2aa0b51c09
10 changed files with 105 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
api/service/dalle/service.go
View File

@@ -237,13 +237,9 @@ func (s *Service) CheckTaskStatus() {
go func() {
logger.Info("Running DALL-E task status checking ...")
for {
// 检查未完成任务进度
var jobs []model.DallJob
res := s.db.Where("progress < ?", 100).Find(&jobs)
if res.Error != nil {
time.Sleep(5 * time.Second)
continue
}
s.db.Where("progress < ?", 100).Find(&jobs)
for _, job := range jobs {
// 超时的任务标记为失败
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
@@ -252,6 +248,21 @@ func (s *Service) CheckTaskStatus() {
s.db.Updates(&job)
}
}
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "dall-e-3",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
continue
}
// 更新任务状态
s.db.Model(&job).UpdateColumn("power", 0)
}
time.Sleep(time.Second * 10)
}
}()

19
api/service/mj/service.go
View File

@@ -30,10 +30,11 @@ type Service struct {
db *gorm.DB
wsService *service.WebsocketService
uploaderManager *oss.UploaderManager
userService *service.UserService
clientIds map[uint]string
}
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager, wsService *service.WebsocketService) *Service {
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
db: db,
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
@@ -42,6 +43,7 @@ func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *os
wsService: wsService,
uploaderManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -313,6 +315,21 @@ func (s *Service) SyncTaskProgress() {
}
}
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "mid-journey",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
continue
}
// 更新任务状态
s.db.Model(&job).UpdateColumn("power", 0)
}
time.Sleep(time.Second * 5)
}
}()

19
api/service/sd/service.go
View File

@@ -34,9 +34,10 @@ type Service struct {
db *gorm.DB
uploadManager *oss.UploaderManager
wsService *service.WebsocketService
userService *service.UserService
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelDB, redisCli *redis.Client, wsService *service.WebsocketService) *Service {
func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelDB, redisCli *redis.Client, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
httpClient: req.C(),
taskQueue: store.NewRedisQueue("StableDiffusion_Task_Queue", redisCli),
@@ -44,6 +45,7 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, levelDB *store.LevelD
db: db,
wsService: wsService,
uploadManager: manager,
userService: userService,
}
}
@@ -301,6 +303,21 @@ func (s *Service) CheckTaskStatus() {
s.db.Updates(&job)
}
}
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "stable-diffusion",
Remark: fmt.Sprintf("任务失败，退回算力。任务ID：%d， Err: %s", job.Id, job.ErrMsg),
})
if err != nil {
continue
}
// 更新任务状态
s.db.Model(&job).UpdateColumn("power", 0)
}
time.Sleep(time.Second * 5)
}
}()

18
api/service/suno/service.go
View File

@@ -36,9 +36,10 @@ type Service struct {
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[string]string
userService *service.UserService
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService) *Service {
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
@@ -47,6 +48,7 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
uploadManager: manager,
wsService: wsService,
clientIds: map[string]string{},
userService: userService,
}
}
@@ -384,6 +386,20 @@ func (s *Service) SyncTaskProgress() {
}
}
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: job.ModelName,
Remark: fmt.Sprintf("Suno 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
if err != nil {
continue
}
// 更新任务状态
s.db.Model(&job).UpdateColumn("power", 0)
}
time.Sleep(time.Second * 10)
}
}()

18
api/service/video/luma.go
View File

@@ -36,9 +36,10 @@ type Service struct {
notifyQueue *store.RedisQueue
wsService *service.WebsocketService
clientIds map[uint]string
userService *service.UserService
}
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService) *Service {
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, wsService *service.WebsocketService, userService *service.UserService) *Service {
return &Service{
httpClient: req.C().SetTimeout(time.Minute * 3),
db: db,
@@ -47,6 +48,7 @@ func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Clien
wsService: wsService,
uploadManager: manager,
clientIds: map[uint]string{},
userService: userService,
}
}
@@ -286,6 +288,20 @@ func (s *Service) SyncTaskProgress() {
}
// 找出失败的任务，并恢复其扣减算力
s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
for _, job := range jobs {
err := s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
Type: types.PowerRefund,
Model: "luma",
Remark: fmt.Sprintf("Luma 任务失败，退回算力。任务ID：%s，Err:%s", job.TaskId, job.ErrMsg),
})
if err != nil {
continue
}
// 更新任务状态
s.db.Model(&job).UpdateColumn("power", 0)
}
time.Sleep(time.Second * 10)
}
}()