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

package storeLink

import (
	"context"
	"github.com/pkg/errors"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/svc"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/types"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/models"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/utils/timeutils"
	"gitlink.org.cn/jcce-pcm/pcm-participant-ac/hpcAC"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/modelarts"
	"gitlink.org.cn/jcce-pcm/pcm-participant-octopus/octopus"
	"gorm.io/gorm"
	"strconv"
)

type Linkage interface {
	UploadImage(path string) (interface{}, error)
	DeleteImage(imageId string) (interface{}, error)
	QueryImageList() (interface{}, error)
	SubmitTask(imageId string, cmd string, params []string, resourceId string) (interface{}, error)
	QueryTask(taskId string) (interface{}, error)
	QuerySpecs() (interface{}, error)
	DeleteTask(taskId string) (interface{}, error)
}

const (
	COMMA               = ","
	TYPE_OCTOPUS        = "1"
	TYPE_MODELARTS      = "2"
	TYPE_SHUGUANGAI     = "3"
	OCTOPUS             = "Octopus"
	MODELARTS           = "Modelarts"
	SHUGUANGAI          = "ShuguangAi"
	DCU                 = "dcu"
	PYTORCH             = "Pytorch"
	TASK_PYTORCH_PREFIX = "PytorchTask"
	TENSORFLOW          = "Tensorflow"
	RESOURCE_GROUP      = "wzhdtest"
	WorkPath            = "/work/home/acgnnmfbwo/111111/py/"
	TimeoutLimit        = "10:00:00"
	PythonCodePath      = "/work/home/acgnnmfbwo/111111/py/test.py"
)

var (
	OctImgStatus = map[int32]string{
		1: "未上传",
		3: "制作完成",
		4: "制作失败",
	}
	AITYPE = map[string]string{
		"1": OCTOPUS,
		"2": MODELARTS,
		"3": SHUGUANGAI,
	}
)

type StoreLink struct {
	ILinkage Linkage
}

func NewStoreLink(ctx context.Context, svcCtx *svc.ServiceContext, participant *models.ScParticipantPhyInfo) *StoreLink {
	switch participant.Type {
	case TYPE_OCTOPUS:
		linkStruct := NewOctopusLink(ctx, svcCtx, participant)
		return &StoreLink{ILinkage: linkStruct}
	case TYPE_MODELARTS:
		linkStruct := NewModelArtsLink(ctx, svcCtx, participant)
		return &StoreLink{ILinkage: linkStruct}
	case TYPE_SHUGUANGAI:
		linkStruct := NewShuguangAi(ctx, svcCtx, participant)
		return &StoreLink{ILinkage: linkStruct}
	default:
		return nil
	}
}

func GetParticipants(dbEngin *gorm.DB) []*models.ScParticipantPhyInfo {
	var participants []*models.ScParticipantPhyInfo
	dbEngin.Raw("select * from sc_participant_phy_info where type = 1").Scan(&participants)
	return participants
}

func GetParticipantById(partId int64, dbEngin *gorm.DB) *models.ScParticipantPhyInfo {
	var participant models.ScParticipantPhyInfo
	dbEngin.Raw("select * from sc_participant_phy_info where id = ?", partId).Scan(&participant)
	return &participant
}

