Support for allocating GPUs in Passthrough-Mode

api/nvidia.com/resource/gpu/v1alpha1/driverconfig.go

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package v1alpha1
+
+import "fmt"
+
+// GpuDriver encodes the gpu driver as a string.
+type GpuDriver string
+
+const (
+	NvidiaDriver  GpuDriver = "nvidia"
+	VfioPciDriver GpuDriver = "vfio-pci"
+)
+
+// GpuDriverConfig holds the set of parameters for configuring a GPU with a driver.
+type GpuDriverConfig struct {
+	Driver GpuDriver `json:"driver"`
+}
+
+// DefaultGpuDriverConfig provides the default configuration of a GPU with a driver.
+func DefaultGpuDriverConfig() *GpuDriverConfig {
+	return &GpuDriverConfig{
+		Driver: NvidiaDriver,
+	}
+}
+
+// Normalize updates a GpuDriverConfig config with implied default values based on other settings.
+func (c *GpuDriverConfig) Normalize() error {
+	if c.Driver == "" {
+		c.Driver = NvidiaDriver
+	}
+	return nil
+}
+
+// Validate ensures that GpuDriverConfig has a valid set of values.
+func (c *GpuDriverConfig) Validate() error {
+	switch c.Driver {
+	case NvidiaDriver:
+		fallthrough
+	case VfioPciDriver:
+		break
+	default:
+		return fmt.Errorf("invalid driver '%s' specified in gpu driver configuration", c.Driver)
+	}
+	return nil
+}

api/nvidia.com/resource/gpu/v1alpha1/gpuconfig.go

@@ -29,7 +29,8 @@ import (
 // GpuConfig holds the set of parameters for configuring a GPU.
 type GpuConfig struct {
 	metav1.TypeMeta `json:",inline"`
-	Sharing         *GpuSharing `json:"sharing,omitempty"`
+	Sharing         *GpuSharing      `json:"sharing,omitempty"`
+	DriverConfig    *GpuDriverConfig `json:"driverConfig,omitempty"`
 }
 
 // DefaultGpuConfig provides the default GPU configuration.
@@ -45,11 +46,27 @@ func DefaultGpuConfig() *GpuConfig {
 				Interval: ptr.To(DefaultTimeSlice),
 			},
 		},
