feat: Fabric Manager partition-aware GetPreferredAllocation for Shared NVSwitch multi-tenancy

README.md

@@ -131,7 +131,32 @@ Push docker image to a docker repo
 ```shell
 make push-image DOCKER_REPO=<docker-repo-url> DOCKER_TAG=<image-tag>
 ```
+### Fabric Manager partition-aware allocation (Shared NVSwitch)
+
+On NVSwitch systems running Fabric Manager in the Shared NVSwitch model, NVLink between a VM's
+GPUs only works when those GPUs form an activated FM fabric partition. `GetPreferredAllocation`
+uses the FM SDK (`fmGetSupportedFabricPartitions`) to prefer device sets that match a supported
+partition of the requested size, so multi-GPU pods get working NVLink.
+
+It is opt-in and off by default. To enable:
+
+- Build with libnvfm support: `CGO_ENABLED=1 go build -tags=nvfm ./...` (needs the
+  `nvidia-fabric-manager-devel` package; the default build ships a stub and needs no SDK). A
+  binary built without the tag refuses to enable FM even if the env var is set.
+- Set `ENABLE_FABRIC_MANAGER=true`. Optional: `FM_ADDRESS` (default `127.0.0.1:6666`),
+  `FM_ADDRESS_TYPE` (`inet`|`unix`), `GPU_PCI_MODULE_MAPPING_PATH` (default
+  `/run/nvidia-fabricmanager/gpu-pci-module-mapping.json`).
+- The deployment must let the plugin reach FM's command API and read the mapping file — i.e.
+  the FM endpoint has to be reachable from the pod (host loopback needs `hostNetwork`) and the
+  mapping path host-mounted. The mapping is a JSON object of PCI BDF to GPU module id; it is
+  required because BDF order is not module-id order.
+
+If FM is unreachable or nothing matches, the plugin returns no preference and the devicemanager
+allocates as usual, so this can never fail an allocation. Automatic partition activation and
+generation of the mapping file are out of scope.
+
 ### To Do
 - Improve the healthcheck mechanism for GPUs with VFIO-PCI drivers
-- Support GetPreferredAllocation API of DevicePluginServer. It returns a preferred set of devices to allocate from a list of available ones. The resulting preferred allocation is not guaranteed to be the allocation ultimately performed by the devicemanager. It is only designed to help the devicemanager make a more informed allocation decision when possible. It has not been implemented in sandbox-device-plugin.
+- Build the container image with `-tags=nvfm` on all target arches (the FM SDK install is
+  currently x86_64 only) and add a libnvfm cgo build/test path to CI
 --------------------------------------------------------------

cmd/main.go

@@ -29,11 +29,47 @@
 package main
 
 import (
-	"github.com/nvidia/sandbox-device-plugin/pkg/device_plugin"
+	"log"
 	"os"
 	"strconv"
+
+	"github.com/nvidia/sandbox-device-plugin/pkg/device_plugin"
+	"github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager"
 )
 
