pcm-coordinator/internal/scheduler/strategy/dataLocality.go

package strategy

import (
	"errors"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
)

type DataLocality struct {
	replicas       int32
	dataDistribute *types.DataDistribute
}

func NewDataLocality(replicas int32, dataDistribute *types.DataDistribute) *DataLocality {
	return &DataLocality{dataDistribute: dataDistribute, replicas: replicas}
}

func (d *DataLocality) Schedule() ([]*AssignedCluster, error) {
	if d.replicas < 1 {
		return nil, errors.New("replicas must be greater than 0")
	}

	var results []*AssignedCluster

	clusterDataMap := make(map[string]int32, 0)
	for _, distribute := range d.dataDistribute.Model {
		for _, cluster := range distribute.Clusters {
			clusterDataMap[cluster.ClusterID]++
		}
	}

	for _, distribute := range d.dataDistribute.Image {
		for _, cluster := range distribute.Clusters {
			clusterDataMap[cluster.ClusterID]++
		}
	}

	for _, distribute := range d.dataDistribute.Code {
		for _, cluster := range distribute.Clusters {
			clusterDataMap[cluster.ClusterID]++
		}
	}

	for _, distribute := range d.dataDistribute.Dataset {
		for _, cluster := range distribute.Clusters {
			clusterDataMap[cluster.ClusterID]++
		}
	}

	if d.replicas == 1 {
		cluster := &AssignedCluster{Replicas: d.replicas}
		var largest int32
		for k, v := range clusterDataMap {
			if v > largest {
				largest = v
				cluster.ClusterId = k
			}
		}
		results = append(results, cluster)
	}
	return results, nil
}