volcano.sh/volcano@v1.9.0/pkg/scheduler/plugins/task-topology/util.go (about)

     1  /*
     2  Copyright 2021 The Volcano 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 tasktopology
    18  
    19  import (
    20  	"volcano.sh/apis/pkg/apis/batch/v1alpha1"
    21  	"volcano.sh/volcano/pkg/scheduler/api"
    22  	"volcano.sh/volcano/pkg/scheduler/framework"
    23  )
    24  
    25  const (
    26  	// PluginName indicates name of volcano scheduler plugin
    27  	PluginName = "task-topology"
    28  	// PluginWeight is task-topology plugin weight in nodeOrderFn
    29  	PluginWeight = "task-topology.weight"
    30  	// JobAffinityKey is the key to read in task-topology arguments from job annotations
    31  	JobAffinityKey = "volcano.sh/task-topology"
    32  	// OutOfBucket indicates task is outside of any bucket
    33  	OutOfBucket = -1
    34  
    35  	// JobAffinityAnnotations is the key to read in task-topology affinity arguments from podgroup annotations
    36  	JobAffinityAnnotations = "volcano.sh/task-topology-affinity"
    37  	// JobAntiAffinityAnnotations is the key to read in task-topology anti-affinity arguments from podgroup annotations
    38  	JobAntiAffinityAnnotations = "volcano.sh/task-topology-anti-affinity"
    39  	// TaskOrderAnnotations is the key to read in task-topology task order arguments from podgroup annotations
    40  	TaskOrderAnnotations = "volcano.sh/task-topology-task-order"
    41  )
    42  
    43  // TaskTopology is struct used to save affinity infos of a job read from job plugin or annotations
    44  type TaskTopology struct {
    45  	Affinity     [][]string `json:"affinity,omitempty"`
    46  	AntiAffinity [][]string `json:"antiAffinity,omitempty"`
    47  	TaskOrder    []string   `json:"taskOrder,omitempty"`
    48  }
    49  
    50  func calculateWeight(args framework.Arguments) int {
    51  	/*
    52  	   User Should give taskTopologyWeight in this format(task-topology.weight).
    53  
    54  	   actions: "enqueue, reclaim, allocate, backfill, preempt"
    55  	   tiers:
    56  	   - plugins:
    57  	     - name: task-topology
    58  	       arguments:
    59  	         task-topology.weight: 10
    60  	*/
    61  	// Values are initialized to 1.
    62  	weight := 1
    63  
    64  	args.GetInt(&weight, PluginWeight)
    65  
    66  	return weight
    67  }
    68  
    69  func getTaskName(task *api.TaskInfo) string {
    70  	return task.Pod.Annotations[v1alpha1.TaskSpecKey]
    71  }
    72  
    73  func addAffinity(m map[string]map[string]struct{}, src, dst string) {
    74  	srcMap, ok := m[src]
    75  	if !ok {
    76  		srcMap = make(map[string]struct{})
    77  		m[src] = srcMap
    78  	}
    79  	srcMap[dst] = struct{}{}
    80  }
    81  
    82  // TaskOrder is struct used to save task order
    83  type TaskOrder struct {
    84  	tasks   []*api.TaskInfo
    85  	manager *JobManager
    86  }
    87  
    88  func (p *TaskOrder) Len() int { return len(p.tasks) }
    89  
    90  func (p *TaskOrder) Swap(l, r int) {
    91  	p.tasks[l], p.tasks[r] = p.tasks[r], p.tasks[l]
    92  }
    93  
    94  func (p *TaskOrder) Less(l, r int) bool {
    95  	L := p.tasks[l]
    96  	R := p.tasks[r]
    97  
    98  	LHasNode := L.NodeName != ""
    99  	RHasNode := R.NodeName != ""
   100  	if LHasNode || RHasNode {
   101  		// the task bounded would have high priority
   102  		if LHasNode != RHasNode {
   103  			return !LHasNode
   104  		}
   105  		// all bound, any order is alright
   106  		return L.NodeName > R.NodeName
   107  	}
   108  
   109  	result := p.manager.taskAffinityOrder(L, R)
   110  	// they have the same taskAffinity order, any order is alright
   111  	if result == 0 {
   112  		return L.Name > R.Name
   113  	}
   114  	return result < 0
   115  }