pcm-coordinator/internal/participant/ai.go

package participant

import (
	"errors"
	"fmt"
	"github.com/go-resty/resty/v2"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/database"
	"net/http"
	"sync"
)

const (
	// 算法路由
	AlgorithmById       = "/ai/algorithm/get"    //根据Id查询算法列表
	AlgorithmsList      = "/ai/algorithm/list"   //所有算法列表
	AlgorithmCreateById = "/ai/algorithm/create" //根据Id创建算法

	// 数据集路由
	DatasetCreateById = "/ai/dataset/create" //根据Id创建数据集

	// 模型相关路由
	ModelCreateById = "/ai/model/create" //根据Id创建模型

	// 资源相关路由
	ResourceSpecList     = "/ai/resource/specs"      //所有资源列表，根据参数 train or infer 查询资源
	ResourceTrainingById = "/ai/resource/train/get"  //根据Id查询资源列表
	ResourceTrainingList = "/ai/resource/train/list" //所有训练资源列表

	// 任务相关路由
	TaskCreateTrain     = "/ai/task/train"
	TaskResultSync      = "/ai/task/sync"
	TaskLog             = "/ai/task/log"
	TaskTrainingDetail  = "/ai/task/train/detail"
	TaskInferenceDetail = "/ai/task/infer/detail"
)

type Ai struct {
	store  *database.AiStorage
	idAddr sync.Map
}

func New(store *database.AiStorage, adapterId string) (*Ai, error) {
	if store == nil {
		return nil, errors.New("store cannot be nil")
	}

	a := &Ai{
		store: store,
	}

	css, err := store.GetClustersByAdapterId(adapterId)
	if err != nil {
		return nil, fmt.Errorf("failed to get clusters: %w", err)
	}

	for _, info := range css.List {
		a.idAddr.Store(info.Id, info.Server)
	}

	InitClient()
	return a, nil
}

func (a *Ai) UpdateAddr(id string, addr string) {
	a.idAddr.Store(id, addr)
}

func (a *Ai) GetServerAddrById(id string) (string, bool) {
	val, ok := a.idAddr.Load(id)
	if !ok {
		return "", false
	}
	addr, ok := val.(string)
	if !ok {
		return "", false
	}
	return addr, true
}

func (a *Ai) AlgorithmById(platformId string) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+AlgorithmById, http.MethodGet, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) AlgorithmCreateById(platformId string, param *CreateParam) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+AlgorithmCreateById, http.MethodPost, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetBody(param).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) DatasetCreateById(platformId string, param *CreateParam) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+DatasetCreateById, http.MethodPost, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetBody(param).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) ModelCreateById(platformId string, param *CreateParam) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+ModelCreateById, http.MethodPost, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetBody(param).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) TaskCreateTrain(platformId string, param *TaskCreateParam) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+TaskCreateTrain, http.MethodPost, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetBody(param).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) TaskResultSync(platformId string, param *TaskResultSyncParam) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+TaskResultSync, http.MethodPost, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId": platformId,
		}).SetBody(param).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) TaskLog(platformId string, taskId string) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+TaskLog, http.MethodGet, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId":   platformId,
			"taskId": taskId,
		}).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) TaskTrainingDetail(platformId string, taskId string) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+TaskTrainingDetail, http.MethodGet, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId":   platformId,
			"taskId": taskId,
		}).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}

func (a *Ai) TaskInferenceDetail(platformId string, taskId string) (resp *Resp, err error) {
	addr, ok := a.GetServerAddrById(platformId)
	if !ok {
		return nil, fmt.Errorf("clusterId not found: %s", platformId)
	}
	respErr := &RespErr{}
	_, err = Request(addr+TaskInferenceDetail, http.MethodGet, func(req *resty.Request) {
		req.SetQueryParams(map[string]string{
			"pfId":   platformId,
			"taskId": taskId,
		}).SetError(&respErr).SetResult(&resp)
	})
	if err != nil {
		return nil, err
	}
	return
}