Add spread allocation policy for maximum per-pod distinct-GPU coverage

[jonathan-meiri]

Summary

• Add spread as a third value for --shared-devices-allocation-policy, alongside distributed (default) and packed.
• spread prefers physical GPUs the current pod has not yet picked from, maximizing distinct-GPU coverage for a single multi-slot allocation.
• Default behavior (distributed) is unchanged.

Depends on #1826 (heap refactor) — the current diff includes that commit until #1826 lands, then collapses to just the spread policy addition. Opened as draft for now.

Contributed by @Meiri28 on behalf of @runatom-ai.

Motivation

distributed minimizes cluster-wide load imbalance — it prefers the GPU with the fewest replicas already allocated across all pods. packed does the opposite for bin-packing. Neither guarantees that a single pod's multi-slot request touches distinct physical GPUs.

Consider a node where GPU-0 has 3 free replicas (5 allocated) and GPU-1 has 5 free (3 allocated). A pod requests 2:

┌─────────────────────────────────────────────────────────────────────┐
│              distributed (current default)                          │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0                       GPU 1                                  │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ ███░░░░░ │      before                │
│  │  5 alloc │               │  3 alloc │                            │
│  └──────────┘               └──────────┘                            │
│                                                                     │
│  Pod requests 2 slots → distributed picks GPU with fewer allocated  │
│                                                                     │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ █████░░░ │      after                 │
│  │  5 alloc │               │  5 alloc │                            │
│  └──────────┘               └──────────┘                            │
│   ⚠ untouched                 ✓ 2 slots (pod concentrated)          │
│                                                                     │
│  → Pod runs entirely on GPU 1. No NVLink/NCCL across GPUs.          │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│              spread (new)                                           │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0                       GPU 1                                  │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ ███░░░░░ │      before                │
│  │  5 alloc │               │  3 alloc │                            │
│  └──────────┘               └──────────┘                            │
│                                                                     │
│  Pod requests 2 slots → spread prefers GPUs untouched by THIS pod   │
│                                                                     │
│  ┌──────────┐               ┌──────────┐                            │
│  │ ██████░░ │               │ ████░░░░ │      after                 │
│  │  6 alloc │               │  4 alloc │                            │
│  └──────────┘               └──────────┘                            │
│   ★ 1 slot                    ★ 1 slot (pod spans both GPUs)        │
│                                                                     │
│  → Multi-GPU workloads (NCCL, DDP, tensor-parallel) get real        │
│    hardware diversity for a single request.                         │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

Concrete workloads that benefit:

• Data-parallel / distributed training — each rank on a distinct physical GPU so ring-allreduce / NCCL spans real hardware
• Fault-isolated inference — a pod requesting 2 GPUs specifically to survive a single-card failure
• NVLink topology plays — a first step toward topology-aware allocation (weak preference: at least distinct cards)

@rajatchopra explicitly floated additional policies on #1788:

"we can decide on distributed allocation or a packed one, or a topology driven one"

This adds the third policy that #1621's abstraction was built for.

Design

Small refactor: the replicaComparator signature is enriched from func(i, j *replicaCount) bool to func(i, j *gpuAllocState) bool. Each policy now owns both its primary ordering and tie-break — the queue's Less is a pure passthrough:

distributed: func(i, j *gpuAllocState) bool {
    if i.count.allocated() != j.count.allocated() {
        return i.count.allocated() < j.count.allocated()   // primary
    }
    return i.pickedFrom < j.pickedFrom                     // tie-break
}
packed: func(i, j *gpuAllocState) bool {
    if i.count.allocated() != j.count.allocated() {
        return i.count.allocated() > j.count.allocated()   // primary
    }
    return i.pickedFrom < j.pickedFrom                     // tie-break
}
spread: func(i, j *gpuAllocState) bool {
    if i.pickedFrom != j.pickedFrom {
        return i.pickedFrom < j.pickedFrom                 // primary
    }
    return i.count.allocated() < j.count.allocated()       // tie-break
}

distributed and packed are behavior-preserving — the primary+tie-break they used to get from greedyAlloc's wrapping is now spelled out in the comparator body itself. All existing tests (TestDistributedAlloc, TestPackedAlloc, TestPackedVsDistributedContrast) pass unchanged.

Usage

--shared-devices-allocation-policy=spread
# or
SHARED_DEVICES_ALLOCATION_POLICY=spread

Config file:

version: v1
flags:
  plugin:
    sharedDevicesAllocationPolicy: spread

Commits are DCO-signed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.