common/sdks/pcmscheduler/jobset.go

package sch

import (
	"fmt"
	schsdk "gitlink.org.cn/cloudream/common/sdks/scheduler"
	cdssdk "gitlink.org.cn/cloudream/common/sdks/storage"
	"gitlink.org.cn/cloudream/common/utils/http2"
	"gitlink.org.cn/cloudream/common/utils/serder"
	"net/url"
	"strings"
)

type GetClusterInfoReq struct {
	Type string             `json:"type"`
	IDs  []schsdk.ClusterID `json:"clusterIDs"`
}

func (c *Client) GetClusterInfo(req GetClusterInfoReq, token string) ([]ClusterDetail, error) {
	targetUrl, err := url.JoinPath(c.baseURL, "schedule/queryResources")
	if err != nil {
		return nil, err
	}

	resourceType := Train
	if req.Type != "" {
		resourceType = req.Type
	}
	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		Body: GetClusterInfoReq{
			Type: resourceType,
		},
		Header: map[string]string{
			"Authorization": token,
		},
	})
	if err != nil {
		return nil, err
	}

	//all, err := io.ReadAll(resp.Body)
	//str := string(all)
	//println(str)

	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {

		var codeResp response[[]ClusterDetail]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return nil, fmt.Errorf("parsing response: %w", err)
		}

		// 该接口传入参数后查询效率很低，所以需要在这里特殊处理
		if codeResp.Code == ResponseCodeOK {
			var results []ClusterDetail
			for _, cluster := range codeResp.Data {
				for _, id := range req.IDs {
					if cluster.ClusterId == id {
						results = append(results, cluster)
					}
				}
			}
			return results, nil
		}

		return nil, codeResp.ToError()
	}

	return nil, fmt.Errorf("unknow response content type: %s", contType)
}

type CreateInferenceJobResp struct {
	TaskId string `json:"taskId"`
}

type CreateAIJobReq struct {
	Name            string              `json:"name"`
	Description     string              `json:"description"`
	JobResources    schsdk.JobResources `json:"jobResources"`
	DataDistributes DataDistribute      `json:"dataDistributes"`
}

type CommonJsonData struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type DataDistribute struct {
	Dataset []DatasetDistribute `json:"dataset"`
	Code    []CodeDistribute    `json:"code"`
	Image   []ImageDistribute   `json:"image"`
	Model   []ModelDistribute   `json:"model"`
}

type DataDetail struct {
	ClusterID schsdk.ClusterID `json:"clusterID"`
	//StorageID cdssdk.StorageID `json:"storageID"`
	StorageID cdssdk.StorageID
	JsonData  string `json:"jsonData"`
}

type DatasetDistribute struct {
	DataName  string           `json:"dataName"`
	PackageID cdssdk.PackageID `json:"packageID"`
	Clusters  []DataDetail     `json:"clusters"`
}

type CodeDistribute struct {
	DataName  string           `json:"dataName"`
	PackageID cdssdk.PackageID `json:"packageID"`
	Output    string           `json:"output"`
	Clusters  []DataDetail     `json:"clusters"`
}

type ImageDistribute struct {
	DataName string `json:"dataName"`
	//PackageID cdssdk.PackageID `json:"packageID"`
	ImageID  schsdk.ImageID `json:"packageID"`
	Clusters []DataDetail   `json:"clusters"`
}

type ModelDistribute struct {
	DataName  string           `json:"dataName"`
	PackageID cdssdk.PackageID `json:"packageID"`
	Clusters  []DataDetail     `json:"clusters"`
}

type CreateJobResp struct {
	TaskID        TaskID         `json:"taskID"`
	TaskName      string         `json:"taskName"`
	ScheduleDatas []ScheduleData `json:"scheduleDatas"`
}

type ScheduleData struct {
	DataType    string             `json:"dataType"`
	PackageID   cdssdk.PackageID   `json:"packageID"`
	StorageType string             `json:"storageType"`
	ClusterIDs  []schsdk.ClusterID `json:"clusterIDs"`
}

func (c *Client) CreateInferenceJob(req CreateAIJobReq, token string) (*CreateInferenceJobResp, error) {
	targetUrl, err := url.JoinPath(c.baseURL, "inference/createTask")
	if err != nil {
		return nil, err
	}

	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		Body: req,
		Header: map[string]string{
			"Authorization": token,
		},
	})
	if err != nil {
		return nil, err
	}

	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {
		var codeResp response2[CreateInferenceJobResp]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return nil, fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return &codeResp.Data, nil
		}

		return nil, fmt.Errorf("error: %s", codeResp.Message)
	}

	return nil, fmt.Errorf("unknow response content type: %s", contType)
}

