feat(scheduler): account native sidecar container GPU resources

pkg/device/initContainer.go

@@ -20,88 +20,88 @@ import (
 	"sort"
 
 	corev1 "k8s.io/api/core/v1"
+
+	"github.com/Project-HAMi/HAMi/pkg/util"
 )
 
-type deviceKey struct {
-	devType string
-	uuid    string
-}
 type usage struct {
 	mem   int32
 	cores int32
 	slots int32
 }
 
+func isSidecarAt(pod *corev1.Pod, cidx int) bool {
+	if cidx < 0 || cidx >= len(pod.Spec.InitContainers) {
+		return false
+	}
+	return util.IsSidecarContainer(&pod.Spec.InitContainers[cidx])
+}
+
 // CollapseInitContainerUsage returns the effective device usage for a pod.
 func CollapseInitContainerUsage(pod *corev1.Pod, raw PodDevices) PodDevices {
 	if raw == nil {
 		return nil
 	}
 	numInit := len(pod.Spec.InitContainers)
-	initPeak := make(map[deviceKey]usage)
-	appSum := make(map[deviceKey]usage)
 
-	for devType, podSingle := range raw {
-		for cidx, ctrDevs := range podSingle {
-			for _, dev := range ctrDevs {
-				key := deviceKey{devType: devType, uuid: dev.UUID}
-				if cidx < numInit {
-					cur := initPeak[key]
-					if dev.Usedmem > cur.mem {
-						cur.mem = dev.Usedmem
-					}
-					if dev.Usedcores > cur.cores {
-						cur.cores = dev.Usedcores
-					}
-					// Init containers run sequentially: peak concurrency is one slot per device.
-					cur.slots = 1
-					initPeak[key] = cur
-				} else {
-					cur := appSum[key]
-					cur.mem += dev.Usedmem
-					cur.cores += dev.Usedcores
-					// App containers run concurrently: each occurrence is an additional slot.
-					// TODO: If PR #2584 changes sidecars to app-sum accounting, update this
-					// slot calculation to follow the same classification.
-					cur.slots++
-					appSum[key] = cur
-				}
-			}
-		}
+	type devState struct {
+		sc   usage // running sum of sidecars declared so far
+		peak usage // peak concurrent usage observed during the init phase
+		app  usage // sum over app containers
 	}
 
 	collapsed := make(PodDevices)
-	for devType := range raw {
-		uuidSet := make(map[string]struct{})
-		for k := range initPeak {
-			if k.devType == devType {
-				uuidSet[k.uuid] = struct{}{}
+	for devType, podSingle := range raw {
+		states := make(map[string]*devState)
+		get := func(uuid string) *devState {
+			s, ok := states[uuid]
+			if !ok {
+				s = &devState{}
+				states[uuid] = s
 			}
+			return s
 		}
-		for k := range appSum {
-			if k.devType == devType {
-				uuidSet[k.uuid] = struct{}{}
+
+		for cidx, ctrDevs := range podSingle {
+			switch {
+			case cidx < numInit && isSidecarAt(pod, cidx):
+				for _, dev := range ctrDevs {
+					s := get(dev.UUID)
+					// A sidecar starts and never exits: it permanently
+					// joins the set of running containers.
+					s.sc.mem += dev.Usedmem
+					s.sc.cores += dev.Usedcores
+					s.sc.slots++
+					s.peak.mem = max(s.peak.mem, s.sc.mem)
+					s.peak.cores = max(s.peak.cores, s.sc.cores)
+					s.peak.slots = max(s.peak.slots, s.sc.slots)
+				}
+			case cidx < numInit:
+				for _, dev := range ctrDevs {
+					s := get(dev.UUID)
+					s.peak.mem = max(s.peak.mem, s.sc.mem+dev.Usedmem)
+					s.peak.cores = max(s.peak.cores, s.sc.cores+dev.Usedcores)
+					s.peak.slots = max(s.peak.slots, s.sc.slots+1)
+				}
+			default:
+				for _, dev := range ctrDevs {
+					s := get(dev.UUID)
+					s.app.mem += dev.Usedmem
+					s.app.cores += dev.Usedcores
+					s.app.slots++
+				}
 			}
 		}
 
 		collapsedSingle := make(PodSingleDevice, 1)
 		var containerDevs ContainerDevices
-
-		for uuid := range uuidSet {
-			initU := initPeak[deviceKey{devType, uuid}]
-			appU := appSum[deviceKey{devType, uuid}]
-
-			effMem := max(initU.mem, appU.mem)
-			effCores := max(initU.cores, appU.cores)
-			// Raw entries carry no slot count; clamp to at least one.
-			effSlots := max(initU.slots, appU.slots, 1)
-
+		for uuid, s := range states {
 			containerDevs = append(containerDevs, ContainerDevice{
 				UUID:      uuid,
 				Type:      devType,
-				Usedmem:   effMem,
-				Usedcores: effCores,
-				Slots:     effSlots,
+				Usedmem:   max(s.peak.mem, s.sc.mem+s.app.mem),
+				Usedcores: max(s.peak.cores, s.sc.cores+s.app.cores),
+				Slots:     max(max(s.peak.slots, s.sc.slots+s.app.slots), 1),
 			})
 		}
 		sort.Slice(containerDevs, func(i, j int) bool {
@@ -113,9 +113,7 @@ func CollapseInitContainerUsage(pod *corev1.Pod, raw PodDevices) PodDevices {
 	return collapsed
 }
 
-// AppContainersOnlyDeviceUsage returns the device usage considering only app containers.
-// Used when init containers have finished and we want to shrink to app-only footprint.
-func AppContainersOnlyDeviceUsage(pod *corev1.Pod, raw PodDevices) PodDevices {
+func SteadyStateDeviceUsage(pod *corev1.Pod, raw PodDevices) PodDevices {
 	if raw == nil {
 		return nil
 	}
@@ -125,8 +123,8 @@ func AppContainersOnlyDeviceUsage(pod *corev1.Pod, raw PodDevices) PodDevices {
 	for devType, podSingle := range raw {
 		sums := make(map[string]usage)
 		for cidx, ctrDevs := range podSingle {
-			if cidx < numInit {
-				continue // skip init containers
+			if cidx < numInit && !isSidecarAt(pod, cidx) {
+				continue
 			}
 			for _, dev := range ctrDevs {
 				s := sums[dev.UUID]