Harden vGPU cleanup failure handling

cmd/api/api/instances.go

@@ -343,7 +343,16 @@ func (s *ApiService) CreateInstance(ctx context.Context, request oapi.CreateInst
 
 	inst, err := s.InstanceManager.CreateInstance(ctx, domainReq)
 	if err != nil {
+		var vgpuPending *instances.VGPUCleanupPendingError
 		switch {
+		// Checked first: it wraps the original create error, so a later
+		// errors.Is case would match the cause and hide the retained instance.
+		case errors.As(err, &vgpuPending):
+			log.ErrorContext(ctx, "failed to create instance", "error", err, "image", request.Body.Image)
+			return oapi.CreateInstance500JSONResponse{
+				Code:    "vgpu_cleanup_pending",
+				Message: fmt.Sprintf("failed to create instance; vGPU release failed during rollback and instance %s retains the assignment, delete it to retry", vgpuPending.InstanceID),
+			}, nil
 		case errors.Is(err, instances.ErrImageNotReady):
 			return oapi.CreateInstance400JSONResponse{
 				Code:    "image_not_ready",

cmd/api/api/instances_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/kernel/hypeman/lib/instances"
 	"github.com/kernel/hypeman/lib/instances/phasetracking"
 	mw "github.com/kernel/hypeman/lib/middleware"
+	"github.com/kernel/hypeman/lib/network"
 	"github.com/kernel/hypeman/lib/oapi"
 	"github.com/kernel/hypeman/lib/paths"
 	restartpolicy "github.com/kernel/hypeman/lib/restart-policy"
@@ -46,6 +47,36 @@ func TestGetInstance_NotFound(t *testing.T) {
 	require.Error(t, err)
 }
 
+type createErrorInstanceManager struct {
+	instances.Manager
+	err error
+}
+
+func (m createErrorInstanceManager) CreateInstance(context.Context, instances.CreateInstanceRequest) (*instances.Instance, error) {
+	return nil, m.err
+}
+
+// A retained-assignment error must win over the mapping of the create error
+// it wraps, or the response omits the instance the caller has to delete.
+func TestCreateInstance_VGPUCleanupPendingBeatsWrappedErrorMapping(t *testing.T) {
+	t.Parallel()
+	svc := newTestService(t)
+	svc.InstanceManager = createErrorInstanceManager{err: &instances.VGPUCleanupPendingError{
+		InstanceID: "inst-1",
+		Err:        network.ErrNameExists,
+	}}
+
+	resp, err := svc.CreateInstance(ctx(), oapi.CreateInstanceRequestObject{
+		Body: &oapi.CreateInstanceRequest{Image: "test-image"},
+	})
+	require.NoError(t, err)
+
+	pending, ok := resp.(oapi.CreateInstance500JSONResponse)
+	require.True(t, ok, "expected 500 vgpu_cleanup_pending, got %T", resp)
+	assert.EqualValues(t, "vgpu_cleanup_pending", pending.Code)
+	assert.Contains(t, pending.Message, "inst-1")
+}
+
 func TestCreateInstance_AutoPullImage(t *testing.T) {
 	t.Parallel()
 	if _, err := os.Stat("/dev/kvm"); os.IsNotExist(err) {

cmd/api/main.go

@@ -172,6 +172,28 @@ func configureUFFDGraduationController(cfg *config.Config, instanceManager insta
 	}, logger), nil
 }
 
+func liveInstanceVGPUDevicePaths(ctx context.Context, instanceManager instances.Manager) (map[string]struct{}, error) {
+	allInstances, err := instanceManager.ListInstancesForReconcile(ctx)
+	if err != nil {
+		return nil, err
+	}
+	protected := make(map[string]struct{})
+	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) {
+			continue
+		}
+		protected[inst.GPUDevicePath] = struct{}{}
+	}
+	return protected, nil
+}
+
 func run() error {
 	// Load config early for OTel initialization
 	// Config path can be specified via CONFIG_PATH env var or defaults to platform-specific locations
@@ -362,11 +384,16 @@ func run() error {
 		return fmt.Errorf("reconcile device state: %w", err)
 	}
 
-	// Reconcile mdev devices (clears orphaned vGPUs from previous runs)
-	logger.Info("Reconciling mdev devices...")
-	if err := devices.ReconcileMdevs(app.Ctx, nil); err != nil {
-		// Log but don't fail - mdev cleanup is best-effort
-		logger.Warn("failed to reconcile mdev devices", "error", err)
+	// 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)
 	}
 
 	// Wire up resource validator for aggregate limit checking

cmd/api/main_test.go

@@ -2,15 +2,18 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"os/exec"
 	"testing"
 	"time"
 
 	"github.com/getkin/kin-openapi/openapi3filter"
 	"github.com/go-chi/chi/v5"
 	"github.com/golang-jwt/jwt/v5"
+	"github.com/kernel/hypeman/lib/instances"
 	mw "github.com/kernel/hypeman/lib/middleware"
 	"github.com/kernel/hypeman/lib/oapi"
 	nethttpmiddleware "github.com/oapi-codegen/nethttp-middleware"
@@ -338,3 +341,30 @@ func TestImageNameWithSlashes_URLEncoding(t *testing.T) {
 		})
 	}
 }
