github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/requester/ranking/labels.go (about)

     1  package ranking
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/filecoin-project/bacalhau/pkg/model"
     7  	"github.com/filecoin-project/bacalhau/pkg/requester"
     8  	"github.com/rs/zerolog/log"
     9  	"k8s.io/apimachinery/pkg/labels"
    10  )
    11  
    12  type LabelsNodeRanker struct {
    13  }
    14  
    15  func NewLabelsNodeRanker() *LabelsNodeRanker {
    16  	return &LabelsNodeRanker{}
    17  }
    18  
    19  // RankNodes ranks nodes based on the node labels and job selectors:
    20  // - Rank 1: Selectors match node labels.
    21  // - Rank -1: Selectors don't match node labels.
    22  // - Rank 0: Job selectors are not set.
    23  func (s *LabelsNodeRanker) RankNodes(ctx context.Context, job model.Job, nodes []model.NodeInfo) ([]requester.NodeRank, error) {
    24  	ranks := make([]requester.NodeRank, len(nodes))
    25  	var selector labels.Selector
    26  	if len(job.Spec.NodeSelectors) > 0 {
    27  		requirements, err := model.FromLabelSelectorRequirements(job.Spec.NodeSelectors...)
    28  		if err != nil {
    29  			return nil, err
    30  		}
    31  		selector = labels.NewSelector().Add(requirements...)
    32  	}
    33  	for i, node := range nodes {
    34  		rank := 0
    35  		if selector != nil {
    36  			if selector.Matches(labels.Set(node.Labels)) {
    37  				rank = 1
    38  			} else {
    39  				log.Ctx(ctx).Trace().Msgf("filtering node %s with labels %s doesn't match selectors %+v",
    40  					node.PeerInfo.ID, node.Labels, job.Spec.NodeSelectors)
    41  				rank = -1
    42  			}
    43  		}
    44  		ranks[i] = requester.NodeRank{
    45  			NodeInfo: node,
    46  			Rank:     rank,
    47  		}
    48  	}
    49  	return ranks, nil
    50  }