func ConvertType[T any](in *T, participant *models.ScParticipantPhyInfo) (interface{}, error) {

	switch (interface{})(in).(type) {
	case *octopus.UploadImageResp:
		var resp types.UploadLinkImageResp
		inresp := (interface{})(in).(*octopus.UploadImageResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			return resp, nil
		}

		return resp, nil
	case *octopus.DeleteImageResp:
		var resp types.DeleteLinkImageResp
		inresp := (interface{})(in).(*octopus.DeleteImageResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			return resp, nil
		}

		return resp, nil

	case *octopus.GetUserImageListResp:
		var resp types.GetLinkImageListResp
		inresp := (interface{})(in).(*octopus.GetUserImageListResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			resp.Images = nil
			return resp, nil
		}

		for _, v := range inresp.Payload.Images {
			var image types.ImageSl
			image.ImageId = v.Image.Id
			image.ImageName = v.Image.ImageName
			image.ImageStatus = OctImgStatus[v.Image.ImageStatus]
			resp.Images = append(resp.Images, &image)
		}
		return resp, nil
	case *modelarts.ListReposDetailsResp:
		var resp types.GetLinkImageListResp
		inresp := (interface{})(in).(*modelarts.ListReposDetailsResp)

		if inresp.Errors != nil {
			resp.Success = false
			resp.ErrorMsg = inresp.Errors[0].ErrorMessage
			resp.Images = nil
			return resp, nil
		}

		for _, v := range inresp.Items {
			for _, r := range v.Tags {
				var image types.ImageSl
				image.ImageId = v.Namespace + "/" + v.Name + ":" + r
				image.ImageName = v.Name
				image.ImageStatus = "created"
				resp.Images = append(resp.Images, &image)
			}
		}
		return resp, nil
	case *hpcAC.GetImageListAiResp:
		var resp types.GetLinkImageListResp
		inresp := (interface{})(in).(*hpcAC.GetImageListAiResp)

		if inresp.Code == "0" {
			resp.Success = true
			for _, img := range inresp.Data {
				var image types.ImageSl
				image.ImageId = img.ImageId
				image.ImageName = img.Name
				image.ImageStatus = "created"
				resp.Images = append(resp.Images, &image)
			}
		} else {
			resp.Success = false
			resp.ErrorMsg = inresp.Msg
			resp.Images = nil
		}
		return resp, nil

	case *octopus.CreateTrainJobResp:
		var resp types.SubmitLinkTaskResp
		inresp := (interface{})(in).(*octopus.CreateTrainJobResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			return resp, nil
		}

		resp.TaskId = inresp.Payload.JobId

		return resp, nil
	case *modelarts.CreateTrainingJobResp:
		var resp types.SubmitLinkTaskResp
		inresp := (interface{})(in).(*modelarts.CreateTrainingJobResp)

		if inresp.ErrorMsg != "" {
			resp.ErrorMsg = inresp.ErrorMsg
			resp.Success = false
			return resp, nil
		}

		resp.TaskId = inresp.Metadata.Id

		return resp, nil
	case *hpcAC.SubmitTaskAiResp:
		var resp types.SubmitLinkTaskResp
		inresp := (interface{})(in).(*hpcAC.SubmitTaskAiResp)

		if inresp.Code == "0" {
			resp.Success = true
			resp.TaskId = inresp.Data
		} else {
			resp.Success = false
			resp.ErrorMsg = inresp.Msg
		}
		return resp, nil
	case *octopus.GetTrainJobResp:
		var resp types.GetLinkTaskResp
		inresp := (interface{})(in).(*octopus.GetTrainJobResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			return resp, nil
		}

		resp.Task.TaskId = inresp.Payload.TrainJob.Id
		resp.Task.TaskName = inresp.Payload.TrainJob.Name
		resp.Task.StartedAt = inresp.Payload.TrainJob.StartedAt
		resp.Task.CompletedAt = inresp.Payload.TrainJob.CompletedAt
		resp.Task.TaskStatus = inresp.Payload.TrainJob.Status

		return resp, nil
	case *modelarts.JobResponse:
		var resp types.GetLinkTaskResp
		inresp := (interface{})(in).(*modelarts.JobResponse)
		if inresp.ErrorMsg != "" {
			resp.ErrorMsg = inresp.ErrorMsg
			resp.Success = false
			return resp, nil
		}

		resp.Task.TaskId = inresp.Metadata.Id
		resp.Task.TaskName = inresp.Metadata.Name
		resp.Task.StartedAt = int64(inresp.Status.StartTime)
		resp.Task.CompletedAt = int64(inresp.Status.Duration)
		resp.Task.TaskStatus = inresp.Status.Phase

		return resp, nil
	case *hpcAC.GetPytorchTaskResp:
		var resp types.GetLinkTaskResp
		inresp := (interface{})(in).(*hpcAC.GetPytorchTaskResp)
		if inresp.Code == "0" {
			resp.Success = true
			resp.Task.TaskId = inresp.Data.Id
			resp.Task.TaskName = inresp.Data.TaskName
			resp.Task.TaskStatus = inresp.Data.Status
			resp.Task.StartedAt = timeutils.StringToUnixTime(inresp.Data.StartTime)
			resp.Task.CompletedAt = timeutils.StringToUnixTime(inresp.Data.EndTime)
		} else {
			resp.Success = false
			resp.ErrorMsg = inresp.Msg
			resp.Task = nil
		}

		return resp, nil
	case *octopus.DeleteTrainJobResp:
		var resp types.DeleteLinkTaskResp
		inresp := (interface{})(in).(*octopus.DeleteTrainJobResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ErrorMsg = inresp.Error.Message
			return resp, nil
		}

		return resp, nil
	case *modelarts.DeleteTrainingJobResp:
		var resp types.DeleteLinkTaskResp
		inresp := (interface{})(in).(*modelarts.DeleteTrainingJobResp)
		if inresp.ErrorMsg != "" {
			resp.ErrorMsg = inresp.ErrorMsg
			resp.Success = false
			return resp, nil
		}
		return resp, nil
	case *hpcAC.DeleteTaskAiResp:
		var resp types.DeleteLinkTaskResp
		inresp := (interface{})(in).(*hpcAC.DeleteTaskAiResp)
		if inresp.Code == "0" {
			resp.Success = true
		} else {
			resp.Success = false
			resp.ErrorMsg = inresp.Msg
		}
		return resp, nil
	case *octopus.GetResourceSpecsResp:
		var resp types.GetResourceSpecsResp
		inresp := (interface{})(in).(*octopus.GetResourceSpecsResp)
		resp.Success = inresp.Success
		if !resp.Success {
			resp.ResourceSpecs = nil
			return resp, nil
		}

		for _, spec := range inresp.TrainResourceSpecs {
			var respec types.ResourceSpecSl
			respec.SpecId = spec.Id
			respec.SpecName = spec.Name
			respec.ParticipantId = strconv.FormatInt(participant.Id, 10)
			respec.ParticipantName = participant.Name
			respec.SpecPrice = spec.Price
			resp.ResourceSpecs = append(resp.ResourceSpecs, &respec)
		}

		return resp, nil
	case *hpcAC.GetResourceSpecResp:
		var resp types.GetResourceSpecsResp
		inresp := (interface{})(in).(*hpcAC.GetResourceSpecResp)

		if inresp.Code != "0" {
			resp.Success = false
			resp.ResourceSpecs = nil
		} else {
			var spec types.ResourceSpecSl
			resp.Success = true
			spec.ParticipantName = participant.Name
			spec.ParticipantId = strconv.FormatInt(participant.Id, 10)
			spec.SpecName = SHUGUANGAI_CUSTOM_RESOURCE_NAME
			spec.SpecId = SHUGUANGAI_CUSTOM_RESOURCE_ID
			resp.ResourceSpecs = append(resp.ResourceSpecs, &spec)
		}
		return resp, nil
	case *modelarts.TrainingJobFlavorsResp:
		var resp types.GetResourceSpecsResp
		resp.Success = true
		inresp := (interface{})(in).(*modelarts.TrainingJobFlavorsResp)
		if inresp.Flavors == nil {
			resp.Success = false
			resp.ResourceSpecs = nil
			return resp, nil
		}

		for _, spec := range inresp.Flavors {
			var respec types.ResourceSpecSl
			respec.SpecId = spec.FlavorId
			respec.SpecName = spec.FlavorName
			respec.ParticipantId = strconv.FormatInt(participant.Id, 10)
			respec.ParticipantName = participant.Name
			respec.SpecPrice = 0
			resp.ResourceSpecs = append(resp.ResourceSpecs, &respec)
		}

		return resp, nil
	default:
		return nil, errors.New("type convert fail")
	}
}