Harden vGPU startup rollback recovery

Author: yummybomb

cmd/api/main.go

@@ -172,26 +172,62 @@ func configureUFFDGraduationController(cfg *config.Config, instanceManager insta
 	}, logger), nil
 }
 
-func liveInstanceVGPUDevicePaths(ctx context.Context, instanceManager instances.Manager) (map[string]struct{}, error) {
+const vgpuAssignmentStartupGracePeriod = 5 * time.Minute
+
+func liveInstanceVGPUDevicePaths(ctx context.Context, instanceManager instances.Manager) (map[string]struct{}, time.Duration, error) {
 	allInstances, err := instanceManager.ListInstancesForReconcile(ctx)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 	protected := make(map[string]struct{})
+	var retryAfter time.Duration
 	for _, inst := range allInstances {
 		if inst.GPUDevicePath == "" {
 			continue
 		}
-		// A nil PID does not mean the assignment is orphaned: the PID is
-		// persisted only after the hypervisor starts, so a crash during boot
-		// leaves the device path without one. Only skip protection when the
-		// recorded hypervisor is known to be gone.
-		if inst.HypervisorPID != nil && !instances.HypervisorProcessExists(*inst.HypervisorPID, inst.SocketPath) {
+		if inst.HypervisorPID != nil {
+			if !instances.HypervisorProcessIdentityExists(*inst.HypervisorPID, inst.HypervisorStartTime, inst.SocketPath) {
+				continue
+			}
+			protected[inst.GPUDevicePath] = struct{}{}
+			continue
+		}
+		if inst.GPUAssignedAt == nil {
+			continue
+		}
+		remaining := vgpuAssignmentStartupGracePeriod - time.Since(*inst.GPUAssignedAt)
+		if remaining <= 0 {
 			continue
 		}
 		protected[inst.GPUDevicePath] = struct{}{}
+		if retryAfter == 0 || remaining < retryAfter {
+			retryAfter = remaining
+		}
 	}
-	return protected, nil
+	return protected, retryAfter, nil
+}
+
+func reconcileVGPUs(ctx context.Context, instanceManager instances.Manager, logger *slog.Logger) {
+	protected, retryAfter, err := liveInstanceVGPUDevicePaths(ctx, instanceManager)
+	if err != nil {
+		logger.Warn("failed to list instances for vGPU reconcile protection; skipping vendor VFIO reconciliation", "error", err)
+		return
+	}
+	if err := devices.ReconcileVGPUs(ctx, protected); err != nil {
+		logger.Warn("failed to reconcile vGPU devices", "error", err)
+	}
+	if retryAfter <= 0 {
+		return
+	}
+	go func() {
+		timer := time.NewTimer(retryAfter)
+		defer timer.Stop()
+		select {
+		case <-ctx.Done():
+		case <-timer.C:
+			reconcileVGPUs(ctx, instanceManager, logger)
+		}
+	}()
 }
 
 func run() error {
@@ -386,15 +422,7 @@ func run() error {
 
 	// Reconcile vGPU devices (clears orphaned vGPUs from previous runs)
 	logger.Info("Reconciling vGPU devices...")
-	protected, err := liveInstanceVGPUDevicePaths(app.Ctx, app.InstanceManager)
-	if err != nil {
-		logger.Warn("failed to list instances for vGPU reconcile protection; skipping vendor VFIO reconciliation", "error", err)
-		protected = nil
-	}
-	if err := devices.ReconcileVGPUs(app.Ctx, protected); err != nil {
-		// Log but don't fail - vGPU cleanup is best-effort
-		logger.Warn("failed to reconcile vGPU devices", "error", err)
-	}
+	reconcileVGPUs(ctx, app.InstanceManager, logger)
 
 	// Wire up resource validator for aggregate limit checking
 	// This enables the instance manager to validate CPU, memory, network, and GPU

cmd/api/main_test.go

@@ -351,20 +351,26 @@ func (s vgpuReconcileManagerStub) ListInstancesForReconcile(context.Context) ([]
 	return s.list, nil
 }
 
-// The hypervisor PID is persisted only after boot, so an assignment without
-// one may belong to a VM that is still starting and must stay protected.
-func TestLiveInstanceVGPUDevicePathsProtectsAssignmentsWithoutPID(t *testing.T) {
+func TestLiveInstanceVGPUDevicePathsBoundsProtectionWithoutPID(t *testing.T) {
 	dead := exec.Command("true")
 	require.NoError(t, dead.Run())
 	deadPID := dead.Process.Pid
+	recent := time.Now().Add(-time.Minute)
+	stale := time.Now().Add(-vgpuAssignmentStartupGracePeriod - time.Minute)
 
 	manager := vgpuReconcileManagerStub{list: []instances.Instance{
-		{StoredMetadata: instances.StoredMetadata{Id: "booting", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.4"}},
-		{StoredMetadata: instances.StoredMetadata{Id: "dead", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.5", HypervisorPID: &deadPID}},
+		{StoredMetadata: instances.StoredMetadata{Id: "booting", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.4", GPUAssignedAt: &recent}},
+		{StoredMetadata: instances.StoredMetadata{Id: "orphaned", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.5", GPUAssignedAt: &stale}},
+		{StoredMetadata: instances.StoredMetadata{Id: "legacy", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.6"}},
+		{StoredMetadata: instances.StoredMetadata{Id: "dead", GPUDevicePath: "/sys/bus/pci/devices/0000:82:00.7", HypervisorPID: &deadPID}},
 	}}
 
-	protected, err := liveInstanceVGPUDevicePaths(context.Background(), manager)
+	protected, retryAfter, err := liveInstanceVGPUDevicePaths(context.Background(), manager)
 	require.NoError(t, err)
+	require.Positive(t, retryAfter)
+	require.LessOrEqual(t, retryAfter, vgpuAssignmentStartupGracePeriod)
 	assert.Contains(t, protected, "/sys/bus/pci/devices/0000:82:00.4")
 	assert.NotContains(t, protected, "/sys/bus/pci/devices/0000:82:00.5")
+	assert.NotContains(t, protected, "/sys/bus/pci/devices/0000:82:00.6")
+	assert.NotContains(t, protected, "/sys/bus/pci/devices/0000:82:00.7")
 }

lib/instances/create.go

@@ -265,6 +265,7 @@ func (m *manager) createInstance(
 	var gpuFramework devices.VGPUFramework
 	var gpuDevicePath string
 	var gpuMdevUUID string
+	var gpuAssignedAt *time.Time
 	var stored *StoredMetadata
 	var retainedVGPU *StoredMetadata
 
@@ -306,6 +307,8 @@ func (m *manager) createInstance(
 		gpuFramework = gpuDevice.Framework
 		gpuDevicePath = gpuDevice.SysfsPath
 		gpuMdevUUID = gpuDevice.MdevUUID
+		assignedAt := m.nowUTC()
+		gpuAssignedAt = &assignedAt
 
 		// Add vGPU cleanup to stack
 		cu.Add(func() {
@@ -406,6 +409,7 @@ func (m *manager) createInstance(
 		GPUFramework:             gpuFramework,
 		GPUDevicePath:            gpuDevicePath,
 		GPUMdevUUID:              gpuMdevUUID,
+		GPUAssignedAt:            gpuAssignedAt,
 		Entrypoint:               req.Entrypoint,
 		Cmd:                      req.Cmd,
 		SkipKernelHeaders:        req.SkipKernelHeaders,

lib/instances/start.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/kernel/hypeman/lib/devices"
 	"github.com/kernel/hypeman/lib/egressproxy"
 	"github.com/kernel/hypeman/lib/instances/phasetracking"
 	"github.com/kernel/hypeman/lib/logger"
@@ -63,6 +62,8 @@ func (m *manager) startInstance(
 		}
 	}
 
+	rollbackMeta := *meta
+
 	// 2a. Clear stale exit info from previous run and apply command overrides
 	stored.ExitCode = nil
 	stored.ExitMessage = ""
@@ -167,26 +168,11 @@ func (m *manager) startInstance(
 			log.ErrorContext(ctx, "failed to create vGPU", "instance_id", id, "profile", stored.GPUProfile, "error", err)
 			return nil, fmt.Errorf("create vGPU for profile %s: %w", stored.GPUProfile, err)
 		}
-		setStoredVGPUDevice(stored, device)
+		assignedAt := m.nowUTC()
+		setStoredVGPUDevice(stored, device, assignedAt)
 		// Add vGPU cleanup to stack
 		cu.Add(func() {
-			assignment := devices.VGPUAssignment{
-				Framework:  device.Framework,
-				DevicePath: device.SysfsPath,
-				MdevUUID:   device.MdevUUID,
-				InstanceID: id,
-			}
-			if err := m.destroyVGPUAssignment(ctx, assignment); err != nil {
-				log.WarnContext(ctx, "failed to destroy vGPU on cleanup", "instance_id", id, "error", err)
-				if saveErr := m.saveMetadata(meta); saveErr != nil {
-					log.ErrorContext(ctx, "failed to retain vGPU assignment metadata after cleanup failure", "instance_id", id, "error", saveErr)
-				}
-			} else {
-				clearStoredVGPUDevice(stored)
-				if saveErr := m.saveMetadata(meta); saveErr != nil {
-					log.ErrorContext(ctx, "failed to save metadata after vGPU cleanup", "instance_id", id, "error", saveErr)
-				}
-			}
+			m.cleanupStartVGPU(ctx, id, device, assignedAt, rollbackMeta)
 		})
 		if err := m.saveMetadata(meta); err != nil {
 			log.ErrorContext(ctx, "failed to save metadata after vGPU creation", "instance_id", id, "error", err)

lib/instances/types.go

@@ -153,7 +153,8 @@ type StoredMetadata struct {
 	GPUProfile    string // vGPU profile name (e.g., "L40S-1Q")
 	GPUFramework  devices.VGPUFramework
 	GPUDevicePath string
-	GPUMdevUUID   string // populated for mdev-backed vGPUs
+	GPUMdevUUID   string     // populated for mdev-backed vGPUs
+	GPUAssignedAt *time.Time // set before hypervisor startup to bound crash recovery protection
 
 	// Command overrides (like docker run <image> <command>)
 	Entrypoint []string // Override image entrypoint (nil = use image default)

lib/instances/vgpu.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"time"
 
 	"github.com/kernel/hypeman/lib/devices"
 	"github.com/kernel/hypeman/lib/logger"
@@ -43,16 +44,40 @@ func (m *manager) destroyVGPUAssignment(ctx context.Context, assignment devices.
 	return destroy(ctx, assignment)
 }
 
-func setStoredVGPUDevice(stored *StoredMetadata, device *devices.VGPUDevice) {
+func setStoredVGPUDevice(stored *StoredMetadata, device *devices.VGPUDevice, assignedAt time.Time) {
 	stored.GPUFramework = device.Framework
 	stored.GPUDevicePath = device.SysfsPath
 	stored.GPUMdevUUID = device.MdevUUID
+	stored.GPUAssignedAt = &assignedAt
 }
 
 func clearStoredVGPUDevice(stored *StoredMetadata) {
 	stored.GPUFramework = devices.VGPUFrameworkNone
 	stored.GPUDevicePath = ""
 	stored.GPUMdevUUID = ""
+	stored.GPUAssignedAt = nil
+}
+
+func (m *manager) cleanupStartVGPU(ctx context.Context, instanceID string, device *devices.VGPUDevice, assignedAt time.Time, rollbackMeta metadata) {
+	assignment := devices.VGPUAssignment{
+		Framework:  device.Framework,
+		DevicePath: device.SysfsPath,
+		MdevUUID:   device.MdevUUID,
+		InstanceID: instanceID,
+	}
+	cleanupMeta := rollbackMeta
+	releaseErr := m.destroyVGPUAssignment(ctx, assignment)
+	if releaseErr != nil {
+		logger.FromContext(ctx).WarnContext(ctx, "failed to destroy vGPU on cleanup", "instance_id", instanceID, "error", releaseErr)
+		setStoredVGPUDevice(&cleanupMeta.StoredMetadata, device, assignedAt)
+	}
+	if err := m.saveMetadata(&cleanupMeta); err != nil {
+		message := "failed to save metadata after vGPU cleanup"
+		if releaseErr != nil {
+			message = "failed to retain vGPU assignment metadata after cleanup failure"
+		}
+		logger.FromContext(ctx).ErrorContext(ctx, message, "instance_id", instanceID, "error", err)
+	}
 }
 
 func (m *manager) releaseStoredVGPU(ctx context.Context, stored *StoredMetadata) error {

lib/instances/vgpu_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/kernel/hypeman/lib/devices"
 	"github.com/kernel/hypeman/lib/paths"
@@ -180,13 +181,68 @@ func TestStartRollbackRetainsVGPUAssignmentAfterFailedDestroy(t *testing.T) {
 	})
 
 	t.Setenv("TMPDIR", filepath.Join(t.TempDir(), "missing"))
-	_, err := m.startInstance(context.Background(), id, StartInstanceRequest{})
+	_, err := m.startInstance(context.Background(), id, StartInstanceRequest{Entrypoint: []string{"new-entrypoint"}})
 	require.Error(t, err)
 
 	stored, err := m.loadMetadata(id)
 	require.NoError(t, err)
 	assert.Equal(t, devices.VGPUFrameworkVendorVFIO, stored.GPUFramework)
 	assert.Equal(t, "/sys/bus/pci/devices/0000:82:00.4", stored.GPUDevicePath)
+	assert.NotNil(t, stored.GPUAssignedAt)
+	assert.Empty(t, stored.Entrypoint)
+}
+
+func TestCleanupStartVGPURestoresMetadataAfterBootFailure(t *testing.T) {
+	m := &manager{
+		paths: paths.New(t.TempDir()),
+		destroyVGPU: func(context.Context, devices.VGPUAssignment) error {
+			return nil
+		},
+	}
+	const id = "failed-start"
+	require.NoError(t, m.ensureDirectories(id))
+
+	previousStart := time.Now().Add(-time.Hour).UTC()
+	previousProgramStart := previousStart.Add(time.Second)
+	exitCode := 1
+	rollbackMeta := metadata{StoredMetadata: StoredMetadata{
+		Id:               id,
+		GPUProfile:       "NVIDIA L40S-2Q",
+		Entrypoint:       []string{"old-entrypoint"},
+		Cmd:              []string{"old-command"},
+		StartedAt:        &previousStart,
+		ProgramStartedAt: &previousProgramStart,
+		ExitCode:         &exitCode,
+		ExitMessage:      "previous exit",
+	}}
+
+	partial := rollbackMeta
+	partial.Entrypoint = []string{"new-entrypoint"}
+	partial.Cmd = []string{"new-command"}
+	partial.StartedAt = ptr(time.Now().UTC())
+	partial.ProgramStartedAt = nil
+	partial.ExitCode = nil
+	partial.ExitMessage = ""
+	assignedAt := time.Now().UTC()
+	device := &devices.VGPUDevice{
+		Framework: devices.VGPUFrameworkVendorVFIO,
+		SysfsPath: "/sys/bus/pci/devices/0000:82:00.4",
+	}
+	setStoredVGPUDevice(&partial.StoredMetadata, device, assignedAt)
+	require.NoError(t, m.saveMetadata(&partial))
+
+	m.cleanupStartVGPU(context.Background(), id, device, assignedAt, rollbackMeta)
+
+	stored, err := m.loadMetadata(id)
+	require.NoError(t, err)
+	assert.Equal(t, rollbackMeta.Entrypoint, stored.Entrypoint)
+	assert.Equal(t, rollbackMeta.Cmd, stored.Cmd)
+	assert.Equal(t, rollbackMeta.StartedAt, stored.StartedAt)
+	assert.Equal(t, rollbackMeta.ProgramStartedAt, stored.ProgramStartedAt)
+	assert.Equal(t, rollbackMeta.ExitCode, stored.ExitCode)
+	assert.Equal(t, rollbackMeta.ExitMessage, stored.ExitMessage)
+	assert.Empty(t, stored.GPUDevicePath)
+	assert.Nil(t, stored.GPUAssignedAt)
 }
 
 func TestVGPUAssignmentClaimedByLiveInstanceFailsOnInvalidMetadata(t *testing.T) {
@@ -333,16 +389,19 @@ func TestReleaseStoredVGPURetainsMetadataOnFailure(t *testing.T) {
 func TestSetAndClearStoredVGPUDevice(t *testing.T) {
 	t.Parallel()
 
+	assignedAt := time.Now().UTC()
 	stored := &StoredMetadata{}
 	setStoredVGPUDevice(stored, &devices.VGPUDevice{
 		Framework: devices.VGPUFrameworkVendorVFIO,
 		SysfsPath: "/sys/bus/pci/devices/0000:82:00.4",
-	})
+	}, assignedAt)
 	assert.Equal(t, devices.VGPUFrameworkVendorVFIO, stored.GPUFramework)
 	assert.Equal(t, "/sys/bus/pci/devices/0000:82:00.4", stored.GPUDevicePath)
+	assert.Equal(t, assignedAt, *stored.GPUAssignedAt)
 
 	clearStoredVGPUDevice(stored)
 	assert.Empty(t, stored.GPUFramework)
 	assert.Empty(t, stored.GPUDevicePath)
 	assert.Empty(t, stored.GPUMdevUUID)
+	assert.Nil(t, stored.GPUAssignedAt)
 }