Add configurable allocation policy for replicated and MIG resources

The existing distributedAlloc spreads replicated devices evenly across
physical GPUs, 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 --allocation-policy flag (env: ALLOCATION_POLICY) with two
options: "distributed" (default, preserving current behavior) and
"packed" (bin-packing onto fewest physical GPUs). The packed policy
helps free 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.

Author: wkd-woo

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

api/config/v1/flags.go

@@ -79,6 +79,7 @@ type PluginCommandLineFlags struct {
 	CDIAnnotationPrefix *string                 `json:"cdiAnnotationPrefix" yaml:"cdiAnnotationPrefix"`
 	NvidiaCTKPath       *string                 `json:"nvidiaCTKPath"       yaml:"nvidiaCTKPath"`
 	ContainerDriverRoot *string                 `json:"containerDriverRoot" yaml:"containerDriverRoot"`
+	AllocationPolicy    *string                 `json:"allocationPolicy"   yaml:"allocationPolicy"`
 }
 
 // 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 "allocation-policy":
+				updateFromCLIFlag(&f.Plugin.AllocationPolicy, c, n)
 			}
 			// GFD specific flags
 			if f.GFD == nil {

cmd/nvidia-device-plugin/main.go

@@ -144,6 +144,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:    "allocation-policy",
+			Value:   spec.AllocationPolicyDistributed,
+			Usage:   "the allocation policy for replicated and MIG resources:\n\t\t[distributed | packed]",
+			EnvVars: []string{"ALLOCATION_POLICY"},
+		},
 		&cli.StringFlag{
 			Name:    "device-discovery-strategy",
 			Value:   "auto",
@@ -205,6 +211,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.AllocationPolicy != nil {
+		switch *config.Flags.Plugin.AllocationPolicy {
+		case spec.AllocationPolicyDistributed:
+		case spec.AllocationPolicyPacked:
+		default:
+			return fmt.Errorf("invalid --allocation-policy option: %v", *config.Flags.Plugin.AllocationPolicy)
+		}
+	}
+
 	if config.Sharing.SharingStrategy() == spec.SharingStrategyMPS {
 		if *config.Flags.MigStrategy == spec.MigStrategyMixed {
 			return fmt.Errorf("using --mig-strategy=mixed is not supported with MPS")

internal/rm/allocate.go

@@ -78,3 +78,61 @@ func (r *resourceManager) distributedAlloc(available, required []string, size in
 
 	return devices, nil
 }
+
+// packedAlloc returns a list of devices such that any replicated devices are
+// packed onto as few physical GPUs as possible. It preferentially allocates
+// replicas from GPUs that already have the most allocated replicas, which
+// helps consolidate workloads and free up entire GPUs for other uses.
+func (r *resourceManager) packedAlloc(available, required []string, size int) ([]string, error) {
+	// Get the set of candidate devices as the difference between available and required.
+	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")
+	}
+
+	// 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 })
+	for _, c := range candidates {
+		id := AnnotatedID(c).GetID()
+		if _, exists := replicas[id]; !exists {
+			replicas[id] = &struct{ total, available int }{}
+		}
+		replicas[id].available++
+	}
+	for d := range r.devices {
+		id := AnnotatedID(d).GetID()
+		if _, exists := replicas[id]; !exists {
+			continue
+		}
+		replicas[id].total++
+	}
+
+	// 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 greatest difference between total and available
+	// replications (i.e. the most already allocated). This packs allocations
+	// onto GPUs that are already in use, freeing up other GPUs entirely.
+	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
+			return idiff > jdiff
+		})
+		id := AnnotatedID(candidates[0]).GetID()
+		replicas[id].available--
+		devices = append(devices, candidates[0])
+		candidates = candidates[1:]
+	}
+
+	// Add the set of required devices to this list and return it.
+	devices = append(required, devices...)
+
+	return devices, nil
+}