fix(bedrock): treat a claimed GPU as unavailable in the capacity gate (#748)

checkCapacity read only PoolEntry.Available, but that field tracks device
health: Release clears it when a vfio rebind fails, so a device that cannot
be handed back is not offered again. Occupancy is the separate InstanceID
field, which PoolEntry documents as empty when free, and gpu.Manager.Claim
tests both before taking a device.

Reading Available alone let a GPU already held by a running serving VM admit.
Ensure never took its capacity refusal, so evictForCapacity was never called
and LRU eviction never fired: on a single-GPU node the launch instead
proceeded to STARTING and failed later with InsufficientInstanceCapacity,
with no eviction attempted and no reason logged.

The existing table had a case named 'already claimed' that set Available to
false, which is not the shape Claim leaves behind. That encoded the wrong
model and is why the gap went unnoticed, so it is renamed to what it really
covers and the real claimed-device shapes are added alongside.

Also stop daemonGPUSnapshotter capturing d.gpuManager by value.
applyGPUConfig swaps that pointer on a passthrough toggle and re-points
daemonGPUClaimer but not the snapshotter, which would leave bedrock reading
an orphaned pool that sees no claims and admits everything.
This commit is contained in:
Josh Siviour
2026-08-07 14:37:32 +10:00
committed by GitHub
parent c7fdd23b7e
commit c0f3e0bee2
3 changed files with 80 additions and 11 deletions
+13 -3
View File
@@ -70,7 +70,7 @@ func (d *Daemon) buildBedrockServiceDeps() handlers_bedrock.ServiceDeps {
Replicas: clusterSize,
}
if d.gpuManager != nil {
deps.GPU = &daemonGPUSnapshotter{mgr: d.gpuManager}
deps.GPU = &daemonGPUSnapshotter{d: d}
}
return deps
}
@@ -78,10 +78,20 @@ func (d *Daemon) buildBedrockServiceDeps() handlers_bedrock.ServiceDeps {
// daemonGPUSnapshotter adapts *gpu.Manager to handlers_bedrock's Snapshot-only
// capacity-check surface, keeping the gpu package out of that handler package
// the same way daemonGPUClaimer keeps it out of handlers_ec2_instance.
//
// It holds the daemon rather than the manager because applyGPUConfig swaps
// d.gpuManager on a passthrough toggle. Capturing the pointer here would leave
// this reading an orphaned pool that sees no claims and so admits everything.
type daemonGPUSnapshotter struct {
mgr *gpu.Manager
d *Daemon
}
func (g *daemonGPUSnapshotter) Snapshot() []gpu.PoolEntry {
return g.mgr.Snapshot()
g.d.mu.Lock()
mgr := g.d.gpuManager
g.d.mu.Unlock()
if mgr == nil {
return nil
}
return mgr.Snapshot()
}
+11 -7
View File
@@ -15,12 +15,16 @@ type gpuSnapshotter interface {
Snapshot() []gpu.PoolEntry
}
// checkCapacity reports whether the free pool (whole devices or MIG slices
// with Available=true) has at least one entry with minVRAMMiB free. It is a
// fast pre-admission check only: the deeper claim/allocate decision still
// happens inside the shared RunInstances path when the VM actually launches,
// so a device can still be lost to a racing launch between this check and
// that claim — the guest fails to boot in that rare case, not silently OOMs.
// checkCapacity reports whether the free pool (whole devices or MIG slices)
// has at least one entry with minVRAMMiB free. It is a fast pre-admission
// check only: the deeper claim/allocate decision still happens inside the
// shared RunInstances path when the VM actually launches, so a device can
// still be lost to a racing launch between this check and that claim — the
// guest fails to boot in that rare case, not silently OOMs.
//
// Free means both axes the pool tracks: Available is device health, which
// Release clears when a vfio rebind fails, while InstanceID is occupancy.
// Claim tests both, and so must this — a claimed device keeps Available=true.
func checkCapacity(snapshot []gpu.PoolEntry, minVRAMMiB int) error {
if minVRAMMiB <= 0 {
return fmt.Errorf("bedrock: invalid MinVRAMMiB %d", minVRAMMiB)
@@ -31,7 +35,7 @@ func checkCapacity(snapshot []gpu.PoolEntry, minVRAMMiB int) error {
var unknownVRAM *gpu.PoolEntry
for i := range snapshot {
entry := &snapshot[i]
if !entry.Available {
if !entry.Available || entry.InstanceID != "" {
continue
}
memMiB := entry.Device.MemoryMiB
+56 -1
View File
@@ -31,13 +31,40 @@ func TestCheckCapacity(t *testing.T) {
minVRAM: 5120,
},
{
name: "device exists but not available (already claimed)",
name: "device unhealthy (rebind failed, Available cleared)",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: false},
},
minVRAM: 5120,
wantErr: true,
},
{
// The shape Claim actually leaves behind: it records InstanceID and
// leaves Available alone, so occupancy is invisible to a check that
// reads Available on its own.
name: "whole device claimed by a running instance",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true, InstanceID: "i-abc"},
},
minVRAM: 5120,
wantErr: true,
},
{
name: "MIG slice claimed by a running instance",
snapshot: []gpu.PoolEntry{
{MIGInstance: &gpu.MIGInstance{Profile: gpu.MIGProfile{MemoryMiB: 10240}}, Available: true, InstanceID: "i-abc"},
},
minVRAM: 5120,
wantErr: true,
},
{
name: "a claimed device does not hide a free sibling",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 16384}, Available: true, InstanceID: "i-abc"},
{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true},
},
minVRAM: 5120,
},
{
name: "available device too small",
snapshot: []gpu.PoolEntry{
@@ -79,6 +106,34 @@ func TestCheckCapacity(t *testing.T) {
}
}
// TestCheckCapacity_ClaimedDeviceIsNotFree pins the defect that kept Ochre's
// LRU eviction dormant: Claim records InstanceID and deliberately leaves
// Available set, because Available tracks device health and Release clears it
// only when a vfio rebind fails. Reading Available alone let a GPU held by a
// running serving VM admit, so Ensure never took its capacity refusal and
// evictForCapacity was never reached.
func TestCheckCapacity_ClaimedDeviceIsNotFree(t *testing.T) {
pool := []gpu.PoolEntry{
{Device: gpu.GPUDevice{PCIAddress: "0000:5e:00.0", MemoryMiB: 8188}, Available: true, InstanceID: "i-abc"},
}
require.Error(t, checkCapacity(pool, 7168), "a GPU held by a running instance must not admit")
pool[0].InstanceID = ""
assert.NoError(t, checkCapacity(pool, 7168), "the same device must admit once released")
}
// TestCheckCapacity_UnknownVRAM_IgnoredWhenClaimed proves occupancy is settled
// before the 0 MiB data-gap branch. A claimed device is not a candidate at all,
// so an operator gets plain exhaustion rather than being sent to chase a
// missing VRAM figure on a card that is merely busy.
func TestCheckCapacity_UnknownVRAM_IgnoredWhenClaimed(t *testing.T) {
err := checkCapacity([]gpu.PoolEntry{
{Device: gpu.GPUDevice{PCIAddress: "0000:03:00.0", MemoryMiB: 0}, Available: true, InstanceID: "i-abc"},
}, 5120)
require.Error(t, err)
assert.NotContains(t, err.Error(), "undiscoverable")
}
func TestCheckCapacity_ErrorCode(t *testing.T) {
err := checkCapacity(nil, 5120)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))