From 196c7c0d458ae74365cfe3602e73fa182db4c5fe Mon Sep 17 00:00:00 2001 From: wkd-woo Date: Wed, 1 Jul 2026 12:01:00 +0900 Subject: [PATCH] Add configurable allocation policy (packed/distributed) for replicated and MIG resources The existing allocation spreads replicated devices evenly across physical GPUs (distributed), which was designed for time-slicing where workloads compete for shared compute. MIG instances, however, are hardware-isolated and do not suffer from contention when packed onto the same GPU. This adds a --shared-devices-allocation-policy flag (env: SHARED_DEVICES_ALLOCATION_POLICY) with two options: "distributed" (default, preserving current behavior) and "packed" (bin-packing onto the fewest physical GPUs). The packed policy frees up entire GPUs for full-GPU workloads in mixed clusters. The flag applies uniformly to all non-aligned allocation paths (MIG, time-slicing, MPS) and can be configured per-node via ConfigMap and the nvidia.com/device-plugin.config node label. Full-GPU nodes are unaffected: alignedAlloc is selected before the policy is ever consulted. Both policies share a single greedyAlloc helper that differs only in a replicaComparator, selected from an allocationComparators map keyed by policy (comparatorForPolicy falls back to distributed for unknown values). This keeps the strategies expressed as a flipped comparator and lets future policies be added as new map entries without duplicating the selection loop. When the comparator ranks two physical GPUs equally, a per-allocation pickedFrom tie-break rotates to the least-touched sibling, preserving distribution across physical GPUs. Relates to #491 Signed-off-by: wkd-woo --- api/config/v1/consts.go | 6 + api/config/v1/flags.go | 15 +- cmd/nvidia-device-plugin/main.go | 15 + internal/rm/allocate.go | 107 +++++-- internal/rm/allocate_test.go | 475 ++++++++++++++++++++++++++++++- internal/rm/nvml_manager.go | 8 +- internal/rm/rm.go | 2 +- 7 files changed, 585 insertions(+), 43 deletions(-) diff --git a/api/config/v1/consts.go b/api/config/v1/consts.go index 43b5267077..925b9507a0 100644 --- a/api/config/v1/consts.go +++ b/api/config/v1/consts.go @@ -48,6 +48,12 @@ const ( DeviceIDStrategyIndex = "index" ) +// Constants to represent the various allocation policies +const ( + AllocationPolicyDistributed = "distributed" + AllocationPolicyPacked = "packed" +) + // Constants related to generating CDI specifications const ( DefaultCDIAnnotationPrefix = cdiapi.AnnotationPrefix diff --git a/api/config/v1/flags.go b/api/config/v1/flags.go index 4695960e51..1d5499ca32 100644 --- a/api/config/v1/flags.go +++ b/api/config/v1/flags.go @@ -73,12 +73,13 @@ type CommandLineFlags struct { // PluginCommandLineFlags holds the list of command line flags specific to the device plugin. type PluginCommandLineFlags struct { - PassDeviceSpecs *bool `json:"passDeviceSpecs" yaml:"passDeviceSpecs"` - DeviceListStrategy *deviceListStrategyFlag `json:"deviceListStrategy" yaml:"deviceListStrategy"` - DeviceIDStrategy *string `json:"deviceIDStrategy" yaml:"deviceIDStrategy"` - CDIAnnotationPrefix *string `json:"cdiAnnotationPrefix" yaml:"cdiAnnotationPrefix"` - NvidiaCTKPath *string `json:"nvidiaCTKPath" yaml:"nvidiaCTKPath"` - ContainerDriverRoot *string `json:"containerDriverRoot" yaml:"containerDriverRoot"` + PassDeviceSpecs *bool `json:"passDeviceSpecs" yaml:"passDeviceSpecs"` + DeviceListStrategy *deviceListStrategyFlag `json:"deviceListStrategy" yaml:"deviceListStrategy"` + DeviceIDStrategy *string `json:"deviceIDStrategy" yaml:"deviceIDStrategy"` + CDIAnnotationPrefix *string `json:"cdiAnnotationPrefix" yaml:"cdiAnnotationPrefix"` + NvidiaCTKPath *string `json:"nvidiaCTKPath" yaml:"nvidiaCTKPath"` + ContainerDriverRoot *string `json:"containerDriverRoot" yaml:"containerDriverRoot"` + SharedDevicesAllocationPolicy *string `json:"sharedDevicesAllocationPolicy" yaml:"sharedDevicesAllocationPolicy"` } // deviceListStrategyFlag is a custom type for parsing the deviceListStrategy flag. @@ -157,6 +158,8 @@ func (f *Flags) UpdateFromCLIFlags(c *cli.Context, flags []cli.Flag) { updateFromCLIFlag(&f.Plugin.NvidiaCTKPath, c, n) case "container-driver-root": updateFromCLIFlag(&f.Plugin.ContainerDriverRoot, c, n) + case "shared-devices-allocation-policy": + updateFromCLIFlag(&f.Plugin.SharedDevicesAllocationPolicy, c, n) } // GFD specific flags if f.GFD == nil { diff --git a/cmd/nvidia-device-plugin/main.go b/cmd/nvidia-device-plugin/main.go index 23d6ff9c5d..ebacd3eddc 100644 --- a/cmd/nvidia-device-plugin/main.go +++ b/cmd/nvidia-device-plugin/main.go @@ -142,6 +142,12 @@ func main() { Usage: "the path on the host where MPS-specific mounts and files are created by the MPS control daemon manager", EnvVars: []string{"MPS_ROOT"}, }, + &cli.StringFlag{ + Name: "shared-devices-allocation-policy", + Value: spec.AllocationPolicyDistributed, + Usage: "the allocation policy for replicated and MIG resources:\n\t\t[distributed | packed]", + EnvVars: []string{"SHARED_DEVICES_ALLOCATION_POLICY"}, + }, &cli.StringFlag{ Name: "device-discovery-strategy", Value: "auto", @@ -209,6 +215,15 @@ func validateFlags(infolib nvinfo.Interface, config *spec.Config) error { return fmt.Errorf("invalid --device-id-strategy option: %v", *config.Flags.Plugin.DeviceIDStrategy) } + if config.Flags.Plugin.SharedDevicesAllocationPolicy != nil { + switch *config.Flags.Plugin.SharedDevicesAllocationPolicy { + case spec.AllocationPolicyDistributed: + case spec.AllocationPolicyPacked: + default: + return fmt.Errorf("invalid --shared-devices-allocation-policy option: %s", *config.Flags.Plugin.SharedDevicesAllocationPolicy) + } + } + if config.Sharing.SharingStrategy() == spec.SharingStrategyMPS { if *config.Flags.MigStrategy == spec.MigStrategyMixed { return fmt.Errorf("using --mig-strategy=mixed is not supported with MPS") diff --git a/internal/rm/allocate.go b/internal/rm/allocate.go index 8f994c8cfe..64a6863634 100644 --- a/internal/rm/allocate.go +++ b/internal/rm/allocate.go @@ -19,27 +19,67 @@ package rm import ( "fmt" "sort" + + spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) -// distributedAlloc returns a list of devices such that any replicated -// devices are distributed across all replicated GPUs equally. It takes into -// account already allocated replicas to ensure a proper balance across them. -func (r *resourceManager) distributedAlloc(available, required []string, size int) ([]string, error) { - // Get the set of candidate devices as the difference between available and required. +// replicaCount tracks the total and available replica counts for a physical GPU. +type replicaCount struct { + total int + available int +} + +// allocated returns the number of replicas currently allocated on the GPU. +func (rc *replicaCount) allocated() int { + return rc.total - rc.available +} + +// replicaComparator decides whether the physical GPU represented by i should +// be preferred over the one represented by j when greedily selecting the next +// device to allocate. +type replicaComparator func(i, j *replicaCount) bool + +// allocationComparators maps each allocation policy to the comparator that +// implements it. All policies share the same greedy selection loop +// (greedyAlloc) and differ only in how the next best candidate is chosen. +var allocationComparators = map[string]replicaComparator{ + // distributed prefers GPUs with the fewest allocated replicas to spread + // workload evenly across physical GPUs. + spec.AllocationPolicyDistributed: func(i, j *replicaCount) bool { + return i.allocated() < j.allocated() + }, + // packed prefers GPUs with the most allocated replicas to consolidate + // workloads onto fewer physical GPUs. + spec.AllocationPolicyPacked: func(i, j *replicaCount) bool { + return i.allocated() > j.allocated() + }, +} + +// comparatorForPolicy returns the comparator implementing the given +// allocation policy. Unknown policies are rejected at startup, but fall back +// to the default distributed policy here as a safety net. +func comparatorForPolicy(policy string) replicaComparator { + if comparator, ok := allocationComparators[policy]; ok { + return comparator + } + return allocationComparators[spec.AllocationPolicyDistributed] +} + +// prepareCandidates filters candidates from available devices (excluding required), +// validates there are enough, and builds a per-GPU replica count map. +func (r *resourceManager) prepareCandidates(available, required []string, size int) ([]string, map[string]*replicaCount, int, error) { candidates := r.devices.Subset(available).Difference(r.devices.Subset(required)).GetIDs() needed := size - len(required) if len(candidates) < needed { - return nil, fmt.Errorf("not enough available devices to satisfy allocation") + return nil, nil, 0, fmt.Errorf("not enough available devices to satisfy allocation") } - // For each candidate device, build a mapping of (stripped) device ID to - // total / available replicas for that device. - replicas := make(map[string]*struct{ total, available int }) + replicas := make(map[string]*replicaCount) for _, c := range candidates { id := AnnotatedID(c).GetID() if _, exists := replicas[id]; !exists { - replicas[id] = &struct{ total, available int }{} + replicas[id] = &replicaCount{} } replicas[id].available++ } @@ -51,30 +91,40 @@ func (r *resourceManager) distributedAlloc(available, required []string, size in replicas[id].total++ } + return candidates, replicas, needed, nil +} + +// greedyAlloc returns a list of devices by repeatedly selecting the best +// remaining candidate according to the supplied comparator. It takes into +// account already allocated replicas so that consecutive allocations keep +// following the policy the comparator implements. +func (r *resourceManager) greedyAlloc(available, required []string, size int, preferred replicaComparator) ([]string, error) { + candidates, replicas, needed, err := r.prepareCandidates(available, required, size) + if err != nil { + return nil, err + } + // Track how many slots have already been picked from each physical device - // during this allocation. Used as the tie-break sort key below so the - // allocator rotates to a sibling physical device when the underlying - // "used" counts would otherwise tie. + // during this allocation. Used as the tie-break sort key below so that, + // when the comparator ranks two physical GPUs equally, the allocator + // rotates to a sibling device it has touched the least this round. This + // keeps the distributed policy spreading replicas across physical GPUs + // even when their allocated counts tie. pickedFrom := make(map[string]int) - // Grab the set of 'needed' devices one-by-one from the candidates list. - // Before selecting each candidate, first sort the candidate list using the - // replicas map above. After sorting, the first element in the list will - // contain the device with the least difference between total and available - // replications (based on what's already been allocated). When two devices - // tie on that count, prefer the physical device we have not touched (or - // have touched the least) during this allocation. Add this device to the - // list of devices to allocate, remove it from the candidate list, down - // its available count in the replicas map, and repeat. + // Select devices one-by-one. The supplied comparator decides which + // physical GPU is preferred for the current policy. Comparators order + // solely by allocated() (see TestComparatorsOrderSolelyByAllocated), so + // equal allocated counts mean the comparator has no preference and the + // pickedFrom tie-break above applies. var devices []string for i := 0; i < needed; i++ { sort.Slice(candidates, func(i, j int) bool { iid := AnnotatedID(candidates[i]).GetID() jid := AnnotatedID(candidates[j]).GetID() - idiff := replicas[iid].total - replicas[iid].available - jdiff := replicas[jid].total - replicas[jid].available - if idiff != jdiff { - return idiff < jdiff + ri, rj := replicas[iid], replicas[jid] + if ri.allocated() != rj.allocated() { + return preferred(ri, rj) } return pickedFrom[iid] < pickedFrom[jid] }) @@ -85,8 +135,5 @@ func (r *resourceManager) distributedAlloc(available, required []string, size in candidates = candidates[1:] } - // Add the set of required devices to this list and return it. - devices = append(required, devices...) - - return devices, nil + return append(required, devices...), nil } diff --git a/internal/rm/allocate_test.go b/internal/rm/allocate_test.go index 4edc38e4a1..7344149f03 100644 --- a/internal/rm/allocate_test.go +++ b/internal/rm/allocate_test.go @@ -17,10 +17,13 @@ package rm import ( + "fmt" "testing" "github.com/stretchr/testify/require" pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) func makeReplicatedDevices(t *testing.T, gpuToReplicas map[string]int) Devices { @@ -39,10 +42,40 @@ func makeReplicatedDevices(t *testing.T, gpuToReplicas map[string]int) Devices { return ds } -func countPerGPU(annotatedIDs []string) map[string]int { +// newTestDevices creates a Devices map with replicated devices for testing. +// Each GPU gets 'replicas' number of annotated device entries. +func newTestDevices(gpuIDs []string, replicas int) Devices { + devices := make(Devices) + for _, id := range gpuIDs { + for r := 0; r < replicas; r++ { + annotatedID := string(NewAnnotatedID(id, r)) + devices[annotatedID] = &Device{ + Device: pluginapi.Device{ + ID: annotatedID, + Health: pluginapi.Healthy, + }, + Index: id, + } + } + } + return devices +} + +// getDeviceIDs returns all device IDs from a Devices map as a string slice. +func getDeviceIDs(devices Devices) []string { + var ids []string + for id := range devices { + ids = append(ids, id) + } + return ids +} + +// countPerGPU counts how many allocated device IDs belong to each physical GPU. +func countPerGPU(allocated []string) map[string]int { counts := make(map[string]int) - for _, id := range annotatedIDs { - counts[AnnotatedID(id).GetID()]++ + for _, id := range allocated { + gpuID := AnnotatedID(id).GetID() + counts[gpuID]++ } return counts } @@ -59,7 +92,7 @@ func TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs(t *te "GPU-1::1", } - allocated, err := r.distributedAlloc(available, nil, 2) + allocated, err := r.greedyAlloc(available, nil, 2, comparatorForPolicy(spec.AllocationPolicyDistributed)) require.NoError(t, err) require.Len(t, allocated, 2) @@ -71,3 +104,437 @@ func TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs(t *te "expected 1 slot from GPU-1 (the still-available second physical GPU) instead of stacking both on GPU-0; got: %v", counts) } + +func TestDistributedAlloc(t *testing.T) { + testCases := []struct { + description string + gpuIDs []string + replicas int + available []string // if nil, use all devices + required []string + size int + expectError bool + validate func(t *testing.T, allocated []string, allDevices Devices) + }{ + { + description: "2 GPUs, 4 replicas each, allocate 2 — should distribute across GPUs", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 2, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 2) + // distributed: should pick one from each GPU + require.Equal(t, 1, counts["gpu0"], "expected 1 allocation from gpu0") + require.Equal(t, 1, counts["gpu1"], "expected 1 allocation from gpu1") + }, + }, + { + description: "3 GPUs, 2 replicas each, allocate 3 — should distribute across all GPUs", + gpuIDs: []string{"gpu0", "gpu1", "gpu2"}, + replicas: 2, + required: []string{}, + size: 3, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 3) + require.Equal(t, 1, counts["gpu0"]) + require.Equal(t, 1, counts["gpu1"]) + require.Equal(t, 1, counts["gpu2"]) + }, + }, + { + description: "2 GPUs, 4 replicas each, allocate 4 — should distribute 2 per GPU", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 4, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 4) + require.Equal(t, 2, counts["gpu0"]) + require.Equal(t, 2, counts["gpu1"]) + }, + }, + { + description: "allocate 1 from single GPU — trivial case", + gpuIDs: []string{"gpu0"}, + replicas: 4, + required: []string{}, + size: 1, + validate: func(t *testing.T, allocated []string, _ Devices) { + require.Len(t, allocated, 1) + counts := countPerGPU(allocated) + require.Equal(t, 1, counts["gpu0"]) + }, + }, + { + description: "not enough devices — should return error", + gpuIDs: []string{"gpu0"}, + replicas: 2, + required: []string{}, + size: 5, + expectError: true, + }, + { + description: "partial availability simulates pre-allocated state — should still distribute", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 2, + // Only gpu1 replicas are available (simulates gpu0 already fully allocated) + available: nil, // will be overridden in test body + validate: func(t *testing.T, allocated []string, _ Devices) { + require.Len(t, allocated, 2) + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices := newTestDevices(tc.gpuIDs, tc.replicas) + available := tc.available + if available == nil { + available = getDeviceIDs(devices) + } + + rm := resourceManager{ + config: &spec.Config{}, + devices: devices, + } + + allocated, err := rm.greedyAlloc(available, tc.required, tc.size, comparatorForPolicy(spec.AllocationPolicyDistributed)) + if tc.expectError { + require.Error(t, err) + return + } + require.NoError(t, err) + if tc.validate != nil { + tc.validate(t, allocated, devices) + } + }) + } +} + +func TestPackedAlloc(t *testing.T) { + testCases := []struct { + description string + gpuIDs []string + replicas int + available []string // if nil, use all devices + required []string + size int + expectError bool + validate func(t *testing.T, allocated []string, allDevices Devices) + }{ + { + description: "2 GPUs, 4 replicas each, allocate 2 — should pack onto same GPU", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 2, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 2) + // packed: should pick 2 from the same GPU + require.Len(t, counts, 1, "expected all allocations from a single GPU") + }, + }, + { + description: "3 GPUs, 2 replicas each, allocate 3 — should fill one GPU first", + gpuIDs: []string{"gpu0", "gpu1", "gpu2"}, + replicas: 2, + required: []string{}, + size: 3, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 3) + // One GPU should have 2, another should have 1, third should have 0 + maxCount := 0 + for _, c := range counts { + if c > maxCount { + maxCount = c + } + } + require.Equal(t, 2, maxCount, "expected one GPU to be fully packed with 2 allocations") + }, + }, + { + description: "2 GPUs, 4 replicas each, allocate 4 — should pack onto single GPU", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 4, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 4) + // packed: should fill one GPU entirely + require.Len(t, counts, 1, "expected all 4 allocations from a single GPU") + }, + }, + { + description: "2 GPUs, 4 replicas each, allocate 5 — should fill one GPU then overflow", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 5, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 5) + // One GPU should have 4, the other should have 1 + maxCount := 0 + minCount := 999 + for _, c := range counts { + if c > maxCount { + maxCount = c + } + if c < minCount { + minCount = c + } + } + require.Equal(t, 4, maxCount, "expected one GPU fully packed") + require.Equal(t, 1, minCount, "expected overflow to second GPU") + }, + }, + { + description: "allocate 1 from single GPU — trivial case", + gpuIDs: []string{"gpu0"}, + replicas: 4, + required: []string{}, + size: 1, + validate: func(t *testing.T, allocated []string, _ Devices) { + require.Len(t, allocated, 1) + }, + }, + { + description: "not enough devices — should return error", + gpuIDs: []string{"gpu0"}, + replicas: 2, + required: []string{}, + size: 5, + expectError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices := newTestDevices(tc.gpuIDs, tc.replicas) + available := tc.available + if available == nil { + available = getDeviceIDs(devices) + } + + rm := resourceManager{ + config: &spec.Config{}, + devices: devices, + } + + allocated, err := rm.greedyAlloc(available, tc.required, tc.size, comparatorForPolicy(spec.AllocationPolicyPacked)) + if tc.expectError { + require.Error(t, err) + return + } + require.NoError(t, err) + if tc.validate != nil { + tc.validate(t, allocated, devices) + } + }) + } +} + +// TestDistributedAllocIsDefault verifies that the default allocation policy +// (no SharedDevicesAllocationPolicy set, or explicitly "distributed") produces +// the same distributed behavior as the existing implementation. +func TestDistributedAllocIsDefault(t *testing.T) { + gpuIDs := []string{"gpu0", "gpu1", "gpu2"} + replicas := 4 + devices := newTestDevices(gpuIDs, replicas) + available := getDeviceIDs(devices) + + // Run distributed allocation multiple times to verify consistent behavior + for i := 0; i < 10; i++ { + rm := resourceManager{ + config: &spec.Config{}, + devices: devices, + } + + allocated, err := rm.greedyAlloc(available, []string{}, 3, comparatorForPolicy(spec.AllocationPolicyDistributed)) + require.NoError(t, err) + require.Len(t, allocated, 3) + + counts := countPerGPU(allocated) + // distributed: should always pick 1 from each of the 3 GPUs + require.Equal(t, 1, counts["gpu0"], "iteration %d: expected 1 from gpu0", i) + require.Equal(t, 1, counts["gpu1"], "iteration %d: expected 1 from gpu1", i) + require.Equal(t, 1, counts["gpu2"], "iteration %d: expected 1 from gpu2", i) + } +} + +// TestPackedVsDistributedContrast directly compares the two allocation +// strategies on the same input to verify they produce meaningfully different results. +func TestPackedVsDistributedContrast(t *testing.T) { + gpuIDs := []string{"gpu0", "gpu1"} + replicas := 4 + devices := newTestDevices(gpuIDs, replicas) + available := getDeviceIDs(devices) + + rm := resourceManager{ + config: &spec.Config{}, + devices: devices, + } + + // Distributed: allocate 4 → should be 2 per GPU + distAllocated, err := rm.greedyAlloc(available, []string{}, 4, comparatorForPolicy(spec.AllocationPolicyDistributed)) + require.NoError(t, err) + distCounts := countPerGPU(distAllocated) + require.Equal(t, 2, distCounts["gpu0"]) + require.Equal(t, 2, distCounts["gpu1"]) + + // Packed: allocate 4 → should be 4 on one GPU + packAllocated, err := rm.greedyAlloc(available, []string{}, 4, comparatorForPolicy(spec.AllocationPolicyPacked)) + require.NoError(t, err) + packCounts := countPerGPU(packAllocated) + require.Len(t, packCounts, 1, "packed should use only 1 GPU") + + // Verify they are actually different + require.NotEqual(t, distCounts, packCounts, "distributed and packed should produce different allocation patterns") +} + +// TestComparatorForPolicy verifies that each allocation policy resolves to +// the comparator implementing it, and that unknown or empty policies fall +// back to the default distributed comparator. +func TestComparatorForPolicy(t *testing.T) { + moreAllocated := &replicaCount{total: 4, available: 1} // 3 allocated + lessAllocated := &replicaCount{total: 4, available: 3} // 1 allocated + + testCases := []struct { + description string + policy string + // expectPrefersLessAllocated is true for distributed-style comparators, + // which prefer the GPU with fewer allocated replicas. + expectPrefersLessAllocated bool + }{ + { + description: "distributed prefers GPU with fewer allocated replicas", + policy: spec.AllocationPolicyDistributed, + expectPrefersLessAllocated: true, + }, + { + description: "packed prefers GPU with more allocated replicas", + policy: spec.AllocationPolicyPacked, + expectPrefersLessAllocated: false, + }, + { + description: "empty policy falls back to distributed", + policy: "", + expectPrefersLessAllocated: true, + }, + { + description: "unknown policy falls back to distributed", + policy: "no-such-policy", + expectPrefersLessAllocated: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + comparator := comparatorForPolicy(tc.policy) + require.NotNil(t, comparator) + require.Equal(t, tc.expectPrefersLessAllocated, comparator(lessAllocated, moreAllocated)) + require.Equal(t, !tc.expectPrefersLessAllocated, comparator(moreAllocated, lessAllocated)) + }) + } +} + +// TestComparatorsOrderSolelyByAllocated pins the invariant that every +// allocation comparator orders physical GPUs solely by their allocated() +// count. The tie-break in greedyAlloc depends on this: it treats equal +// allocated counts as "the comparator has no preference" and falls back to +// the pickedFrom rotation, so a comparator that distinguishes GPUs by +// anything else would be silently ignored there. +func TestComparatorsOrderSolelyByAllocated(t *testing.T) { + for policy, preferred := range allocationComparators { + t.Run(policy, func(t *testing.T) { + // Equal allocated counts with different total/available shapes + // must rank equal so the tie-break applies. + a := &replicaCount{total: 8, available: 6} // 2 allocated + b := &replicaCount{total: 4, available: 2} // 2 allocated + require.False(t, preferred(a, b), "GPUs with equal allocated counts must rank equal") + require.False(t, preferred(b, a), "GPUs with equal allocated counts must rank equal") + + // Different allocated counts must be strictly ordered. + c := &replicaCount{total: 8, available: 5} // 3 allocated + require.NotEqual(t, preferred(a, c), preferred(c, a), "GPUs with different allocated counts must be strictly ordered") + }) + } +} + +// newFullGPUDevices creates a Devices map with non-replicated full GPU devices. +// These devices have no annotations and support aligned allocation. +func newFullGPUDevices(uuids []string) Devices { + devices := make(Devices) + for i, uuid := range uuids { + devices[uuid] = &Device{ + Device: pluginapi.Device{ + ID: uuid, + Health: pluginapi.Healthy, + }, + Index: fmt.Sprintf("%d", i), + Paths: []string{fmt.Sprintf("/dev/nvidia%d", i)}, + } + } + return devices +} + +// TestFullGPUNodeIgnoresAllocationPolicy verifies that on a node with full +// (non-MIG, non-replicated) GPUs, the allocation policy setting has no effect. +// This is critical for mixed clusters where the same DaemonSet deploys +// nvidia-device-plugin with identical flags to both full GPU and MIG nodes. +func TestFullGPUNodeIgnoresAllocationPolicy(t *testing.T) { + uuids := []string{"GPU-aaa", "GPU-bbb", "GPU-ccc", "GPU-ddd"} + devices := newFullGPUDevices(uuids) + available := getDeviceIDs(devices) + + // Verify precondition: these devices support aligned allocation + require.True(t, devices.AlignedAllocationSupported(), "full GPU devices should support aligned allocation") + + // Verify precondition: no annotations on available IDs + require.False(t, AnnotatedIDs(available).AnyHasAnnotations(), "full GPU device IDs should not have annotations") + + // With packed policy set, getPreferredAllocation should still go to alignedAlloc + // (not packedAlloc). Since alignedAlloc requires NVML, we verify the branching + // logic at the condition level: the aligned path is taken before + // sharedDevicesAllocationPolicy is ever checked. + t.Run("AlignedAllocation is selected regardless of sharedDevicesAllocationPolicy", func(t *testing.T) { + // The condition that selects alignedAlloc: + // r.Devices().AlignedAllocationSupported() && !AnnotatedIDs(available).AnyHasAnnotations() + // must be true for full GPU devices, ensuring packedAlloc is never reached. + isAlignedPath := devices.AlignedAllocationSupported() && !AnnotatedIDs(available).AnyHasAnnotations() + require.True(t, isAlignedPath, "full GPU nodes must always take the aligned allocation path") + }) + + // Verify that MIG devices do NOT take the aligned path + t.Run("MIG devices do not take aligned allocation path", func(t *testing.T) { + migDevices := make(Devices) + migUUIDs := []string{"MIG-aaa"} + for _, uuid := range migUUIDs { + migDevices[uuid] = &Device{ + Device: pluginapi.Device{ + ID: uuid, + Health: pluginapi.Healthy, + }, + Index: "0:0", // MIG index contains ":" + Paths: []string{"/dev/nvidia0"}, + } + } + require.False(t, migDevices.AlignedAllocationSupported(), "MIG devices should not support aligned allocation") + }) + + // Verify that replicated (annotated) devices do NOT take the aligned path + t.Run("replicated devices do not take aligned allocation path", func(t *testing.T) { + replicatedDevices := newTestDevices([]string{"gpu0"}, 4) + replicatedAvailable := getDeviceIDs(replicatedDevices) + require.True(t, AnnotatedIDs(replicatedAvailable).AnyHasAnnotations(), "replicated device IDs should have annotations") + }) +} diff --git a/internal/rm/nvml_manager.go b/internal/rm/nvml_manager.go index f0ca775d1b..299806510a 100644 --- a/internal/rm/nvml_manager.go +++ b/internal/rm/nvml_manager.go @@ -116,8 +116,12 @@ func (r *nvmlResourceManager) getPreferredAllocation(available, required []strin return r.alignedAlloc(available, required, size) } - // Otherwise, distribute them evenly across all replicated GPUs - return r.distributedAlloc(available, required, size) + // Otherwise, apply the configured allocation policy for replicated/MIG resources. + policy := spec.AllocationPolicyDistributed + if r.config.Flags.Plugin != nil && r.config.Flags.Plugin.SharedDevicesAllocationPolicy != nil { + policy = *r.config.Flags.Plugin.SharedDevicesAllocationPolicy + } + return r.greedyAlloc(available, required, size, comparatorForPolicy(policy)) } // alignedAlloc shells out to the alignedAllocationPolicy that is set in diff --git a/internal/rm/rm.go b/internal/rm/rm.go index 8ef6b8874c..9c7414dab8 100644 --- a/internal/rm/rm.go +++ b/internal/rm/rm.go @@ -62,7 +62,7 @@ func (r *resourceManager) GetDevicePaths(ids []string) []string { // GetPreferredAllocation runs an allocation algorithm over the inputs. func (r *resourceManager) GetPreferredAllocation(available, required []string, size int) ([]string, error) { - return r.distributedAlloc(available, required, size) + return r.greedyAlloc(available, required, size, comparatorForPolicy(spec.AllocationPolicyDistributed)) } // Resource gets the resource name associated with the ResourceManager