pcm-coordinator/api/internal/logic/ai/getcentertasklistlogic.go

package ai

import (
	"context"
	"errors"
	"fmt"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"
	"strconv"
	"sync"
	"time"

	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/types"

	"github.com/zeromicro/go-zero/core/logx"
)

type GetCenterTaskListLogic struct {
	logx.Logger
	ctx    context.Context
	svcCtx *svc.ServiceContext
}

func NewGetCenterTaskListLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetCenterTaskListLogic {
	return &GetCenterTaskListLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: svcCtx,
	}
}

func (l *GetCenterTaskListLogic) GetCenterTaskList() (resp *types.CenterTaskListResp, err error) {
	resp = &types.CenterTaskListResp{}
	var mu sync.RWMutex
	ch := make(chan struct{})

	adapterList, err := l.svcCtx.Scheduler.AiStorages.GetAdaptersByType("1")
	if err != nil {
		return nil, err
	}

	go l.updateAiTaskStatus(&mu, ch, adapterList)

	for _, adapter := range adapterList {
		mu.RLock()
		taskList, err := l.svcCtx.Scheduler.AiStorages.GetAiTasksByAdapterId(adapter.Id)
		mu.RUnlock()
		if err != nil {
			continue
		}
		if len(taskList) == 0 {
			continue
		}
		for _, task := range taskList {
			var elapsed time.Duration
			switch task.Status {
			case constants.Completed:
				end, err := time.ParseInLocation(constants.Layout, task.EndTime, time.Local)
				if err != nil {
					elapsed = time.Duration(0)
				}
				start, err := time.ParseInLocation(constants.Layout, task.StartTime, time.Local)
				if err != nil {
					elapsed = time.Duration(0)
				}
				elapsed = end.Sub(start)
			case constants.Running:
				elapsed = time.Now().Sub(task.CommitTime)
			default:
				elapsed = 0
			}

			t := &types.AiTask{
				Name:        task.Name,
				Status:      task.Status,
				Cluster:     task.ClusterName,
				Card:        task.Card,
				TimeElapsed: int32(elapsed.Seconds()),
			}
			resp.List = append(resp.List, t)
		}
	}

	select {
	case _ = <-ch:
		return resp, nil
	case <-time.After(2 * time.Second):
		return resp, nil
	}
}

func (l *GetCenterTaskListLogic) updateAiTaskStatus(mu *sync.RWMutex, ch chan<- struct{}, list []*types.AdapterInfo) {
	var wg sync.WaitGroup
	for _, adapter := range list {
		taskList, err := l.svcCtx.Scheduler.AiStorages.GetAiTasksByAdapterId(adapter.Id)
		if err != nil {
			continue
		}
		if len(taskList) == 0 {
			continue
		}
		for _, task := range taskList {
			t := task
			if t.Status == constants.Completed {
				continue
			}
			wg.Add(1)
			go func() {
				trainingTask, err := l.svcCtx.Scheduler.AiService.AiCollectorAdapterMap[adapter.Id][strconv.FormatInt(t.ClusterId, 10)].GetTrainingTask(l.ctx, t.JobId)
				if err != nil {
					msg := fmt.Sprintf("AiTaskId: %v, clusterId: %v , JobId: %v,  error: %v \n", t.Id, t.ClusterId, t.JobId, err.Error())
					logx.Errorf(errors.New(msg).Error())
					wg.Done()
					return
				}
				t.Status = trainingTask.Status
				t.StartTime = trainingTask.Start
				t.EndTime = trainingTask.End
				mu.Lock()
				err = l.svcCtx.Scheduler.AiStorages.UpdateAiTask(t)
				mu.Unlock()
				if err != nil {
					wg.Done()
					return
				}
				wg.Done()
			}()
		}
	}
	wg.Wait()
	ch <- struct{}{}
}