volcano.sh/volcano@v1.9.0/pkg/scheduler/plugins/task-topology/bucket.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  	"k8s.io/apimachinery/pkg/types"
    21  	"k8s.io/klog/v2"
    22  
    23  	"volcano.sh/volcano/pkg/scheduler/api"
    24  )
    25  
    26  type reqAction int
    27  
    28  const (
    29  	reqSub reqAction = iota
    30  	reqAdd
    31  )
    32  
    33  // Bucket is struct used to classify tasks by affinity and anti-affinity
    34  type Bucket struct {
    35  	index       int
    36  	tasks       map[types.UID]*api.TaskInfo
    37  	taskNameSet map[string]int
    38  
    39  	// reqScore is score of resource
    40  	// now, we regard 1 CPU and 1 GPU and 1Gi memory as the same score.
    41  	reqScore float64
    42  	request  *api.Resource
    43  
    44  	boundTask int
    45  	node      map[string]int
    46  }
    47  
    48  // NewBucket create a new empty bucket
    49  func NewBucket() *Bucket {
    50  	return &Bucket{
    51  		index:       0,
    52  		tasks:       make(map[types.UID]*api.TaskInfo),
    53  		taskNameSet: make(map[string]int),
    54  
    55  		reqScore: 0,
    56  		request:  api.EmptyResource(),
    57  
    58  		boundTask: 0,
    59  		node:      make(map[string]int),
    60  	}
    61  }
    62  
    63  // CalcResReq calculates task resources request
    64  func (b *Bucket) CalcResReq(req *api.Resource, action reqAction) {
    65  	if req == nil {
    66  		return
    67  	}
    68  
    69  	cpu := req.MilliCPU
    70  	// treat 1Mi the same as 1m cpu 1m gpu
    71  	mem := req.Memory / 1024 / 1024
    72  	score := cpu + mem
    73  	for _, request := range req.ScalarResources {
    74  		score += request
    75  	}
    76  
    77  	switch action {
    78  	case reqSub:
    79  		b.reqScore -= score
    80  		b.request.Sub(req)
    81  	case reqAdd:
    82  		b.reqScore += score
    83  		b.request.Add(req)
    84  	default:
    85  		klog.V(3).Infof("Invalid action <%v> for resource <%v>", action, req)
    86  	}
    87  }
    88  
    89  // AddTask adds task into bucket
    90  func (b *Bucket) AddTask(taskName string, task *api.TaskInfo) {
    91  	b.taskNameSet[taskName]++
    92  	if task.NodeName != "" {
    93  		b.node[task.NodeName]++
    94  		b.boundTask++
    95  		return
    96  	}
    97  
    98  	b.tasks[task.Pod.UID] = task
    99  	b.CalcResReq(task.Resreq, reqAdd)
   100  }
   101  
   102  // TaskBound binds task to bucket
   103  func (b *Bucket) TaskBound(task *api.TaskInfo) {
   104  	b.node[task.NodeName]++
   105  	b.boundTask++
   106  
   107  	delete(b.tasks, task.Pod.UID)
   108  	b.CalcResReq(task.Resreq, reqSub)
   109  }