+		DriverConfig: &GpuDriverConfig{
+			Driver: NvidiaDriver,
+		},
 	}
 }
 
 // Normalize updates a GpuConfig config with implied default values based on other settings.
 func (c *GpuConfig) Normalize() error {
+	if c.DriverConfig == nil {
+		c.DriverConfig = DefaultGpuDriverConfig()
+	}
+
+	if err := c.DriverConfig.Normalize(); err != nil {
+		return err
+	}
+
+	// If sharing is not supported, don't proceed with normalizing its configuration.
+	if !c.SupportsSharing() {
+		return nil
+	}
+
 	if c.Sharing == nil {
 		c.Sharing = &GpuSharing{
 			Strategy: TimeSlicingStrategy,
@@ -68,8 +85,23 @@ func (c *GpuConfig) Normalize() error {
 
 // Validate ensures that GpuConfig has a valid set of values.
 func (c *GpuConfig) Validate() error {
-	if c.Sharing == nil {
-		return fmt.Errorf("no sharing strategy set")
+	if err := c.DriverConfig.Validate(); err != nil {
+		return err
 	}
-	return c.Sharing.Validate()
+
+	if c.SupportsSharing() {
+		if c.Sharing == nil {
+			return fmt.Errorf("no sharing strategy set")
+		}
+		if err := c.Sharing.Validate(); err != nil {
+			return err
+		}
+	} else if c.Sharing != nil {
+		return fmt.Errorf("sharing strategy cannot be provided while using non-nvidia driver")
+	}
+	return nil
+}
+
+func (c *GpuConfig) SupportsSharing() bool {
+	return c.DriverConfig.Driver == NvidiaDriver
 }

cmd/nvidia-dra-plugin/allocatable.go

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -106,3 +106,14 @@ func (d AllocatableDevices) UUIDs() []string {
 	slices.Sort(uuids)
 	return uuids
 }
+
+func (d AllocatableDevices) PciAddresses() []string {
+	var pciAddresses []string
+	for _, device := range d {
+		if device.Type() == GpuDeviceType {
+			pciAddresses = append(pciAddresses, device.Gpu.PciAddress)
+		}
+	}
+	slices.Sort(pciAddresses)
+	return pciAddresses
+}

cmd/nvidia-dra-plugin/device_state.go

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -38,17 +38,19 @@ type OpaqueDeviceConfig struct {
 }
 
 type DeviceConfigState struct {
-	MpsControlDaemonID string `json:"mpsControlDaemonID"`
+	MpsControlDaemonID string               `json:"mpsControlDaemonID"`
+	GpuConfig          *configapi.GpuConfig `json:"gpuConfig,omitempty"`
 	containerEdits     *cdiapi.ContainerEdits
 }
 
 type DeviceState struct {
 	sync.Mutex
-	cdi         *CDIHandler
-	tsManager   *TimeSlicingManager
-	mpsManager  *MpsManager
-	allocatable AllocatableDevices
-	config      *Config
+	cdi            *CDIHandler
+	tsManager      *TimeSlicingManager
+	mpsManager     *MpsManager
+	vfioPciManager *VfioPciManager
+	allocatable    AllocatableDevices
+	config         *Config
 
 	nvdevlib          *deviceLib
 	checkpointManager checkpointmanager.CheckpointManager
@@ -87,6 +89,8 @@ func NewDeviceState(ctx context.Context, config *Config) (*DeviceState, error) {
 	tsManager := NewTimeSlicingManager(nvdevlib)
 	mpsManager := NewMpsManager(config, nvdevlib, MpsRoot, hostDriverRoot, MpsControlDaemonTemplatePath)
 
+	vfioPciManager := NewVfioPciManager()
+
 	if err := cdi.CreateStandardDeviceSpecFile(allocatable); err != nil {
 		return nil, fmt.Errorf("unable to create base CDI spec file: %v", err)
 	}
@@ -100,12 +104,18 @@ func NewDeviceState(ctx context.Context, config *Config) (*DeviceState, error) {
 		cdi:               cdi,
 		tsManager:         tsManager,
 		mpsManager:        mpsManager,
+		vfioPciManager:    vfioPciManager,
 		allocatable:       allocatable,
 		config:            config,
 		nvdevlib:          nvdevlib,
 		checkpointManager: checkpointManager,
 	}
 
+	// Initialize the vfio-pci driver manager.
+	if err := vfioPciManager.Init(); err != nil {
+		return nil, fmt.Errorf("unable to initialize vfio-pci manager: %v", err)
+	}
+
 	checkpoints, err := state.checkpointManager.ListCheckpoints()
 	if err != nil {
 		return nil, fmt.Errorf("unable to list checkpoints: %v", err)
@@ -349,35 +359,67 @@ func (s *DeviceState) prepareDevices(ctx context.Context, claim *resourceapi.Res
 
 func (s *DeviceState) unprepareDevices(ctx context.Context, claimUID string, devices PreparedDevices) error {
 	for _, group := range devices {
+		if group.ConfigState.GpuConfig != nil {
+			err := s.unprepareGpus(ctx, group.ConfigState.GpuConfig, group.Devices.Gpus())
+			if err != nil {
+				return err
+			}
+		}
 		// Stop any MPS control daemons started for each group of prepared devices.
 		mpsControlDaemon := s.mpsManager.NewMpsControlDaemon(claimUID, group)
 		if err := mpsControlDaemon.Stop(ctx); err != nil {
 			return fmt.Errorf("error stopping MPS control daemon: %w", err)
 		}
-
 		// Go back to default time-slicing for all full GPUs.
 		tsc := configapi.DefaultGpuConfig().Sharing.TimeSlicingConfig
-		if err := s.tsManager.SetTimeSlice(group.Devices.Gpus(), tsc); err != nil {
+		if err := s.tsManager.SetTimeSlice(devices, tsc); err != nil {
 			return fmt.Errorf("error setting timeslice for devices: %w", err)
 		}
 	}
 	return nil
 }
+func (s *DeviceState) unprepareGpus(ctx context.Context, config *configapi.GpuConfig, devices PreparedDeviceList) error {
+	if config.DriverConfig.Driver == configapi.VfioPciDriver {
+		for _, device := range devices {
+			if err := s.vfioPciManager.Unconfigure(device.Gpu.Info); err != nil {
+				return fmt.Errorf("error unconfiguring vfio-pci device: %w", err)
+			}
+		}
+	}
+	return nil
+}
 
 func (s *DeviceState) applyConfig(ctx context.Context, config configapi.Interface, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult) (*DeviceConfigState, error) {
+	var configState DeviceConfigState
 	switch castConfig := config.(type) {
 	case *configapi.GpuConfig:
-		return s.applySharingConfig(ctx, castConfig.Sharing, claim, results)
+		configState.GpuConfig = castConfig
+		return s.applyGpuConfig(ctx, castConfig, claim, results, &configState)
 	case *configapi.MigDeviceConfig:
-		return s.applySharingConfig(ctx, castConfig.Sharing, claim, results)
+		return s.applySharingConfig(ctx, castConfig.Sharing, claim, results, &configState)
 	case *configapi.ImexChannelConfig:
-		return s.applyImexChannelConfig(ctx, castConfig, claim, results)
+		return s.applyImexChannelConfig(ctx, castConfig, claim, results, &configState)
 	default:
 		return nil, fmt.Errorf("unknown config type: %T", castConfig)
 	}
 }
 
-func (s *DeviceState) applySharingConfig(ctx context.Context, config configapi.Sharing, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult) (*DeviceConfigState, error) {
+func (s *DeviceState) applyGpuConfig(ctx context.Context, config *configapi.GpuConfig, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult, configState *DeviceConfigState) (*DeviceConfigState, error) {
+	var err error
+	configState, err = s.applyGpuDriverConfig(ctx, config.DriverConfig, results, configState)
+	if err != nil {
+		return nil, err
+	}
+	if config.SupportsSharing() {
+		configState, err = s.applySharingConfig(ctx, config.Sharing, claim, results, configState)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return configState, nil
+}
+
+func (s *DeviceState) applySharingConfig(ctx context.Context, config configapi.Sharing, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult, configState *DeviceConfigState) (*DeviceConfigState, error) {
 	// Get the list of claim requests this config is being applied over.
 	var requests []string
 	for _, r := range results {
@@ -390,9 +432,6 @@ func (s *DeviceState) applySharingConfig(ctx context.Context, config configapi.S
 		allocatableDevices[r.Device] = s.allocatable[r.Device]
 	}
 
-	// Declare a device group state object to populate.
-	var configState DeviceConfigState
-
 	// Apply time-slicing settings (if available).
 	if config.IsTimeSlicing() {
 		tsc, err := config.GetTimeSlicingConfig()
@@ -424,13 +463,10 @@ func (s *DeviceState) applySharingConfig(ctx context.Context, config configapi.S
 		configState.containerEdits = mpsControlDaemon.GetCDIContainerEdits()
 	}
 
-	return &configState, nil
+	return configState, nil
 }
 
-func (s *DeviceState) applyImexChannelConfig(ctx context.Context, config *configapi.ImexChannelConfig, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult) (*DeviceConfigState, error) {
-	// Declare a device group state object to populate.
-	var configState DeviceConfigState
-
+func (s *DeviceState) applyImexChannelConfig(ctx context.Context, config *configapi.ImexChannelConfig, claim *resourceapi.ResourceClaim, results []*resourceapi.DeviceRequestAllocationResult, configState *DeviceConfigState) (*DeviceConfigState, error) {
 	// Create any necessary IMEX channels and gather their CDI container edits.
 	for _, r := range results {
 		imexChannel := s.allocatable[r.Device].ImexChannel
@@ -440,7 +476,25 @@ func (s *DeviceState) applyImexChannelConfig(ctx context.Context, config *config
 		configState.containerEdits = configState.containerEdits.Append(s.cdi.GetImexChannelContainerEdits(imexChannel))
 	}
 
-	return &configState, nil
+	return configState, nil
+}
+
+func (s *DeviceState) applyGpuDriverConfig(ctx context.Context, config *configapi.GpuDriverConfig, results []*resourceapi.DeviceRequestAllocationResult, configState *DeviceConfigState) (*DeviceConfigState, error) {
+	if config.Driver != configapi.VfioPciDriver {
+		return configState, nil
+	}
+
+	// Apply vfio-pci driver settings.
+	for _, r := range results {
+		info := s.allocatable[r.Device]
+		err := s.vfioPciManager.Configure(info.Gpu)
+		if err != nil {
+			return nil, err
+		}
+		configState.containerEdits = configState.containerEdits.Append(s.vfioPciManager.GetCDIContainerEdits(info.Gpu))
+	}
+
+	return configState, nil
 }
 
 // GetOpaqueDeviceConfigs returns an ordered list of the configs contained in possibleConfigs for this driver.

cmd/nvidia-dra-plugin/deviceinfo.go

@@ -29,6 +29,7 @@ import (
 
 type GpuInfo struct {
 	UUID                  string `json:"uuid"`
+	PciAddress            string `json:"pciAddress"`
 	index                 int
 	minor                 int
 	migEnabled            bool
@@ -130,6 +131,9 @@ func (d *GpuInfo) GetDevice() resourceapi.Device {
 				"cudaDriverVersion": {
 					VersionValue: ptr.To(semver.MustParse(d.cudaDriverVersion).String()),
 				},
+				"pciAddress": {
+					StringValue: &d.PciAddress,
+				},
 			},
 			Capacity: map[resourceapi.QualifiedName]resource.Quantity{
 				"memory": *resource.NewQuantity(int64(d.memoryBytes), resource.BinarySI),