forked from JointCloud/pcm-coordinator
80 lines
1.9 KiB
Go
80 lines
1.9 KiB
Go
package schedule
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/utils/status"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"
|
|
"strconv"
|
|
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
)
|
|
|
|
type ScheduleCancelTaskLogic struct {
|
|
logx.Logger
|
|
ctx context.Context
|
|
svcCtx *svc.ServiceContext
|
|
}
|
|
|
|
func NewScheduleCancelTaskLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ScheduleCancelTaskLogic {
|
|
return &ScheduleCancelTaskLogic{
|
|
Logger: logx.WithContext(ctx),
|
|
ctx: ctx,
|
|
svcCtx: svcCtx,
|
|
}
|
|
}
|
|
|
|
func (l *ScheduleCancelTaskLogic) ScheduleCancelTask(req *types.CancelTaskReq) (resp *types.CancelTaskResp, err error) {
|
|
task, err := l.svcCtx.Scheduler.AiService.Storage.GetTaskById(req.TaskId)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if task == nil {
|
|
return nil, errors.New("failed to cancel task, task not found")
|
|
}
|
|
|
|
// find ai tasks
|
|
aitasks, err := l.svcCtx.Scheduler.AiStorages.GetAiTaskListById(task.Id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if len(aitasks) == 0 {
|
|
return nil, errors.New("failed to cancel task, ai sub tasks have not been created")
|
|
}
|
|
|
|
// update status
|
|
status.UpdateAiTask(l.svcCtx, aitasks...)
|
|
|
|
t := aitasks[0]
|
|
|
|
if t.Status != constants.Running {
|
|
return nil, fmt.Errorf("failed to cancel task, ai sub tasks is %s", t.Status)
|
|
}
|
|
|
|
// assume a task has only one sub ai task
|
|
err = l.svcCtx.Scheduler.AiService.AiExecutorAdapterMap[strconv.FormatInt(t.AdapterId, 10)][strconv.FormatInt(t.ClusterId, 10)].Stop(l.ctx, t.JobId)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
t.Status = constants.Cancelled
|
|
err = l.svcCtx.Scheduler.AiService.Storage.UpdateAiTask(t)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
task.Status = constants.Cancelled
|
|
err = l.svcCtx.Scheduler.AiService.Storage.UpdateTaskByModel(task)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return resp, nil
|
|
}
|