+
+type vgpuReconcileManagerStub struct {
+	instances.Manager
+	list []instances.Instance
+}
+
+func (s vgpuReconcileManagerStub) ListInstancesForReconcile(context.Context) ([]instances.Instance, error) {
+	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) {
+	dead := exec.Command("true")
+	require.NoError(t, dead.Run())
+	deadPID := dead.Process.Pid
+
+	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}},
+	}}
+
+	protected, err := liveInstanceVGPUDevicePaths(context.Background(), manager)
+	require.NoError(t, err)
+	assert.Contains(t, protected, "/sys/bus/pci/devices/0000:82:00.4")
+	assert.NotContains(t, protected, "/sys/bus/pci/devices/0000:82:00.5")
+}

integration/vgpu_test.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"context"
 	"os"
+	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -21,21 +23,23 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// TestVGPU is an integration test that verifies vGPU (SR-IOV mdev) support works.
+// TestVGPU is an integration test that verifies vGPU (SR-IOV) support works
+// on the host's framework: mdev or NVIDIA's vendor-specific VFIO.
 //
 // This test automatically detects vGPU availability and skips if:
-//   - No SR-IOV VFs are found in /sys/class/mdev_bus/
+//   - No vGPU framework (mdev or vendor VFIO) is discovered
 //   - No vGPU profiles are available
-//   - Not running as root (required for mdev creation)
+//   - Not running as root (required for sysfs vGPU assignment)
 //   - KVM is not available
 //
 // To run manually:
 //
 //	sudo go test -v -run TestVGPU -timeout 5m ./integration/...
 //
