pcm-coordinator/api/internal/storeLink/modelarts.go

package storeLink

import (
	"context"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/svc"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/models"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/modelarts"
	"strconv"
	"strings"
)

type ModelArtsLink struct {
	ctx         context.Context
	svcCtx      *svc.ServiceContext
	platform    string
	pageIndex   int32
	pageSize    int32
	participant *models.StorelinkCenter
}

//const (
//	IMG_NAME_PREFIX    = "oct_"
//	IMG_VERSION_PREFIX = "version_"
//	TASK_NAME_PREFIX   = "trainJob_"
//	RESOURCE_POOL      = "common-pool"
//)

func NewModelArtsLink(ctx context.Context, svcCtx *svc.ServiceContext, participant *models.StorelinkCenter) *ModelArtsLink {
	return &ModelArtsLink{ctx: ctx, svcCtx: svcCtx, participant: participant, platform: participant.Name, pageIndex: 1, pageSize: 100}
}

func (o *ModelArtsLink) UploadImage(path string) (interface{}, error) {
	//TODO modelArts上传镜像
	return nil, nil
}

func (o *ModelArtsLink) DeleteImage(imageId string) (interface{}, error) {
	// TODO modelArts删除镜像

	return nil, nil
}

func (o *ModelArtsLink) QueryImageList() (interface{}, error) {
	// modelArts获取镜像列表
	req := &modelarts.ListRepoReq{
		Offset:   "0",
		Limit:    strconv.Itoa(int(o.pageSize)),
		Platform: o.platform,
	}
	resp, err := o.svcCtx.ModelArtsImgRpc.ListReposDetails(o.ctx, req)
	if err != nil {
		return nil, err
	}

	//转换成统一返回类型
	imgListResp, err := ConvertType[modelarts.ListReposDetailsResp](resp, nil)
	if err != nil {
		return nil, err
	}

	return imgListResp, nil
}

func (o *ModelArtsLink) SubmitTask(imageId string, cmd string, params []string, resourceId string) (interface{}, error) {
	// modelArts提交任务
	environments := make(map[string]string)

	for _, param := range params {
		s := strings.Split(param, COMMA)
		environments[s[0]] = s[1]
	}

	req := &modelarts.CreateTrainingJobReq{
		Kind: "job",
		Metadata: &modelarts.MetadataS{
			Name:        TASK_NAME_PREFIX + utils.RandomString(10),
			WorkspaceId: "0",
		},
		Algorithm: &modelarts.Algorithms{
			Engine: &modelarts.EngineCreateTraining{
				ImageUrl: imageId,
			},
			Command:      cmd,
			Environments: environments,
		},
		Spec: &modelarts.SpecsC{
			Resource: &modelarts.ResourceCreateTraining{
				FlavorId:  resourceId,
				NodeCount: 1,
			},
		},
		Platform: o.platform,
	}
	resp, err := o.svcCtx.ModelArtsRpc.CreateTrainingJob(o.ctx, req)
	if err != nil {
		return nil, err
	}

	//转换成统一返回类型
	submitResp, err := ConvertType[modelarts.CreateTrainingJobResp](resp, nil)
	if err != nil {
		return nil, err
	}

	return submitResp, nil
}

func (o *ModelArtsLink) QueryTask(taskId string) (interface{}, error) {
	// 获取任务
	req := &modelarts.DetailTrainingJobsReq{
		TrainingJobId: taskId,
		Platform:      o.platform,
	}
	resp, err := o.svcCtx.ModelArtsRpc.GetTrainingJobs(o.ctx, req)
	if err != nil {
		return nil, err
	}

	//转换成统一返回类型
	taskResp, err := ConvertType[modelarts.JobResponse](resp, o.participant)
	if err != nil {
		return nil, err
	}

	return taskResp, nil
}

func (o *ModelArtsLink) DeleteTask(taskId string) (interface{}, error) {
	// 删除任务
	req := &modelarts.DeleteTrainingJobReq{
		TrainingJobId: taskId,
		Platform:      o.platform,
	}
	resp, err := o.svcCtx.ModelArtsRpc.DeleteTrainingJob(o.ctx, req)
	if err != nil {
		return nil, err
	}

	//转换成统一返回类型
	deleteResp, err := ConvertType[modelarts.DeleteTrainingJobResp](resp, nil)
	if err != nil {
		return nil, err
	}

	return deleteResp, nil
}

func (o *ModelArtsLink) QuerySpecs() (interface{}, error) {
	// octopus查询资源规格
	req := &modelarts.TrainingJobFlavorsReq{
		Platform: o.platform,
	}
	resp, err := o.svcCtx.ModelArtsRpc.GetTrainingJobFlavors(o.ctx, req)
	if err != nil {
		return nil, err
	}

	//转换成统一返回类型
	specsResp, err := ConvertType[modelarts.TrainingJobFlavorsResp](resp, o.participant)
	if err != nil {
		return nil, err
	}

	return specsResp, nil
}