pcm-coordinator/internal/scheduler/service/inference/textInference/textInference.go

package textInference

import (
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/database"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/schedulers/option"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/inference"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/models"
)

type ITextInference interface {
	SaveAiTask(id int64, adapterName string) error
	UpdateStatus(aiTaskList []*models.TaskAi, adapterName string) error
	GetAiType() string
}

type FilteredCluster struct {
	urls        []*inference.InferUrl
	clusterId   string
	clusterName string
}

type TextInference struct {
	inference    ITextInference
	opt          *option.InferOption
	storage      *database.AiStorage
	inferAdapter map[string]map[string]inference.ICluster
	errMap       map[string]string
	adapterName  string
}

func New(
	inference ITextInference,
	opt *option.InferOption,
	storage *database.AiStorage,
	inferAdapter map[string]map[string]inference.ICluster,
	adapterName string) (*TextInference, error) {
	return &TextInference{
		inference:    inference,
		opt:          opt,
		storage:      storage,
		inferAdapter: inferAdapter,
		adapterName:  adapterName,
		errMap:       make(map[string]string),
	}, nil
}

func (ti *TextInference) CreateTask() (int64, error) {
	id, err := ti.saveTask()
	if err != nil {
		return 0, err
	}
	err = ti.saveAiTask(id)
	if err != nil {
		return 0, err
	}
	return id, nil
}

func (ti *TextInference) InferTask(id int64) error {
	aiTaskList, err := ti.storage.GetAiTaskListById(id)
	if err != nil || len(aiTaskList) == 0 {
		return err
	}
	err = ti.updateStatus(aiTaskList)
	if err != nil {
		return err
	}
	return nil
}

func (ti *TextInference) saveTask() (int64, error) {
	var synergystatus int64
	var strategyCode int64

	id, err := ti.storage.SaveTask(ti.opt.TaskName, strategyCode, synergystatus, ti.inference.GetAiType())
	if err != nil {
		return 0, err
	}
	return id, nil
}

func (ti *TextInference) saveAiTask(id int64) error {
	err := ti.inference.SaveAiTask(id, ti.adapterName)
	if err != nil {
		return err
	}
	return nil
}

func (ti *TextInference) updateStatus(aiTaskList []*models.TaskAi) error {
	err := ti.inference.UpdateStatus(aiTaskList, ti.adapterName)
	if err != nil {
		return err
	}
	return nil
}