Wait for mdev parents in vGPU manager validation

cmd/nvidia-validator/main.go

@@ -240,8 +240,8 @@ const (
 	wslNvidiaSMIPath = "/usr/lib/wsl/lib/nvidia-smi"
 	// shell indicates what shell to use when invoking commands in a subprocess
 	shell = "sh"
-	// defaultVFWaitTimeout is the default timeout for waiting for VFs to be created
-	defaultVFWaitTimeout = 5 * time.Minute
+	// defaultVGPUReadinessTimeout is the default timeout for waiting for the vGPU stack to be ready
+	defaultVGPUReadinessTimeout = 5 * time.Minute
 	// constants for driver components
 	GDRCOPY       = "gdrcopy"
 	NVIDIAFS      = "nvidia-fs"
@@ -1741,9 +1741,9 @@ func (v *VGPUManager) validate() error {
 		return err
 	}
 
-	log.Info("Waiting for VFs to be available...")
-	if err := waitForVFs(ctx, defaultVFWaitTimeout); err != nil {
-		return fmt.Errorf("vGPU Manager VFs not ready: %w", err)
+	log.Info("Waiting for parent devices to be available...")
+	if err := waitForParentDevices(ctx, defaultVGPUReadinessTimeout); err != nil {
+		return fmt.Errorf("vGPU Manager parent devices not ready: %w", err)
 	}
 
 	statusFile := vGPUManagerStatusFile
@@ -1777,43 +1777,77 @@ func (v *VGPUManager) runValidation(silent bool) (hostDriver bool, err error) {
 	return hostDriver, runCommand(command, args, silent)
 }
 
-// waitForVFs waits for Virtual Functions to be created on all NVIDIA GPUs.
-// It polls sriov_numvfs until all GPUs have their full VF count enabled.
-func waitForVFs(ctx context.Context, timeout time.Duration) error {
+// waitForParentDevices polls until the vGPU stack is ready — either NVIDIA
+// mdev parent devices have been registered (PF on Turing, VFs on Ampere+
+// SR-IOV) or all SR-IOV VFs are enabled.
+func waitForParentDevices(ctx context.Context, timeout time.Duration) error {
 	pollInterval := time.Duration(sleepIntervalSecondsFlag) * time.Second
 	nvpciLib := nvpci.New()
+	nvmdevLib := nvmdev.New()
 
 	return wait.PollUntilContextTimeout(ctx, pollInterval, timeout, true, func(ctx context.Context) (bool, error) {
 		gpus, err := nvpciLib.GetGPUs()
 		if err != nil {
-			log.Warnf("Error getting GPUs: %v", err)
+			log.Warnf("error getting GPUs: %v", err)
 			return false, nil
 		}
-
-		var totalExpected, totalEnabled uint64
-		var pfCount int
-		for _, gpu := range gpus {
-			sriovInfo := gpu.SriovInfo
-			if sriovInfo.IsPF() {
-				pfCount++
-				totalExpected += sriovInfo.PhysicalFunction.TotalVFs
-				totalEnabled += sriovInfo.PhysicalFunction.NumVFs
-			}
+		if isDriverUsingSRIOV(gpus) {
+			return AreAllVFsReady(gpus), nil
 		}
+		return mdevParentDevicesExist(nvmdevLib), nil
+	})
+}
 
-		if totalExpected == 0 {
-			log.Info("No SR-IOV capable GPUs found, skipping VF wait")
-			return true, nil
+// isDriverUsingSRIOV reports whether any GPU is an SR-IOV PF with VFs enabled;
+// AreAllVFsReady checks completeness across all PFs.
+func isDriverUsingSRIOV(gpus []*nvpci.NvidiaPCIDevice) bool {
+	for _, gpu := range gpus {
+		if gpu.SriovInfo.IsPF() && gpu.SriovInfo.PhysicalFunction.NumVFs > 0 {
+			return true
 		}
+	}
+	return false
+}
 
-		if totalEnabled == totalExpected {
-			log.Infof("All %d VF(s) enabled on %d NVIDIA GPU(s)", totalEnabled, pfCount)
-			return true, nil
+func mdevParentDevicesExist(nvmdevLib nvmdev.Interface) bool {
+	parents, err := nvmdevLib.GetAllParentDevices()
+	if err != nil {
+		log.Warnf("could not get mdev parent devices: %v", err)
+		return false
+	}
+	if len(parents) == 0 {
+		log.Info("found 0 mdev parent devices")
+		return false
+	}
+	log.Infof("found %d mdev parent devices", len(parents))
+	return true
+}
+
+// AreAllVFsReady reports whether every SR-IOV PF has all of its VFs enabled.
+func AreAllVFsReady(gpus []*nvpci.NvidiaPCIDevice) bool {
+	var totalExpected, totalEnabled uint64
+	var pfCount int
+	for _, gpu := range gpus {
+		sriovInfo := gpu.SriovInfo
+		if sriovInfo.IsPF() {
+			pfCount++
+			totalExpected += sriovInfo.PhysicalFunction.TotalVFs
+			totalEnabled += sriovInfo.PhysicalFunction.NumVFs
 		}
+	}
 
-		log.Infof("Waiting for VFs: %d/%d enabled across %d GPU(s)", totalEnabled, totalExpected, pfCount)
-		return false, nil
-	})
+	if totalExpected == 0 {
+		log.Info("no SR-IOV capable GPUs found")
+		return false
+	}
+
+	if totalEnabled == totalExpected {
+		log.Infof("all %d VF(s) enabled on %d NVIDIA GPU(s)", totalEnabled, pfCount)
+		return true
+	}
+
+	log.Infof("not all VFs have been created. %d/%d enabled across %d GPU(s)", totalEnabled, totalExpected, pfCount)
+	return false
 }
 
 func (c *CCManager) validate() error {

cmd/nvidia-validator/main_test.go

@@ -25,6 +25,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/NVIDIA/go-nvlib/pkg/nvmdev"
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/stretchr/testify/require"
 )
 
@@ -376,3 +378,110 @@ UNKNOWN_FEATURE: true`,
 		})
 	}
 }
+
+func newTestPF(totalVFs, numVFs uint64) *nvpci.NvidiaPCIDevice {
+	return &nvpci.NvidiaPCIDevice{
+		SriovInfo: nvpci.SriovInfo{
+			PhysicalFunction: &nvpci.SriovPhysicalFunction{
+				TotalVFs: totalVFs,
+				NumVFs:   numVFs,
+			},
+		},
+	}
+}
+
+func TestIsDriverUsingSRIOV(t *testing.T) {
+	nonSriovGPU := &nvpci.NvidiaPCIDevice{}
+
+	testCases := []struct {
+		description string
+		gpus        []*nvpci.NvidiaPCIDevice
+		expected    bool
+	}{
+		{
+			description: "no GPUs",
+			gpus:        nil,
+			expected:    false,
+		},
+		{
+			description: "non-SRIOV GPU",
+			gpus:        []*nvpci.NvidiaPCIDevice{nonSriovGPU},
+			expected:    false,
+		},
+		{
+			description: "SRIOV-capable PF with no VFs enabled",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 0)},
+			expected:    false,
+		},
+		{
+			description: "PF with VFs enabled",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 16)},
+			expected:    true,
+		},
+		{
+			description: "mixed non-SRIOV GPU and PF with VFs enabled",
+			gpus:        []*nvpci.NvidiaPCIDevice{nonSriovGPU, newTestPF(16, 4)},
+			expected:    true,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			require.Equal(t, tc.expected, isDriverUsingSRIOV(tc.gpus))
+		})
+	}
+}
+
+func TestAreAllVFsReady(t *testing.T) {
+	testCases := []struct {
+		description string
+		gpus        []*nvpci.NvidiaPCIDevice
+		expected    bool
+	}{
+		{
+			description: "no GPUs",
+			gpus:        nil,
+			expected:    false,
+		},
+		{
+			description: "non-SRIOV GPU only",
+			gpus:        []*nvpci.NvidiaPCIDevice{{}},
+			expected:    false,
+		},
+		{
+			description: "PF with only some VFs enabled",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 4)},
+			expected:    false,
+		},
+		{
+			description: "PF with all VFs enabled",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 16)},
+			expected:    true,
+		},
+		{
+			description: "one PF complete, one PF incomplete",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 16), newTestPF(16, 0)},
+			expected:    false,
+		},
+		{
+			description: "all PFs complete",
+			gpus:        []*nvpci.NvidiaPCIDevice{newTestPF(16, 16), newTestPF(8, 8)},
+			expected:    true,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			require.Equal(t, tc.expected, AreAllVFsReady(tc.gpus))
+		})
+	}
+}
+
+func TestMdevParentDevicesExist(t *testing.T) {
+	mock, err := nvmdev.NewMock()
+	require.NoError(t, err)
+	defer mock.Cleanup()
+
+	require.False(t, mdevParentDevicesExist(mock))
+
+	require.NoError(t, mock.AddMockA100Parent("0000:3b:00.0", 0))
+	require.True(t, mdevParentDevicesExist(mock))
+}