-// Note: This test verifies mdev creation and PCI device visibility inside the VM.
-// It does NOT test nvidia-smi or CUDA functionality since that requires NVIDIA
-// guest drivers pre-installed in the image.
+// Note: This test verifies vGPU assignment, release on stop, reacquisition on
+// start, and PCI device visibility inside the VM. It does NOT test nvidia-smi
+// or CUDA functionality since that requires NVIDIA guest drivers pre-installed
+// in the image.
 func TestVGPU(t *testing.T) {
 	t.Parallel()
 	if testing.Short() {
@@ -84,9 +88,15 @@ func TestVGPU(t *testing.T) {
 
 	// Cleanup any orphaned instances and mdevs
 	t.Cleanup(func() {
-		if instanceID != "" {
-			t.Log("Cleanup: Deleting instance...")
-			instanceManager.DeleteInstance(ctx, instanceID)
+		if instanceID == "" {
+			return
+		}
+		if _, err := instanceManager.StopInstance(ctx, instanceID); err != nil {
+			t.Logf("Cleanup: stop instance: %v", err)
+		}
+		t.Log("Cleanup: Deleting instance...")
+		if err := instanceManager.DeleteInstance(ctx, instanceID); err != nil {
+			t.Errorf("cleanup: delete instance: %v", err)
 		}
 	})
 
@@ -159,9 +169,18 @@ func TestVGPU(t *testing.T) {
 	instanceID = inst.Id
 	t.Logf("Instance created: %s", inst.Id)
 
-	// Verify mdev UUID was assigned
-	require.NotEmpty(t, inst.GPUMdevUUID, "Instance should have mdev UUID assigned")
-	t.Logf("mdev UUID: %s", inst.GPUMdevUUID)
+	// Verify the assignment matches the host's framework
+	require.NotEmpty(t, inst.GPUDevicePath, "Instance should have a vGPU device path assigned")
+	switch inst.GPUFramework {
+	case devices.VGPUFrameworkMdev:
+		require.NotEmpty(t, inst.GPUMdevUUID, "mdev instance should have a UUID assigned")
+		t.Logf("mdev UUID: %s", inst.GPUMdevUUID)
+	case devices.VGPUFrameworkVendorVFIO:
+		require.Empty(t, inst.GPUMdevUUID, "vendor VFIO instance should not have an mdev UUID")
+		t.Logf("vendor VFIO VF: %s", inst.GPUDevicePath)
+	default:
+		t.Fatalf("unexpected vGPU framework %q", inst.GPUFramework)
+	}
 
 	// Step 5: Check GPU resources AFTER creating instance
 	t.Run("ResourcesDecrementedAfterCreation", func(t *testing.T) {
@@ -180,12 +199,9 @@ func TestVGPU(t *testing.T) {
 		assert.Less(t, availableAfter, availableBefore, "available instances should decrease after creating VM")
 	})
 
-	// Step 6: Verify mdev was created in sysfs
-	t.Run("MdevCreated", func(t *testing.T) {
-		mdevPath := "/sys/bus/mdev/devices/" + inst.GPUMdevUUID
-		_, err := os.Stat(mdevPath)
-		assert.NoError(t, err, "mdev device should exist at %s", mdevPath)
-		t.Logf("mdev exists at: %s", mdevPath)
+	// Step 6: Verify the assignment exists in sysfs
+	t.Run("VGPUAssignedInSysfs", func(t *testing.T) {
+		assertVGPUAssigned(t, inst.GPUFramework, inst.GPUDevicePath)
 	})
 
 	// Step 7: Wait for guest agent to be ready
@@ -225,13 +241,68 @@ func TestVGPU(t *testing.T) {
 		require.NoError(t, err)
 
 		assert.Equal(t, profile, actualInst.GPUProfile, "GPU profile should match")
-		assert.NotEmpty(t, actualInst.GPUMdevUUID, "mdev UUID should be set")
-		t.Logf("Instance GPU: profile=%s, mdev=%s", actualInst.GPUProfile, actualInst.GPUMdevUUID)
+		assert.Equal(t, inst.GPUFramework, actualInst.GPUFramework, "framework should match")
+		assert.NotEmpty(t, actualInst.GPUDevicePath, "device path should be set")
+		if inst.GPUFramework == devices.VGPUFrameworkMdev {
+			assert.NotEmpty(t, actualInst.GPUMdevUUID, "mdev UUID should be set")
+		}
+		t.Logf("Instance GPU: profile=%s, framework=%s, device=%s", actualInst.GPUProfile, actualInst.GPUFramework, actualInst.GPUDevicePath)
+	})
+
+	t.Log("Step 10: Stopping instance...")
+	_, err = instanceManager.StopInstance(ctx, inst.Id)
+	require.NoError(t, err, "stop should succeed")
+
+	t.Run("VGPUReleasedOnStop", func(t *testing.T) {
+		stopped, err := instanceManager.GetInstance(ctx, inst.Id)
+		require.NoError(t, err)
+		assert.Empty(t, stopped.GPUDevicePath, "assignment metadata should be cleared on stop")
+		assertVGPUReleased(t, inst.GPUFramework, inst.GPUDevicePath)
+	})
+
+	t.Log("Step 11: Starting instance to reacquire a vGPU...")
+	started, err := instanceManager.StartInstance(ctx, inst.Id, instances.StartInstanceRequest{})
+	require.NoError(t, err, "start should succeed")
+
+	t.Run("VGPUReacquiredOnStart", func(t *testing.T) {
+		require.NotEmpty(t, started.GPUDevicePath, "start should assign a vGPU")
+		assert.Equal(t, inst.GPUFramework, started.GPUFramework, "framework should match")
+		assertVGPUAssigned(t, started.GPUFramework, started.GPUDevicePath)
 	})
 
 	t.Log("✅ vGPU test PASSED!")
 }
 
+func assertVGPUAssigned(t *testing.T, framework devices.VGPUFramework, devicePath string) {
+	t.Helper()
+	switch framework {
+	case devices.VGPUFrameworkMdev:
+		_, err := os.Stat(devicePath)
+		assert.NoError(t, err, "mdev device should exist at %s", devicePath)
+	case devices.VGPUFrameworkVendorVFIO:
+		data, err := os.ReadFile(filepath.Join(devicePath, "nvidia", "current_vgpu_type"))
+		require.NoError(t, err, "VF should expose current_vgpu_type")
+		assert.NotEqual(t, "0", strings.TrimSpace(string(data)), "VF should have a vGPU type assigned")
+	default:
+		t.Fatalf("unexpected vGPU framework %q", framework)
+	}
+}
+
+func assertVGPUReleased(t *testing.T, framework devices.VGPUFramework, devicePath string) {
+	t.Helper()
+	switch framework {
+	case devices.VGPUFrameworkMdev:
+		_, err := os.Stat(devicePath)
+		assert.True(t, os.IsNotExist(err), "mdev device should be gone from %s", devicePath)
+	case devices.VGPUFrameworkVendorVFIO:
+		data, err := os.ReadFile(filepath.Join(devicePath, "nvidia", "current_vgpu_type"))
+		require.NoError(t, err, "VF should expose current_vgpu_type")
+		assert.Equal(t, "0", strings.TrimSpace(string(data)), "VF assignment should be released")
+	default:
+		t.Fatalf("unexpected vGPU framework %q", framework)
+	}
+}
+
 // checkVGPUTestPrerequisites checks if vGPU test can run.
 // Returns (skipReason, profileName) - skipReason is empty if all prerequisites are met.
 func checkVGPUTestPrerequisites() (string, string) {
@@ -245,10 +316,13 @@ func checkVGPUTestPrerequisites() (string, string) {
 		return "vGPU test requires root (sudo) for mdev creation", ""
 	}
 
-	// Check for vGPU mode (SR-IOV VFs present)
-	mode := devices.DetectHostGPUMode()
-	if mode != devices.GPUModeVGPU {
-		return "vGPU test requires SR-IOV VFs in /sys/class/mdev_bus/", ""
+	// Check for a vGPU framework (mdev or vendor VFIO)
+	framework, _, err := devices.DiscoverVGPU()
+	if err != nil {
+		return "vGPU test failed to discover vGPU framework: " + err.Error(), ""
+	}
+	if framework == devices.VGPUFrameworkNone {
+		return "vGPU test requires SR-IOV VFs with an mdev or vendor VFIO vGPU framework", ""
 	}
 
 	// Check for available profiles

lib/builds/manager_test.go

@@ -51,6 +51,10 @@ func (m *mockInstanceManager) ListInstances(ctx context.Context, filter *instanc
 	return result, nil
 }
 
+func (m *mockInstanceManager) ListInstancesForReconcile(ctx context.Context) ([]instances.Instance, error) {
+	return m.ListInstances(ctx, nil)
+}
+
 func (m *mockInstanceManager) ListSnapshots(ctx context.Context, filter *instances.ListSnapshotsFilter) ([]instances.Snapshot, error) {
 	return nil, nil
 }