red-future/model-gateway
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gatway

Browse Source
This commit is contained in:
王利钊
2026-05-12 13:45:08 +08:00
parent e81df5ce5a
commit 37d3461983
38 changed files with 1721 additions and 1113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
service/worker.go
View File

@@ -5,12 +5,14 @@ import (
"fmt"
"strings"
"time"
"unicode/utf8"
"model-asynch/dao"
"model-asynch/model/entity"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/grpool"
"github.com/tidwall/gjson"
)
var AsyncWorker = &asyncWorker{}
@@ -43,7 +45,7 @@ func (w *asyncWorker) RunOnce(ctx context.Context, batchSize, goroutines int) (c
for _, t := range tasks {
task := t
_ = pool.AddWithRecover(ctx, func(ctx context.Context) {
w.handleOne(ctx, task)
w.handleOne(ctx, task, 0)
done <- struct{}{}
}, func(ctx context.Context, e error) {
if e != nil {
@@ -59,8 +61,23 @@ func (w *asyncWorker) RunOnce(ctx context.Context, batchSize, goroutines int) (c
return claimed, nil
}
func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
// 从任务入库的 request_payload 里恢复 payload + headers，给 OSS 上传透传鉴权用
// RunByTaskID 创建任务后立即异步尝试执行当前任务：
// - 只定向抢占当前 taskId 对应的 pending 任务
// - 若任务已被其它 worker 抢走/已不在 pending，则直接返回
func (w *asyncWorker) RunByTaskID(ctx context.Context, taskID string, epicycleId int64) error {
task, err := dao.Task.ClaimPendingByTaskIDGlobal(ctx, taskID)
if err != nil {
return err
}
if task == nil {
return nil
}
w.handleOne(ctx, task, epicycleId)
return nil
}
func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask, epicycleId int64) {
// 从任务入库的 request_payload 里恢复 payload + headers
payload, headers := parseStoredPayload(t.RequestPayload)
if len(headers) > 0 {
ctx = setTaskHeadersToCtx(ctx, headers)
@@ -71,26 +88,42 @@ func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
if err != nil {
_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// ============ 失败回调 ============
t.State = 3
t.ErrorMsg = err.Error()
go triggerCallback(context.WithoutCancel(ctx), t)
// ================================
return
}
if m == nil || m.Enabled != 1 {
_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, "模型不存在或未启用")
errMsg := "模型不存在或未启用"
_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, errMsg)
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// ============ 失败回调 ============
t.State = 3
t.ErrorMsg = errMsg
go triggerCallback(context.WithoutCancel(ctx), t)
// ================================
return
}
// 2) 分布式并发限制（按 model_name 全局维度）
// 2) 分布式并发限制
semKey := fmt.Sprintf("asynch:sem:%s", t.ModelName)
leaseSeconds := int64(3600) // 兜底1小时
leaseSeconds := int64(3600)
maxC := GetRuntimeMaxConcurrency(ctx, t.ModelName, m.MaxConcurrency)
acquired, err := acquireSemaphore(ctx, semKey, maxC, leaseSeconds)
if err != nil {
_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// ============ 失败回调 ============
t.State = 3
t.ErrorMsg = err.Error()
go triggerCallback(context.WithoutCancel(ctx), t)
// ================================
return
}
if !acquired {
// 并发满了：放回排队（重新置回 state=0），下一轮再抢占
// 并发满了：放回排队，不回调（不是失败）
_ = w.rollbackToPending(ctx, t.Id)
return
}
@@ -109,30 +142,40 @@ func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
data []byte
contentType string
ext string
textResult string
)
// phase=1 表示模型已成功但 OSS 上传失败：优先从临时文件加载，避免重复跑模型
// phase=1 表示模型已成功但 OSS 上传失败：优先从临时文件加载
if t.Phase == 1 && strings.TrimSpace(t.TmpFile) != "" {
data, err = loadTmpResult(t.TmpFile)
if err == nil && len(data) > 0 {
contentType, ext = DetectFileType(data)
} else {
// 临时文件不可用：回退重新调用模型
data = nil
}
}
if data == nil {
// 统计：仅在真正请求模型时 +1（OSS 重试不计入）
// 统计
_ = dao.Stat.IncRequestCount(ctx, time.Now(), int64(t.TenantId), t.Creator, t.ModelName)
// 核心调用
data, err = InvokeModel(ctx, m, payload, t.ModelKey)
if err != nil {
_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// ============ 失败回调 ============
t.State = 3
t.ErrorMsg = err.Error()
go triggerCallback(context.WithoutCancel(ctx), t)
// ================================
return
}
contentType, ext = DetectFileType(data)
// 将模型输出写入临时文件，后续若 OSS 失败可只重试 OSS
if utf8.Valid(data) && (strings.HasPrefix(contentType, "text/") || contentType == "application/json") {
textResult = string(data)
if len(textResult) > 20000 {
textResult = textResult[:20000]
}
}
tmpPath, err := saveTmpResult(t.TaskID, data, ext)
if err == nil && tmpPath != "" {
t.TmpFile = tmpPath
@@ -147,26 +190,46 @@ func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
// OSS 阶段失败：保留临时文件，下一轮仅重试 OSS
_ = dao.Task.UpdateFailedKeepTmpGlobal(ctx, t.Id, err.Error())
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// ============ OSS失败不回调（还会重试） ============
// 注意：OSS失败保留临时文件，下次重试，所以这里不触发最终回调
// 如果已经重试多次还没成功，需要在任务超时或超过最大重试次数时才回调失败
return
}
// 5) 更新任务状态成功
// 注意：expire_at 的计算改为“已下载(state=4)后开始计时”，因此成功(state=2)不写 expire_at。
fileType := strings.TrimPrefix(ext, ".")
if fileType == "" {
fileType = contentType
}
if err := dao.Task.UpdateSuccessGlobal(ctx, t.Id, ossURL, fileType, int64(len(data)), nil); err != nil {
if err := dao.Task.UpdateSuccessGlobal(
ctx,
t.Id,
ossURL,
fileType,
textResult,
int64(len(data)),
nil,
GetExpendTokens(m.TokenMapping, textResult),
); err != nil {
g.Log().Errorf(ctx, "[worker] update success failed: %v", err)
return
}
// 成功/失败均不再占用 queue_limit（state=0/1 才占用）
// 成功/失败均不再占用 queue_limit
ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
// 6) 成功回调（不影响主流程）
// 6) 成功回调
t.State = 2
t.OssFile = ossURL
t.FileType = fileType
go triggerSuccessCallback(context.WithoutCancel(ctx), t)
t.TextResult = textResult
g.Log().Infof(ctx, "[CALLBACK][DISPATCH] taskId=%s bizName=%s callbackUrl=%s", t.TaskID, t.BizName, t.CallbackURL)
go triggerCallback(context.WithoutCancel(ctx), t)
// ============ 如果有 epicycleId，也触发业务回调 ============
if epicycleId != 0 {
go triggerPromptsCallback(context.WithoutCancel(ctx), t, epicycleId)
}
// 成功后清理临时文件
deleteTmpResult(t.TmpFile)
}
@@ -174,3 +237,13 @@ func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
func (w *asyncWorker) rollbackToPending(ctx context.Context, id int64) error {
return dao.Task.RollbackToPendingGlobal(ctx, id)
}
// GetExpendTokens 根据映射路径从 textResult 中提取消耗 token 值
func GetExpendTokens(tokenMapping string, textResult string) int {
value := gjson.Get(textResult, tokenMapping)
if value.Exists() {
return int(value.Int())
} else {
return len(textResult)
}
}