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

/*

 Copyright (c) [2023] [pcm]
 [pcm-coordinator] is licensed under Mulan PSL v2.
 You can use this software according to the terms and conditions of the Mulan PSL v2.
 You may obtain a copy of Mulan PSL v2 at:
          http://license.coscl.org.cn/MulanPSL2
 THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 EITHER EXPaRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 See the Mulan PSL v2 for more details.

*/

package storeLink

import (
	"context"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/schedulers/option"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/service/collector"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/imagesservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/modelartsservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/modelarts"
	"strconv"
	"strings"
)

type ModelArtsLink struct {
	modelArtsRpc    modelartsservice.ModelArtsService
	modelArtsImgRpc imagesservice.ImagesService
	platform        string
	participantId   int64
	pageIndex       int32
	pageSize        int32
}

func NewModelArtsLink(modelArtsRpc modelartsservice.ModelArtsService, modelArtsImgRpc imagesservice.ImagesService, name string, id int64) *ModelArtsLink {
	return &ModelArtsLink{modelArtsRpc: modelArtsRpc, modelArtsImgRpc: modelArtsImgRpc, platform: name, participantId: id, pageIndex: 1, pageSize: 100}
}

func (m *ModelArtsLink) UploadImage(ctx context.Context, path string) (interface{}, error) {
	//TODO modelArts上传镜像
	return nil, nil
}

func (m *ModelArtsLink) DeleteImage(ctx context.Context, imageId string) (interface{}, error) {
	// TODO modelArts删除镜像
	return nil, nil
}

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

	return resp, nil
}

func (m *ModelArtsLink) SubmitTask(ctx context.Context, imageId string, cmd string, envs []string, params []string, resourceId string, datasetsId string, algorithmId string, aiType string) (interface{}, error) {
	// modelArts提交任务
	environments := make(map[string]string)
	parameters := make([]*modelarts.ParametersTrainJob, 0)
	for _, env := range envs {
		s := strings.Split(env, COMMA)
		environments[s[0]] = s[1]
	}
	for _, param := range params {
		s := strings.Split(param, COMMA)
		parameters = append(parameters, &modelarts.ParametersTrainJob{
			Name:  s[0],
			Value: 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,
			Parameters:   parameters,
		},
		Spec: &modelarts.SpecsC{
			Resource: &modelarts.ResourceCreateTraining{
				FlavorId:  resourceId,
				NodeCount: 1,
			},
		},
		Platform: m.platform,
	}
	resp, err := m.modelArtsRpc.CreateTrainingJob(ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

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

	return resp, nil
}

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

	return resp, nil
}

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

	return resp, nil
}

func (m *ModelArtsLink) GetResourceStats(ctx context.Context) (*collector.ResourceStats, error) {
	return nil, nil
}

func (m *ModelArtsLink) GetDatasetsSpecs(ctx context.Context) ([]*collector.DatasetsSpecs, error) {
	return nil, nil
}

func (m *ModelArtsLink) GetAlgorithms(ctx context.Context) ([]*collector.Algorithm, error) {
	return nil, nil
}

func (m *ModelArtsLink) Execute(ctx context.Context, option *option.AiOption) (interface{}, error) {
	err := m.GenerateSubmitParams(ctx, option)
	if err != nil {
		return nil, err
	}
	task, err := m.SubmitTask(ctx, option.ImageId, option.Cmd, option.Envs, option.Params, option.ResourceId, option.DatasetsId, option.AlgorithmId, option.TaskType)
	if err != nil {
		return nil, err
	}
	return task, nil
}

func (m *ModelArtsLink) GenerateSubmitParams(ctx context.Context, option *option.AiOption) error {
	err := m.generateResourceId(ctx, option)
	if err != nil {
		return err
	}
	err = m.generateImageId(option)
	if err != nil {
		return err
	}
	err = m.generateCmd(option)
	if err != nil {
		return err
	}
	err = m.generateEnv(option)
	if err != nil {
		return err
	}
	err = m.generateParams(option)
	if err != nil {
		return err
	}
	return nil
}

func (m *ModelArtsLink) generateResourceId(ctx context.Context, option *option.AiOption) error {
	_, err := m.QuerySpecs(ctx)
	if err != nil {
		return err
	}
	return nil
}

func (m *ModelArtsLink) generateImageId(option *option.AiOption) error {

	return nil
}

func (m *ModelArtsLink) generateCmd(option *option.AiOption) error {

	return nil
}

func (m *ModelArtsLink) generateEnv(option *option.AiOption) error {

	return nil
}

func (m *ModelArtsLink) generateParams(option *option.AiOption) error {

	return nil
}