type StopInferenceJobReq struct {
	AdapterID  string           `json:"adapterId"`
	ClusterID  schsdk.ClusterID `json:"clusterId"`
	ID         string           `json:"id"`
	InstanceID string           `json:"instanceId"`
}

func (c *Client) StopInferenceJob(req StopInferenceJobReq, token string) error {
	targetUrl, err := url.JoinPath(c.baseURL, "inference/stopDeployInstance")
	if err != nil {
		return err
	}

	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		Body: req,
		Header: map[string]string{
			"Authorization": token,
		},
	})
	if err != nil {
		return err
	}

	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {
		var codeResp response2[CreateInferenceJobResp]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return nil
		}

		return fmt.Errorf("error: %s", codeResp.Message)
	}

	return fmt.Errorf("unknow response content type: %s", contType)
}

func (c *Client) CreateJob(req CreateAIJobReq, token string) (*CreateJobResp, error) {
	targetUrl, err := url.JoinPath(c.baseURL, "schedule/createTask")
	if err != nil {
		return nil, err
	}

	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		Body: req,
		Header: map[string]string{
			"Authorization": token,
		},
	})
	if err != nil {
		return nil, err
	}

	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {
		var codeResp response2[CreateJobResp]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return nil, fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return &codeResp.Data, nil
		}

		return nil, fmt.Errorf("error: %s", codeResp.Message)
	}

	return nil, fmt.Errorf("unknow response content type: %s", contType)

}

type RunJobReq struct {
	TaskID         TaskID                `json:"taskID"`
	ScheduledDatas []DataScheduleResults `json:"scheduledDatas"`
}

type DataScheduleResult struct {
	Clusters        []DataDetail     `json:"clusters"`
	PackageID       cdssdk.PackageID `json:"packageID"`
	PackageFullPath string           `json:"packageFullPath"`
	Status          bool             `json:"status"`
	Msg             string           `json:"msg"`
}

type DataScheduleResults struct {
	DataType string               `json:"dataType"`
	Results  []DataScheduleResult `json:"results"`
}

func (c *Client) RunJob(req RunJobReq, token string) error {
	targetUrl, err := url.JoinPath(c.baseURL, "schedule/runTask")
	if err != nil {
		return err
	}

	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		Body: req,
		Header: map[string]string{
			"Authorization": token,
		},
	})
	if err != nil {
		return err
	}

	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {
		var codeResp response2[string]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return nil
		}

		return fmt.Errorf("error: %s", codeResp.Message)
	}

	return fmt.Errorf("unknow response content type: %s", contType)

}

type CancelJobReq struct {
	TaskID TaskID `json:"taskID"`
	Msg    string `json:"msg"`
}

func (c *Client) CancelJob(req CancelJobReq) error {
	targetUrl, err := url.JoinPath(c.baseURL, "schedule/queryResources")
	if err != nil {
		return err
	}
	resp, err := http2.GetJSON(targetUrl, http2.RequestParam{Body: req})
	if err != nil {
		return err
	}
	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {

		var codeResp response[string]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return nil
		}

		return codeResp.ToError()
	}

	return fmt.Errorf("unknow response content type: %s", contType)
}

type GetResourceSpecReq struct {
	ClusterID schsdk.ClusterID `form:"clusterId"`
	Tag       string           `form:"tag"`
	PageNum   int64            `form:"pageNum"`
	PageSize  int64            `form:"pageSize"`
	Status    string           `form:"status"`
}

type GetResourceSpecResp struct {
	Total    int64            `json:"total"`
	PageNum  int64            `json:"pageNum"`
	PageSize int64            `json:"pageSize"`
	List     []ResourceDetail `json:"list"`
}

func (c *Client) GetResourceSpec(req GetResourceSpecReq, token string) ([]ResourceDetail, int64, error) {
	targetUrl, err := url.JoinPath(c.baseURL, "core/ai/resourceSpec/page")
	if err != nil {
		return nil, 0, err
	}

	req.Status = "1"
	resp, err := http2.GetJSON(targetUrl, http2.RequestParam{
		Query: req,
		Header: map[string]string{
			"Authorization": token,
		},
	})

	//all, err := io.ReadAll(resp.Body)
	//println(string(all))

	if err != nil {
		return nil, 0, err
	}
	contType := resp.Header.Get("Content-Type")
	if strings.Contains(contType, http2.ContentTypeJSON) {

		var codeResp response2[GetResourceSpecResp]
		if err := serder.JSONToObjectStream(resp.Body, &codeResp); err != nil {
			return nil, 0, fmt.Errorf("parsing response: %w", err)
		}

		if codeResp.Code == ResponseCodeOK {
			return codeResp.Data.List, codeResp.Data.Total, nil
		}

		return nil, 0, codeResp.ToError()
	}

	return nil, 0, fmt.Errorf("unknow response content type: %s", contType)
}