Merge pull request #143 from NVIDIA/load-modules-in-vfio-manage

[vfio-manage] load vfio-pci module before binding GPUs to it

Author: cdesiniotis

cmd/vfio-manage/bind.go

@@ -30,72 +30,80 @@ import (
 type bindCommand struct {
 	logger   *logrus.Logger
 	nvpciLib nvpci.Interface
+	options  bindOptions
 }
 
 type bindOptions struct {
 	all      bool
 	deviceID string
+	hostRoot string
 }
 
 // newBindCommand constructs a bind command with the specified logger
 func newBindCommand(logger *logrus.Logger) *cli.Command {
 	c := bindCommand{
 		logger: logger,
-		nvpciLib: nvpci.New(
-			nvpci.WithLogger(logger),
-		),
 	}
 	return c.build()
 }
 
 // build the bind command
 func (m bindCommand) build() *cli.Command {
-	cfg := bindOptions{}
-
-	// Create the 'bind' command
 	c := cli.Command{
 		Name:  "bind",
 		Usage: "Bind device(s) to vfio-pci driver",
 		Before: func(c *cli.Context) error {
-			return m.validateFlags(&cfg)
+			return m.validateFlags()
 		},
 		Action: func(c *cli.Context) error {
-			return m.run(&cfg)
+			return m.run()
 		},
 		Flags: []cli.Flag{
 			&cli.BoolFlag{
 				Name:        "all",
 				Aliases:     []string{"a"},
-				Destination: &cfg.all,
+				Destination: &m.options.all,
 				Usage:       "Bind all NVIDIA devices to vfio-pci",
 			},
 			&cli.StringFlag{
 				Name:        "device-id",
 				Aliases:     []string{"d"},
-				Destination: &cfg.deviceID,
+				Destination: &m.options.deviceID,
 				Usage:       "Specific device ID to bind (e.g., 0000:01:00.0)",
 			},
+			&cli.StringFlag{
+				Name:        "host-root",
+				Destination: &m.options.hostRoot,
+				EnvVars:     []string{"HOST_ROOT"},
+				Value:       "/",
+				Usage:       "Path to the host's root filesystem. This is used when loading the vfio-pci module.",
+			},
 		},
 	}
 
 	return &c
 }
 
-func (m bindCommand) validateFlags(cfg *bindOptions) error {
-	if !cfg.all && cfg.deviceID == "" {
+func (m bindCommand) validateFlags() error {
+	if !m.options.all && m.options.deviceID == "" {
 		return fmt.Errorf("either --all or --device-id must be specified")
 	}
 
-	if cfg.all && cfg.deviceID != "" {
+	if m.options.all && m.options.deviceID != "" {
 		return fmt.Errorf("cannot specify both --all and --device-id")
 	}
 
 	return nil
 }
 
-func (m bindCommand) run(cfg *bindOptions) error {
-	if cfg.deviceID != "" {
-		return m.bindDevice(cfg.deviceID)
+func (m bindCommand) run() error {
+	m.nvpciLib = nvpci.New(
+		nvpci.WithLogger(m.logger),
+		nvpci.WithHostRoot(m.options.hostRoot),
+	)
+
+	if m.options.deviceID != "" {
+		return m.bindDevice()
 	}
 
 	return m.bindAll()
@@ -118,7 +126,8 @@ func (m bindCommand) bindAll() error {
 	return nil
 }
 
-func (m bindCommand) bindDevice(device string) error {
+func (m bindCommand) bindDevice() error {
+	device := m.options.deviceID
 	nvdev, err := m.nvpciLib.GetGPUByPciBusID(device)
 	if err != nil {
 		return fmt.Errorf("failed to get NVIDIA GPU device: %w", err)

cmd/vfio-manage/unbind.go

@@ -30,6 +30,7 @@ import (
 type unbindCommand struct {
 	logger   *logrus.Logger
 	nvpciLib nvpci.Interface
+	options  unbindOptions
 }
 
 type unbindOptions struct {
@@ -50,29 +51,26 @@ func newUnbindCommand(logger *logrus.Logger) *cli.Command {
 
 // build the unbind command
 func (m unbindCommand) build() *cli.Command {
-	cfg := unbindOptions{}
-
-	// Create the 'unbind' command
 	c := cli.Command{
 		Name:  "unbind",
 		Usage: "Unbind device(s) from their current driver",
 		Before: func(c *cli.Context) error {
-			return m.validateFlags(&cfg)
+			return m.validateFlags()
 		},
 		Action: func(c *cli.Context) error {
-			return m.run(&cfg)
+			return m.run()
 		},
 		Flags: []cli.Flag{
 			&cli.BoolFlag{
 				Name:        "all",
 				Aliases:     []string{"a"},
-				Destination: &cfg.all,
+				Destination: &m.options.all,
 				Usage:       "Bind all NVIDIA devices to vfio-pci",
 			},
 			&cli.StringFlag{
 				Name:        "device-id",
 				Aliases:     []string{"d"},
-				Destination: &cfg.deviceID,
+				Destination: &m.options.deviceID,
 				Usage:       "Specific device ID to bind (e.g., 0000:01:00.0)",
 			},
 		},
@@ -81,21 +79,21 @@ func (m unbindCommand) build() *cli.Command {
 	return &c
 }
 
-func (m unbindCommand) validateFlags(cfg *unbindOptions) error {
-	if !cfg.all && cfg.deviceID == "" {
+func (m unbindCommand) validateFlags() error {
+	if !m.options.all && m.options.deviceID == "" {
 		return fmt.Errorf("either --all or --device-id must be specified")
 	}
 
-	if cfg.all && cfg.deviceID != "" {
+	if m.options.all && m.options.deviceID != "" {
 		return fmt.Errorf("cannot specify both --all and --device-id")
 	}
 
 	return nil
 }
 
-func (m unbindCommand) run(cfg *unbindOptions) error {
-	if cfg.deviceID != "" {
-		return m.unbindDevice(cfg.deviceID)
+func (m unbindCommand) run() error {
+	if m.options.deviceID != "" {
+		return m.unbindDevice()
 	}
 
 	return m.unbindAll()
@@ -117,7 +115,8 @@ func (m unbindCommand) unbindAll() error {
 	return nil
 }
 
-func (m unbindCommand) unbindDevice(device string) error {
+func (m unbindCommand) unbindDevice() error {
+	device := m.options.deviceID
 	nvdev, err := m.nvpciLib.GetGPUByPciBusID(device)
 	if err != nil {
 		return fmt.Errorf("failed to get NVIDIA GPU device: %w", err)

internal/linuxutils/kmod.go

@@ -20,6 +20,7 @@ import (
 	"bufio"
 	"fmt"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -44,6 +45,9 @@ func NewKernelModules(log *logrus.Logger, options ...func(modules *KernelModules
 	for _, option := range options {
 		option(km)
 	}
+	if km.root == "" {
+		km.root = "/"
+	}
 	return km
 }
 
@@ -105,3 +109,8 @@ func (km *KernelModules) List(searchKey string) error {
 	}
 	return nil
 }
+
+func (km *KernelModules) Load(module string) error {
+	cmd := exec.Command("chroot", km.root, "modprobe", module)
+	return cmd.Run()
+}

internal/nvpci/nvpci.go

@@ -24,6 +24,8 @@ import (
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/sirupsen/logrus"
+
+	"github.com/NVIDIA/k8s-driver-manager/internal/linuxutils"
 )
 
 const (
@@ -43,7 +45,8 @@ type Interface interface {
 
 type nvpciWrapper struct {
 	nvpci.Interface
-	logger *logrus.Logger
+	logger   *logrus.Logger
+	hostRoot string
 }
 
 type nvidiaPCIDevice struct {
@@ -64,6 +67,9 @@ func New(opts ...Option) Interface {
 	if n.logger == nil {
 		n.logger = logrus.New()
 	}
+	if n.hostRoot == "" {
+		n.hostRoot = "/"
+	}
 
 	// (cdesiniotis) Create an identical logger for the underlying nvpci library,
 	// with the exception being the log level. Currently, the nvpci library
@@ -92,6 +98,13 @@ func WithLogger(logger *logrus.Logger) Option {
 	}
 }
 
+// WithHostRoot provides an Option to set the path to the host root filesystem
+func WithHostRoot(hostRoot string) Option {
+	return func(w *nvpciWrapper) {
+		w.hostRoot = hostRoot
+	}
+}
+
 // (cdesiniotis) ideally this method would be attached to the nvcpi.NvidiaPCIDevice struct
 // which removes the need for this wrapper
 func (w *nvpciWrapper) BindToVFIODriver(dev *nvpci.NvidiaPCIDevice) error {
@@ -112,6 +125,11 @@ func (w *nvpciWrapper) bindToVFIODriver(device *nvidiaPCIDevice) error {
 		return fmt.Errorf("failed to find best vfio variant driver: %w", err)
 	}
 
+	km := linuxutils.NewKernelModules(w.logger, linuxutils.WithRoot(w.hostRoot))
+	if err := km.Load(vfioDriverName); err != nil {
+		return fmt.Errorf("failed to load %q driver: %w", vfioDriverName, err)
+	}
+
 	// (cdesiniotis) Module names in the modules.alias file will only ever contain
 	// underscores characters and not dashes -- this aligns with how the linux kernel
 	// stores module names internally. This can sometimes differ from the name of the