pcm-coordinator/internal/logic/xjlab/task_resource_usage.go

package xjlab

import (
	"context"
	"fmt"

	"github.com/pkg/errors"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/models"
	"gorm.io/gorm"
)

type TaskResourceUsageLogic struct {
	logx.Logger
	ctx        context.Context
	svcCtx     *svc.ServiceContext
	hpcService *service.HpcService
}

func NewTaskResourceUsageLogic(ctx context.Context, svcCtx *svc.ServiceContext) *TaskResourceUsageLogic {
	cache := make(map[string]interface{}, 10)
	hpcService, err := service.NewHpcService(&svcCtx.Config, svcCtx.Scheduler.HpcStorages, cache)
	if err != nil {
		return nil
	}
	return &TaskResourceUsageLogic{
		Logger:     logx.WithContext(ctx),
		ctx:        ctx,
		svcCtx:     svcCtx,
		hpcService: hpcService,
	}
}

func (l *TaskResourceUsageLogic) TaskResourceUsage(req *types.FId) (interface{}, error) {
	task := &models.Task{}
	var resp interface{}
	if errors.Is(l.svcCtx.DbEngin.Where("id", req.Id).First(&task).Error, gorm.ErrRecordNotFound) {
		return nil, errors.New("记录不存在")
	}
	switch task.AdapterTypeDict {
	case constants.AdapterTypeCloud:
		return nil, nil
	case constants.AdapterTypeAI:
		return nil, nil
	case constants.AdapterTypeHPC:
		// 获取HPC任务的资源使用情况
		usage, err := l.GetHpcTaskResourceUsage(req)
		if err != nil {
			return nil, err
		}
		resp = usage
	}
	return resp, nil
}

type TaskHPCResult struct {
	ID        uint   `gorm:"column:id"`         // 对应 t.id
	JobID     string `gorm:"column:job_id"`     // 对应 hpc.job_id
	AdapterId string `gorm:"column:adapter_id"` // 对应 hpc.adapter_id
	ClusterId string `gorm:"column:cluster_id"` // 对应 hpc.cluster_id
}

func (l *TaskResourceUsageLogic) GetHpcTaskResourceUsage(req *types.FId) (resp interface{}, err error) {
	var hpcR TaskHPCResult
	tx := l.svcCtx.DbEngin.Raw(
		"SELECT t.id, hpc.job_id ,hpc.adapter_id ,hpc.cluster_id FROM task t "+
			"INNER JOIN task_hpc hpc ON t.id = hpc.task_id "+
			"WHERE adapter_type_dict = 2 AND t.id = ?",
		req.Id,
	).Scan(&hpcR).Error
	if tx != nil {
		return nil, fmt.Errorf("数据库查询失败: %v", tx.Error)
	}
	if hpcR.ID == 0 {
		return nil, fmt.Errorf("任务不存在")
	}

	// 获取资源使用情况
	resp, err = l.hpcService.HpcExecutorAdapterMap[hpcR.AdapterId].GetTaskResourceUsage(l.ctx, hpcR.JobID, hpcR.ClusterId)
	if err != nil {
		return nil, err
	}

	return resp, nil
}