pcm-coordinator/api/internal/storeLink/octopus.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/jcce-pcm/pcm-coordinator/api/internal/scheduler/service/collector"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/svc"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/jcce-pcm/pcm-participant-octopus/octopus"
	"strings"
)

type OctopusLink struct {
	ctx           context.Context
	svcCtx        *svc.ServiceContext
	pageIndex     int32
	pageSize      int32
	platform      string
	participantId int64
}

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

func NewOctopusLink(ctx context.Context, svcCtx *svc.ServiceContext, name string, id int64) *OctopusLink {
	return &OctopusLink{ctx: ctx, svcCtx: svcCtx, platform: name, participantId: id, pageIndex: 1, pageSize: 100}
}

func (o *OctopusLink) UploadImage(path string) (interface{}, error) {
	// octopus创建镜像
	createReq := &octopus.CreateImageReq{
		Platform: o.platform,
		CreateImage: &octopus.CreateImage{
			SourceType:   1,
			ImageName:    IMG_NAME_PREFIX + utils.RandomString(7),
			ImageVersion: IMG_VERSION_PREFIX + utils.RandomString(7),
		},
	}
	createResp, err := o.svcCtx.OctopusRpc.CreateImage(o.ctx, createReq)
	if err != nil {
		return nil, err
	}

	// octopus上传镜像
	uploadReq := &octopus.UploadImageReq{
		Platform: o.platform,
		ImageId:  createResp.Payload.ImageId,
		Params: &octopus.UploadImageParam{
			Domain:   "",
			FileName: "",
		},
	}
	uploadResp, err := o.svcCtx.OctopusRpc.UploadImage(o.ctx, uploadReq)
	if err != nil {
		return nil, err
	}

	// Todo 实际上传

	return uploadResp, nil
}

func (o *OctopusLink) DeleteImage(imageId string) (interface{}, error) {
	// octopus删除镜像
	req := &octopus.DeleteImageReq{
		Platform: o.platform,
		ImageId:  imageId,
	}
	resp, err := o.svcCtx.OctopusRpc.DeleteImage(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) QueryImageList() (interface{}, error) {
	// octopus获取镜像列表
	req := &octopus.GetUserImageListReq{
		Platform:  o.platform,
		PageIndex: o.pageIndex,
		PageSize:  o.pageSize,
	}
	resp, err := o.svcCtx.OctopusRpc.GetUserImageList(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) SubmitTask(imageId string, cmd string, envs []string, params []string, resourceId string) (interface{}, error) {
	// octopus提交任务

	// python参数
	var prms []*octopus.Parameters
	for _, param := range params {
		var p octopus.Parameters
		s := strings.Split(param, COMMA)
		p.Key = s[0]
		p.Value = s[1]
		prms = append(prms, &p)
	}

	//环境变量
	envMap := make(map[string]string)
	for _, env := range envs {
		s := strings.Split(env, COMMA)
		envMap[s[0]] = s[1]
	}

	req := &octopus.CreateTrainJobReq{
		Platform: o.platform,
		Params: &octopus.CreateTrainJobParam{
			ImageId:      imageId,
			Name:         TASK_NAME_PREFIX + UNDERSCORE + utils.RandomString(10),
			ResourcePool: RESOURCE_POOL,
			Config: []*octopus.Config{
				{
					Command:               cmd,
					ResourceSpecId:        resourceId,
					MinFailedTaskCount:    1,
					MinSucceededTaskCount: 1,
					TaskNumber:            1,
					Parameters:            prms,
					Envs:                  envMap,
				},
			},
		},
	}
	resp, err := o.svcCtx.OctopusRpc.CreateTrainJob(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) QueryTask(taskId string) (interface{}, error) {
	// octopus获取任务
	req := &octopus.GetTrainJobReq{
		Platform: o.platform,
		Id:       taskId,
	}
	resp, err := o.svcCtx.OctopusRpc.GetTrainJob(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) DeleteTask(taskId string) (interface{}, error) {
	// octopus删除任务
	req := &octopus.DeleteTrainJobReq{
		Platform: o.platform,
		JobIds:   []string{taskId},
	}
	resp, err := o.svcCtx.OctopusRpc.DeleteTrainJob(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) QuerySpecs() (interface{}, error) {
	// octopus查询资源规格
	req := &octopus.GetResourceSpecsReq{
		Platform:     o.platform,
		ResourcePool: "common-pool",
	}
	resp, err := o.svcCtx.OctopusRpc.GetResourceSpecs(o.ctx, req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (o *OctopusLink) GetResourceSpecs() (*collector.ResourceSpecs, error) {
	return nil, nil
}