pcm-coordinator/api/internal/scheduler/schedulers/aiScheduler.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 schedulers

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"gitlink.org.cn/JointCloud/pcm-ac/hpcAC"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler"
	"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/api/internal/scheduler/strategy"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/strategy/param"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/pkg/response"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/models"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/JointCloud/pcm-octopus/octopus"
	"sync"
)

type AiScheduler struct {
	yamlString string
	task       *response.TaskInfo
	*scheduler.Scheduler
	option *option.AiOption
	ctx    context.Context
}

type AiResult struct {
	TaskId    string
	ClusterId string
	Replica   int32
	Msg       string
}

func NewAiScheduler(ctx context.Context, val string, scheduler *scheduler.Scheduler, option *option.AiOption) (*AiScheduler, error) {
	return &AiScheduler{ctx: ctx, yamlString: val, Scheduler: scheduler, option: option}, nil
}

func (as *AiScheduler) GetNewStructForDb(task *response.TaskInfo, resource string, participantId int64) (interface{}, error) {
	ai := models.Ai{
		ParticipantId: participantId,
		TaskId:        task.TaskId,
		Status:        "Saved",
		YamlString:    as.yamlString,
	}
	utils.Convert(task.Metadata, &ai)
	return ai, nil
}

func (as *AiScheduler) PickOptimalStrategy() (strategy.Strategy, error) {
	if as.option.AiClusterId != "" {
		// TODO database operation Find
		return &strategy.SingleAssignment{Cluster: &strategy.AssignedCluster{ClusterId: "", Replicas: 1}}, nil
	}

	resources, err := as.findClustersWithResources()

	if err != nil {
		return nil, err
	}
	if len(resources) == 0 {
		return nil, errors.New("no cluster has resources")
	}

	if len(resources) == 1 {
		var cluster strategy.AssignedCluster
		cluster.ClusterId = resources[0].ClusterId
		cluster.Replicas = 1
		return &strategy.SingleAssignment{Cluster: &cluster}, nil
	}

	params := &param.Params{Resources: resources}

	switch as.option.StrategyName {
	case strategy.REPLICATION:
		var clusterIds []string
		for _, resource := range resources {
			clusterIds = append(clusterIds, resource.ClusterId)
		}
		strategy := strategy.NewReplicationStrategy(clusterIds, 1)
		return strategy, nil
	case strategy.RESOURCES_PRICING:
		strategy := strategy.NewPricingStrategy(&param.ResourcePricingParams{Params: params, Replicas: 1})
		return strategy, nil
	case strategy.DYNAMIC_RESOURCES:
		strategy := strategy.NewDynamicResourcesStrategy(params.Resources, as.option, 1)
		return strategy, nil
	case strategy.STATIC_WEIGHT:
		//todo resources should match cluster StaticWeightMap
		strategy := strategy.NewStaticWeightStrategy(as.option.ClusterToStaticWeight, 1)
		return strategy, nil
	}

	return nil, errors.New("no strategy has been chosen")
}

func (as *AiScheduler) AssignTask(clusters []*strategy.AssignedCluster) (interface{}, error) {
	if clusters == nil {
		return nil, errors.New("clusters is nil")
	}

	for i := len(clusters) - 1; i >= 0; i-- {
		if clusters[i].Replicas == 0 {
			clusters = append(clusters[:i], clusters[i+1:]...)
		}
	}

	if len(clusters) == 0 {
		return nil, errors.New("clusters is nil")
	}

	var wg sync.WaitGroup
	var results []*AiResult
	var errs []interface{}
	var ch = make(chan *AiResult, len(clusters))
	var errCh = make(chan interface{}, len(clusters))

	executorMap := *as.AiExecutor
	for _, cluster := range clusters {
		c := cluster
		wg.Add(1)
		go func() {
			opt, _ := cloneAiOption(as.option)
			resp, err := executorMap[c.ClusterId].Execute(as.ctx, opt)

			if err != nil {
				e := struct {
					err       error
					clusterId string
				}{
					err:       err,
					clusterId: c.ClusterId,
				}
				errCh <- e
				wg.Done()
				return
			}

			result, _ := convertType(resp)
			result.Replica = c.Replicas
			result.ClusterId = c.ClusterId

			ch <- result
			wg.Done()
		}()
	}
	wg.Wait()
	close(ch)
	close(errCh)

	for e := range errCh {
		errs = append(errs, e)
	}

	if len(errs) != len(clusters) {
		return nil, errors.New("submit task failed")
	}

	if len(errs) != 0 {
		var msg string
		for _, err := range errs {
			e := (err).(struct {
				err       error
				clusterId string
			})
			msg += fmt.Sprintf("clusterId: %v , error: %v \n", e.clusterId, e.err.Error())
		}
		return nil, errors.New(msg)
	}

	for s := range ch {
		// TODO: database operation
		results = append(results, s)
	}

	return results, nil
}

func (as *AiScheduler) findClustersWithResources() ([]*collector.ResourceStats, error) {
	var wg sync.WaitGroup
	var ch = make(chan *collector.ResourceStats, len(*as.ResourceCollector))
	var errCh = make(chan error, len(*as.ResourceCollector))

	var resourceSpecs []*collector.ResourceStats
	var errs []error

	for _, resourceCollector := range *as.ResourceCollector {
		wg.Add(1)
		rc := resourceCollector
		go func() {
			spec, err := rc.GetResourceStats(as.ctx)
			if err != nil {
				errCh <- err
				wg.Done()
				return
			}
			ch <- spec
			wg.Done()
		}()
	}
	wg.Wait()
	close(ch)
	close(errCh)

	for s := range ch {
		resourceSpecs = append(resourceSpecs, s)
	}

	for e := range errCh {
		errs = append(errs, e)
	}

	if len(errs) != 0 {
		return nil, errors.New("get resources failed")
	}

	if len(resourceSpecs) == 0 {
		return nil, errors.New("no resource found")
	}
	return resourceSpecs, nil
}

func convertType(in interface{}) (*AiResult, error) {
	var result AiResult
	switch (in).(type) {
	case *hpcAC.SubmitTaskAiResp:
		resp := (in).(*hpcAC.SubmitTaskAiResp)
		if resp.Code == "0" {
			result.TaskId = resp.Data
		} else {
			result.Msg = resp.Msg
		}
		return &result, nil
	case *octopus.CreateTrainJobResp:
		resp := (in).(*octopus.CreateTrainJobResp)

		if resp.Success {
			result.TaskId = resp.Payload.JobId
		} else {
			result.Msg = resp.Error.Message
		}

		return &result, nil
	default:
		return nil, errors.New("ai task response failed")
	}
}

func cloneAiOption(opt *option.AiOption) (*option.AiOption, error) {
	origJSON, err := json.Marshal(opt)
	if err != nil {
		return nil, err
	}

	clone := option.AiOption{}
	if err = json.Unmarshal(origJSON, &clone); err != nil {
		return nil, err
	}

	return &clone, nil
}