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 {
	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
	}

	resp, err := http2.PostJSON(targetUrl, http2.RequestParam{
		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 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 respons2[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 {
	TaskId string `json:"taskId"`
}

func (c *Client) StopInferenceJob(req StopInferenceJobReq, token string) error {
	targetUrl, err := url.JoinPath(c.baseURL, "inference/createTask")
	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 respons2[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 respons2[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 respons2[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)
}