+// isFabricManagerEnabled returns true if fabric manager integration is
+// enabled via the ENABLE_FABRIC_MANAGER environment variable.
+func isFabricManagerEnabled() bool {
+	enabled, _ := strconv.ParseBool(os.Getenv("ENABLE_FABRIC_MANAGER"))
+	return enabled
+}
+
+// setupFabricManager loads the GPU PCI-to-module mapping and constructs the
+// fabric manager client used for partition-aware preferred allocation. On any
+// failure it logs and returns without enabling fabric manager support, so the
+// device plugin still serves devices normally.
+func setupFabricManager() {
+	if !fabric_manager.Built {
+		log.Printf("ENABLE_FABRIC_MANAGER is set but this binary was built without libnvfm support; partition-aware allocation stays off (rebuild on linux with CGO_ENABLED=1 and -tags=nvfm)")
+		return
+	}
+
+	mappingPath := fabric_manager.PCIModuleMappingPathFromEnv()
+	pciToModule, err := fabric_manager.LoadPCIModuleMapping(mappingPath)
+	if err != nil {
+		log.Printf("Fabric manager enabled but loading PCI module mapping %q failed: %v; continuing without partition-aware allocation",
+			mappingPath, err)
+		return
+	}
+
+	device_plugin.FMPartition = &device_plugin.FMPartitionConfig{
+		Client:      fabric_manager.NewFMClient(fabric_manager.ConfigFromEnv()),
+		PCIToModule: pciToModule,
+	}
+	log.Printf("Fabric manager partition-aware allocation enabled (mapping: %s, %d GPU(s))",
+		mappingPath, len(pciToModule))
+}
+
 func main() {
 	var ok bool
 	var enableGFDFlag string
@@ -50,5 +86,9 @@ func main() {
 	if ok {
 		device_plugin.EnableGFD, _ = strconv.ParseBool(enableGFDFlag)
 	}
+	// default is false
+	if isFabricManagerEnabled() {
+		setupFabricManager()
+	}
 	device_plugin.InitiateDevicePlugin()
 }

pkg/device_plugin/device_plugin.go

@@ -37,6 +37,8 @@ import (
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
+
+	"github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager"
 )
 
 // NvidiaPCIDevice holds details about an NVIDIA PCI device (GPU or NVSwitch)
@@ -66,6 +68,19 @@ var stop = make(chan struct{})
 var PGPUAlias string
 var NVSwitchAlias string
 
+// FMPartitionConfig carries the fabric manager wiring created in cmd/main.go
+// when ENABLE_FABRIC_MANAGER=true.
+type FMPartitionConfig struct {
+	// Client talks to the fabric manager partition API.
+	Client fabric_manager.FMClient
+	// PCIToModule maps GPU PCI BDF addresses to physical module IDs.
+	PCIToModule map[string]uint32
+}
+
+// FMPartition, when non-nil, enables fabric manager partition-aware preferred
+// allocation for GPU (non-NVSwitch) device plugins.
+var FMPartition *FMPartitionConfig
+
 func InitiateDevicePlugin() {
 	// Initialize nvpci library if not already set (allows injection for testing)
 	if nvpciLib == nil {
@@ -124,6 +139,17 @@ func createDevicePlugins() {
 			devicePath = "/dev/vfio/devices/"
 		}
 		dp := NewGenericDevicePlugin(deviceName, devicePath, devs)
+
+		// Enable fabric manager partition-aware preferred allocation for GPU
+		// plugins only; NVSwitch plugins are left without a partition manager.
+		if FMPartition != nil && !isNVSwitchDeviceID(deviceID) {
+			deviceIDToPCI := buildDeviceIDToPCIMap(iommuKeys)
+			dp.SetPartitionManager(fabric_manager.NewPartitionManager(
+				FMPartition.Client, FMPartition.PCIToModule, deviceIDToPCI))
+			log.Printf("Fabric manager partition-aware allocation enabled for %q (%d device(s))",
+				deviceName, len(deviceIDToPCI))
+		}
+
 		err := startDevicePlugin(dp)
 		if err != nil {
 			log.Printf("Error starting %s device plugin: %v", dp.deviceName, err)
@@ -221,6 +247,25 @@ func createIommuDeviceMap() {
 	}
 }
 
+// buildDeviceIDToPCIMap maps plugin device IDs (IOMMU group/fd keys) to the
+// PCI BDF address of the device they represent, for fabric manager partition
+// lookups.
+func buildDeviceIDToPCIMap(iommuKeys []string) map[string]string {
+	deviceIDToPCI := make(map[string]string, len(iommuKeys))
+	for _, iommuKey := range iommuKeys {
+		devs := iommuMap[iommuKey]
+		if len(devs) == 0 {
+			continue
+		}
+		if len(devs) > 1 {
+			log.Printf("IOMMU key %s has %d devices; using %s for fabric partition mapping",
+				iommuKey, len(devs), devs[0].Address)
+		}
+		deviceIDToPCI[iommuKey] = devs[0].Address
+	}
+	return deviceIDToPCI
+}
+
 // getDeviceType returns a human-readable device type string
 func getDeviceType(dev *nvpci.NvidiaPCIDevice) string {
 	if dev.IsNVSwitch() {

pkg/device_plugin/generic_device_plugin.go

@@ -44,6 +44,8 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
+
+	"github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager"
 )
 
 var returnIommuMap = getIommuMap
@@ -60,6 +62,9 @@ type GenericDevicePlugin struct {
 	devicePath string
 	deviceName string
 	devsHealth []*pluginapi.Device
+	// partitionManager, when non-nil, enables fabric manager partition-aware
+	// preferred allocation for this plugin's devices.
+	partitionManager *fabric_manager.PartitionManager
 }
 
 // Returns an initialized instance of GenericDevicePlugin
@@ -78,6 +83,12 @@ func NewGenericDevicePlugin(deviceName string, devicePath string, devices []*plu
 	return dpi
 }
 
+// SetPartitionManager enables fabric manager partition-aware preferred
+// allocation for this device plugin.
+func (dpi *GenericDevicePlugin) SetPartitionManager(pm *fabric_manager.PartitionManager) {
+	dpi.partitionManager = pm
+}
+
 func waitForGrpcServer(socketPath string, timeout time.Duration) error {
 	conn, err := connect(socketPath, timeout)
 	if err != nil {
@@ -300,6 +311,9 @@ func (dpi *GenericDevicePlugin) cleanup() error {
 func (dpi *GenericDevicePlugin) GetDevicePluginOptions(ctx context.Context, e *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
 	options := &pluginapi.DevicePluginOptions{
 		PreStartRequired: false,
+		// Only advertise GetPreferredAllocation when fabric manager
+		// partition-aware allocation is enabled for this plugin.
+		GetPreferredAllocationAvailable: dpi.partitionManager != nil,
 	}
 	return options, nil
 }
@@ -309,15 +323,41 @@ func (dpi *GenericDevicePlugin) PreStartContainer(ctx context.Context, in *plugi
 	return res, nil
 }
 
-// GetPreferredAllocation is for compatible with new DevicePluginServer API for DevicePlugin service. It has not been implemented in kubevrit-gpu-device-plugin
+// GetPreferredAllocation returns a preferred set of devices to allocate from
+// a list of available ones. When fabric manager partition-aware allocation is
+// enabled, it prefers device sets that exactly match an FM fabric partition
+// of the requested size so that NVLink works between the allocated GPUs. The
+// resulting preferred allocation is not guaranteed to be the allocation
+// ultimately performed by the devicemanager; on any failure or when no
+// partition fits, an empty preference is returned so the kubelet falls back
+// to its default allocation instead of failing.
 func (dpi *GenericDevicePlugin) GetPreferredAllocation(ctx context.Context, in *pluginapi.PreferredAllocationRequest) (*pluginapi.PreferredAllocationResponse, error) {
-	// TODO
-	// returns a preferred set of devices to allocate
-	// from a list of available ones. The resulting preferred allocation is not
-	// guaranteed to be the allocation ultimately performed by the
-	// devicemanager. It is only designed to help the devicemanager make a more
-	// informed allocation decision when possible.
-	return nil, nil
+	log.Printf("[%s] GetPreferredAllocation called with %d container request(s)", dpi.deviceName, len(in.ContainerRequests))
+
+	response := &pluginapi.PreferredAllocationResponse{}
+	for idx, req := range in.ContainerRequests {
+		var preferred []string
+		if dpi.partitionManager != nil {
+			log.Printf("[%s] container request %d: available=%v mustInclude=%v size=%d",
+				dpi.deviceName, idx, req.AvailableDeviceIDs, req.MustIncludeDeviceIDs, req.AllocationSize)
+
+			var err error
+			preferred, err = dpi.partitionManager.SelectPreferred(ctx, req.AvailableDeviceIDs, req.MustIncludeDeviceIDs, int(req.AllocationSize))
+			if err != nil {
+				// Never fail the allocation on fabric manager errors; log and
+				// let the kubelet use its default allocation.
+				log.Printf("[%s] fabric partition preferred allocation failed for container request %d: %v",
+					dpi.deviceName, idx, err)
+				preferred = nil
+			}
+			log.Printf("[%s] preferred devices for container request %d: %v", dpi.deviceName, idx, preferred)
+		}
+		response.ContainerResponses = append(response.ContainerResponses,
+			&pluginapi.ContainerPreferredAllocationResponse{
+				DeviceIDs: preferred,
+			})
+	}
+	return response, nil
 }
 
 // Health check of GPU devices

pkg/device_plugin/generic_device_plugin_test.go

@@ -30,6 +30,7 @@ package device_plugin
 
 import (
 	"context"
+	"errors"
 	"os"
 	"path"
 	"path/filepath"
@@ -39,6 +40,8 @@ import (
 	. "github.com/onsi/gomega"
 	"google.golang.org/grpc"
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
+
+	"github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager"
 )
 
 var devices []*pluginapi.Device
@@ -237,3 +240,85 @@ var _ = Describe("Generic Device", func() {
 		Expect(devices[1].Health).To(Equal(pluginapi.Healthy))
 	})
 })
+
+// fakeFMClient is a fake fabric_manager.FMClient returning fixed partitions.
+type fakeFMClient struct {
+	partitions []fabric_manager.Partition
+	err        error
+}
+
+func (f *fakeFMClient) GetSupportedPartitions(ctx context.Context) ([]fabric_manager.Partition, error) {
+	if f.err != nil {
+		return nil, f.err
+	}
+	return f.partitions, nil
+}
+
+var _ = Describe("GetPreferredAllocation() Fabric Manager Tests", func() {
+	// Devices "1" and "2" (PCI addresses pciAddress1/pciAddress2) map to
+	// physical module IDs 1 and 2, which form fabric partition 1.
+	newFMPartitionManager := func(client fabric_manager.FMClient) *fabric_manager.PartitionManager {
+		return fabric_manager.NewPartitionManager(client,
+			map[string]uint32{pciAddress1: 1, pciAddress2: 2},
+			map[string]string{iommuGroup1: pciAddress1, iommuGroup2: pciAddress2},
+		)
+	}
+
+	newRequest := func() *pluginapi.PreferredAllocationRequest {
+		return &pluginapi.PreferredAllocationRequest{
+			ContainerRequests: []*pluginapi.ContainerPreferredAllocationRequest{
+				{
+					AvailableDeviceIDs: []string{iommuGroup1, iommuGroup2},
+					AllocationSize:     2,
+				},
+			},
+		}
+	}
+
+	It("does not advertise GetPreferredAllocation without a partition manager", func() {
+		dpi := NewGenericDevicePlugin("pgpu", "/tmp/pgpu", []*pluginapi.Device{})
+		options, err := dpi.GetDevicePluginOptions(context.Background(), &pluginapi.Empty{})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(options.GetPreferredAllocationAvailable).To(BeFalse())
+	})
+
+	It("advertises GetPreferredAllocation with a partition manager", func() {
+		dpi := NewGenericDevicePlugin("pgpu", "/tmp/pgpu", []*pluginapi.Device{})
+		dpi.SetPartitionManager(newFMPartitionManager(&fakeFMClient{}))
+		options, err := dpi.GetDevicePluginOptions(context.Background(), &pluginapi.Empty{})
+		Expect(err).ToNot(HaveOccurred())
+		Expect(options.GetPreferredAllocationAvailable).To(BeTrue())
+	})
+
+	It("returns an empty preference when no partition manager is set", func() {
+		dpi := NewGenericDevicePlugin("pgpu", "/tmp/pgpu", []*pluginapi.Device{})
+		resp, err := dpi.GetPreferredAllocation(context.Background(), newRequest())
+		Expect(err).ToNot(HaveOccurred())
+		Expect(resp.ContainerResponses).To(HaveLen(1))
+		Expect(resp.ContainerResponses[0].DeviceIDs).To(BeEmpty())
+	})
+
+	It("prefers the devices of a matching fabric partition", func() {
+		dpi := NewGenericDevicePlugin("pgpu", "/tmp/pgpu", []*pluginapi.Device{})
+		dpi.SetPartitionManager(newFMPartitionManager(&fakeFMClient{
+			partitions: []fabric_manager.Partition{
+				{ID: 1, GPUPhysicalIDs: []uint32{1, 2}},
+			},
+		}))
+		resp, err := dpi.GetPreferredAllocation(context.Background(), newRequest())
+		Expect(err).ToNot(HaveOccurred())
+		Expect(resp.ContainerResponses).To(HaveLen(1))
+		Expect(resp.ContainerResponses[0].DeviceIDs).To(Equal([]string{iommuGroup1, iommuGroup2}))
+	})
+
+	It("returns an empty preference instead of failing on fabric manager errors", func() {
+		dpi := NewGenericDevicePlugin("pgpu", "/tmp/pgpu", []*pluginapi.Device{})
+		dpi.SetPartitionManager(newFMPartitionManager(&fakeFMClient{
+			err: errors.New("connection refused"),
+		}))
+		resp, err := dpi.GetPreferredAllocation(context.Background(), newRequest())
+		Expect(err).ToNot(HaveOccurred())
+		Expect(resp.ContainerResponses).To(HaveLen(1))
+		Expect(resp.ContainerResponses[0].DeviceIDs).To(BeEmpty())
+	})
+})