diff --git a/pkg/cache/tas_flavor_snapshot.go b/pkg/cache/tas_flavor_snapshot.go index b02d00897f..38fd32df38 100644 --- a/pkg/cache/tas_flavor_snapshot.go +++ b/pkg/cache/tas_flavor_snapshot.go @@ -241,17 +241,14 @@ func (s *TASFlavorSnapshot) removeTASUsage(domainID utiltas.TopologyDomainID, us s.leaves[domainID].tasUsage.Sub(usage) } -func (s *TASFlavorSnapshot) FreeCapacityPerDomain() map[string]int { - freeCapacityPerDomain := make(map[string]int) +func (s *TASFlavorSnapshot) FreeCapacityPerDomain() map[string]resources.Requests { + freeCapacityPerDomain := make(map[string]resources.Requests) for domainID, leaf := range s.leaves { - // Calculate the total free capacity for the domain - totalFreeCapacity := 0 - for _, quantity := range leaf.freeCapacity { - totalFreeCapacity += int(quantity) - } - - freeCapacityPerDomain[string(domainID)] = totalFreeCapacity + // CPU is represented in millicores + // Memory is represented in bytes + // For other resources, use their raw integer value + freeCapacityPerDomain[string(domainID)] = leaf.freeCapacity.Clone() } return freeCapacityPerDomain diff --git a/pkg/cache/tas_flavor_snapshot_test.go b/pkg/cache/tas_flavor_snapshot_test.go new file mode 100644 index 0000000000..0d88455b79 --- /dev/null +++ b/pkg/cache/tas_flavor_snapshot_test.go @@ -0,0 +1,63 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package cache + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + corev1 "k8s.io/api/core/v1" + + "sigs.k8s.io/kueue/pkg/resources" +) + +func TestFreeCapacityPerDomain(t *testing.T) { + snapshot := &TASFlavorSnapshot{ + leaves: leafDomainByID{ + "domain1": &leafDomain{ + freeCapacity: resources.Requests{ + corev1.ResourceCPU: 2000, + corev1.ResourceMemory: 4 * 1024 * 1024 * 1024, // 4 GiB + "nvidia.com/gpu": 1, + }, + }, + "domain2": &leafDomain{ + freeCapacity: resources.Requests{ + corev1.ResourceCPU: 1000, + corev1.ResourceMemory: 2 * 1024 * 1024 * 1024, // 2 GiB + }, + }, + }, + } + + expected := map[string]resources.Requests{ + "domain1": { + corev1.ResourceCPU: 2000, + corev1.ResourceMemory: 4 * 1024 * 1024 * 1024, + "nvidia.com/gpu": 1, + }, + "domain2": { + corev1.ResourceCPU: 1000, + corev1.ResourceMemory: 2 * 1024 * 1024 * 1024, + }, + } + + actual := snapshot.FreeCapacityPerDomain() + if diff := cmp.Diff(expected, actual); diff != "" { + t.Errorf("FreeCapacityPerDomain() mismatch (-expected +actual):\n%s", diff) + } +}