k8s.io/kubernetes@v1.29.3/test/e2e/framework/pod/node_selection.go (about)

     1  /*
     2  Copyright 2019 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package pod
    18  
    19  import (
    20  	v1 "k8s.io/api/core/v1"
    21  )
    22  
    23  // NodeSelection specifies where to run a pod, using a combination of fixed node name,
    24  // node selector and/or affinity.
    25  type NodeSelection struct {
    26  	Name     string
    27  	Selector map[string]string
    28  	Affinity *v1.Affinity
    29  }
    30  
    31  // setNodeAffinityRequirement sets affinity with specified operator to nodeName to nodeSelection
    32  func setNodeAffinityRequirement(nodeSelection *NodeSelection, operator v1.NodeSelectorOperator, nodeName string) {
    33  	// Add node-anti-affinity.
    34  	if nodeSelection.Affinity == nil {
    35  		nodeSelection.Affinity = &v1.Affinity{}
    36  	}
    37  	if nodeSelection.Affinity.NodeAffinity == nil {
    38  		nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
    39  	}
    40  	if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
    41  		nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
    42  	}
    43  	nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
    44  		v1.NodeSelectorTerm{
    45  			MatchFields: []v1.NodeSelectorRequirement{
    46  				{Key: "metadata.name", Operator: operator, Values: []string{nodeName}},
    47  			},
    48  		})
    49  }
    50  
    51  // SetNodeAffinityTopologyRequirement sets node affinity to a specified topology
    52  func SetNodeAffinityTopologyRequirement(nodeSelection *NodeSelection, topology map[string]string) {
    53  	if nodeSelection.Affinity == nil {
    54  		nodeSelection.Affinity = &v1.Affinity{}
    55  	}
    56  	if nodeSelection.Affinity.NodeAffinity == nil {
    57  		nodeSelection.Affinity.NodeAffinity = &v1.NodeAffinity{}
    58  	}
    59  	if nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
    60  		nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &v1.NodeSelector{}
    61  	}
    62  	for k, v := range topology {
    63  		nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms = append(nodeSelection.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms,
    64  			v1.NodeSelectorTerm{
    65  				MatchExpressions: []v1.NodeSelectorRequirement{
    66  					{Key: k, Operator: v1.NodeSelectorOpIn, Values: []string{v}},
    67  				},
    68  			})
    69  
    70  	}
    71  }
    72  
    73  // SetAffinity sets affinity to nodeName to nodeSelection
    74  func SetAffinity(nodeSelection *NodeSelection, nodeName string) {
    75  	setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpIn, nodeName)
    76  }
    77  
    78  // SetAntiAffinity sets anti-affinity to nodeName to nodeSelection
    79  func SetAntiAffinity(nodeSelection *NodeSelection, nodeName string) {
    80  	setNodeAffinityRequirement(nodeSelection, v1.NodeSelectorOpNotIn, nodeName)
    81  }
    82  
    83  // SetNodeAffinity modifies the given pod object with
    84  // NodeAffinity to the given node name.
    85  func SetNodeAffinity(podSpec *v1.PodSpec, nodeName string) {
    86  	nodeSelection := &NodeSelection{}
    87  	SetAffinity(nodeSelection, nodeName)
    88  	podSpec.Affinity = nodeSelection.Affinity
    89  }
    90  
    91  // SetNodeSelection modifies the given pod object with
    92  // the specified NodeSelection
    93  func SetNodeSelection(podSpec *v1.PodSpec, nodeSelection NodeSelection) {
    94  	podSpec.NodeSelector = nodeSelection.Selector
    95  	podSpec.Affinity = nodeSelection.Affinity
    96  	// pod.Spec.NodeName should not be set directly because
    97  	// it will bypass the scheduler, potentially causing
    98  	// kubelet to Fail the pod immediately if it's out of
    99  	// resources. Instead, we want the pod to remain
   100  	// pending in the scheduler until the node has resources
   101  	// freed up.
   102  	if nodeSelection.Name != "" {
   103  		SetNodeAffinity(podSpec, nodeSelection.Name)
   104  	}
   105  }