From 1edbd6743f2f99aab83dd3c30f32a6bd9136357b Mon Sep 17 00:00:00 2001 From: inerplat Date: Thu, 2 Jul 2026 10:49:20 +0900 Subject: [PATCH] feat: Fabric Manager partition-aware GetPreferredAllocation for Shared NVSwitch multi-tenancy Signed-off-by: inerplat --- README.md | 27 ++- cmd/main.go | 42 +++- pkg/device_plugin/device_plugin.go | 45 ++++ pkg/device_plugin/generic_device_plugin.go | 56 ++++- .../generic_device_plugin_test.go | 85 +++++++ pkg/fabric_manager/client.go | 112 ++++++++++ .../fabric_manager_suite_test.go | 41 ++++ pkg/fabric_manager/mapping.go | 121 ++++++++++ pkg/fabric_manager/mapping_test.go | 115 ++++++++++ pkg/fabric_manager/nvfm_client_linux.go | 174 +++++++++++++++ pkg/fabric_manager/nvfm_client_stub.go | 58 +++++ pkg/fabric_manager/partition_manager.go | 206 +++++++++++++++++ pkg/fabric_manager/partition_manager_test.go | 211 ++++++++++++++++++ 13 files changed, 1283 insertions(+), 10 deletions(-) create mode 100644 pkg/fabric_manager/client.go create mode 100644 pkg/fabric_manager/fabric_manager_suite_test.go create mode 100644 pkg/fabric_manager/mapping.go create mode 100644 pkg/fabric_manager/mapping_test.go create mode 100644 pkg/fabric_manager/nvfm_client_linux.go create mode 100644 pkg/fabric_manager/nvfm_client_stub.go create mode 100644 pkg/fabric_manager/partition_manager.go create mode 100644 pkg/fabric_manager/partition_manager_test.go diff --git a/README.md b/README.md index dd759bdf..a5cb4168 100644 --- a/README.md +++ b/README.md @@ -131,7 +131,32 @@ Push docker image to a docker repo ```shell make push-image DOCKER_REPO= DOCKER_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 -------------------------------------------------------------- diff --git a/cmd/main.go b/cmd/main.go index 35226f56..b74e9b09 100644 --- a/cmd/main.go +++ b/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() } diff --git a/pkg/device_plugin/device_plugin.go b/pkg/device_plugin/device_plugin.go index 98a84785..8d3d846e 100644 --- a/pkg/device_plugin/device_plugin.go +++ b/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() { diff --git a/pkg/device_plugin/generic_device_plugin.go b/pkg/device_plugin/generic_device_plugin.go index 60ecc489..ee313dd4 100644 --- a/pkg/device_plugin/generic_device_plugin.go +++ b/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 diff --git a/pkg/device_plugin/generic_device_plugin_test.go b/pkg/device_plugin/generic_device_plugin_test.go index fb415e20..d989066b 100644 --- a/pkg/device_plugin/generic_device_plugin_test.go +++ b/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()) + }) +}) diff --git a/pkg/fabric_manager/client.go b/pkg/fabric_manager/client.go new file mode 100644 index 00000000..0971e13c --- /dev/null +++ b/pkg/fabric_manager/client.go @@ -0,0 +1,112 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +// Package fabric_manager provides a client for the NVIDIA Fabric Manager (FM) +// Shared NVSwitch partition API and partition-aware preferred-allocation logic +// for the sandbox device plugin. On HGX systems, NVLink between a tenant VM's +// GPUs only works when the assigned GPUs exactly match an FM fabric partition, +// so the device plugin uses this package to steer kubelet allocations towards +// complete partitions. +// +// The design mirrors kubevirt-gpu-device-plugin PR#166 (mresvanis), adapted +// for sandbox-device-plugin. Addresses NVIDIA/sandbox-device-plugin#24. +package fabric_manager + +import ( + "context" + "log" + "os" + "strings" +) + +const ( + // DefaultFMAddress is the default fabric manager command API address. + DefaultFMAddress = "127.0.0.1:6666" + + // DefaultFMTimeoutMs is the default connection timeout in milliseconds. + // fmConnect requires a timeout strictly greater than zero. + DefaultFMTimeoutMs = 5000 + + // AddressTypeInet selects a TCP/IP connection to fabric manager. + AddressTypeInet = "inet" + // AddressTypeUnix selects a Unix domain socket connection to fabric manager. + AddressTypeUnix = "unix" + + fmAddressEnvVar = "FM_ADDRESS" + fmAddressTypeEnvVar = "FM_ADDRESS_TYPE" +) + +// Partition describes a single fabric partition supported by fabric manager. +type Partition struct { + // ID is the fabric partition ID. + ID uint32 + // IsActive reports whether the partition is currently activated. + IsActive bool + // GPUPhysicalIDs lists the physical module IDs of the GPUs in the partition. + GPUPhysicalIDs []uint32 +} + +// FMClient is the interface to the fabric manager partition API. +type FMClient interface { + // GetSupportedPartitions returns all fabric partitions supported by the + // fabric manager instance. + GetSupportedPartitions(ctx context.Context) ([]Partition, error) +} + +// Config contains connection options for the fabric manager client. +type Config struct { + // Address is the fabric manager command API address, e.g. "127.0.0.1:6666" + // for inet or a socket path for unix. + Address string + // AddressType is either AddressTypeInet or AddressTypeUnix. + AddressType string + // TimeoutMs is the connection timeout in milliseconds; must be > 0. + TimeoutMs uint32 +} + +// ConfigFromEnv builds a Config from the FM_ADDRESS and FM_ADDRESS_TYPE +// environment variables, falling back to defaults when unset or invalid. +func ConfigFromEnv() Config { + cfg := Config{ + Address: DefaultFMAddress, + AddressType: AddressTypeInet, + TimeoutMs: DefaultFMTimeoutMs, + } + if v := os.Getenv(fmAddressEnvVar); v != "" { + cfg.Address = v + } + if v := os.Getenv(fmAddressTypeEnvVar); v != "" { + switch strings.ToLower(v) { + case AddressTypeInet, AddressTypeUnix: + cfg.AddressType = strings.ToLower(v) + default: + log.Printf("fabric_manager: unknown %s value %q, using %q", fmAddressTypeEnvVar, v, cfg.AddressType) + } + } + return cfg +} diff --git a/pkg/fabric_manager/fabric_manager_suite_test.go b/pkg/fabric_manager/fabric_manager_suite_test.go new file mode 100644 index 00000000..7313238d --- /dev/null +++ b/pkg/fabric_manager/fabric_manager_suite_test.go @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestFabricManager(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "FabricManager Suite") +} diff --git a/pkg/fabric_manager/mapping.go b/pkg/fabric_manager/mapping.go new file mode 100644 index 00000000..cd61be2b --- /dev/null +++ b/pkg/fabric_manager/mapping.go @@ -0,0 +1,121 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager + +import ( + "bytes" + "encoding/json" + "fmt" + "os" + "strconv" + "strings" +) + +const ( + // DefaultPCIModuleMappingPath is the default location of the GPU + // PCI-to-module mapping JSON file produced on the host. + DefaultPCIModuleMappingPath = "/run/nvidia-fabricmanager/gpu-pci-module-mapping.json" + + mappingPathEnvVar = "GPU_PCI_MODULE_MAPPING_PATH" +) + +// PCIModuleMappingPathFromEnv returns the mapping file path from the +// GPU_PCI_MODULE_MAPPING_PATH environment variable, or the default path. +func PCIModuleMappingPathFromEnv() string { + if v := os.Getenv(mappingPathEnvVar); v != "" { + return v + } + return DefaultPCIModuleMappingPath +} + +// NormalizePCIAddress normalizes a PCI BDF address to the sysfs format used +// by the Linux kernel (4-digit lowercase hex domain). Mapping files produced +// by NVIDIA tooling may use an 8-digit domain and uppercase hex (e.g. +// "00000000:AB:00.0"), while sysfs uses "0000:ab:00.0". +func NormalizePCIAddress(addr string) string { + lower := strings.ToLower(strings.TrimSpace(addr)) + + colonIdx := strings.Index(lower, ":") + if colonIdx < 0 { + return lower + } + + domain := lower[:colonIdx] + rest := lower[colonIdx:] + + domainVal, err := strconv.ParseUint(domain, 16, 32) + if err != nil { + return lower + } + + if domainVal <= 0xFFFF { + return fmt.Sprintf("%04x%s", domainVal, rest) + } + return fmt.Sprintf("%08x%s", domainVal, rest) +} + +// LoadPCIModuleMapping reads a JSON file mapping GPU PCI BDF addresses to +// physical module IDs and returns the mapping with normalized PCI addresses. +// Both value forms are accepted for compatibility: +// +// {"0000:19:00.0": 1} (numeric module IDs) +// {"0000:19:00.0": "1"} (string module IDs, as used by kubevirt-gpu-device-plugin PR#166) +func LoadPCIModuleMapping(path string) (map[string]uint32, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read PCI module mapping file: %w", err) + } + + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + var raw map[string]interface{} + if err := decoder.Decode(&raw); err != nil { + return nil, fmt.Errorf("failed to parse PCI module mapping JSON: %w", err) + } + + pciToModule := make(map[string]uint32, len(raw)) + for pciAddr, value := range raw { + var moduleIDStr string + switch v := value.(type) { + case string: + moduleIDStr = v + case json.Number: + moduleIDStr = v.String() + default: + return nil, fmt.Errorf("invalid module ID %v for PCI address %s: expected number or string", value, pciAddr) + } + moduleID, err := strconv.ParseUint(moduleIDStr, 10, 32) + if err != nil { + return nil, fmt.Errorf("invalid module ID %q for PCI address %s: %w", moduleIDStr, pciAddr, err) + } + pciToModule[NormalizePCIAddress(pciAddr)] = uint32(moduleID) + } + + return pciToModule, nil +} diff --git a/pkg/fabric_manager/mapping_test.go b/pkg/fabric_manager/mapping_test.go new file mode 100644 index 00000000..48576dc8 --- /dev/null +++ b/pkg/fabric_manager/mapping_test.go @@ -0,0 +1,115 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager_test + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + fm "github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager" +) + +var _ = Describe("PCI module mapping", func() { + writeMapping := func(content string) string { + path := filepath.Join(GinkgoT().TempDir(), "mapping.json") + Expect(os.WriteFile(path, []byte(content), 0o644)).To(Succeed()) + return path + } + + Context("LoadPCIModuleMapping()", func() { + It("parses numeric module IDs", func() { + path := writeMapping(`{"0000:19:00.0": 1, "0000:3b:00.0": 2}`) + mapping, err := fm.LoadPCIModuleMapping(path) + Expect(err).ToNot(HaveOccurred()) + Expect(mapping).To(Equal(map[string]uint32{ + "0000:19:00.0": 1, + "0000:3b:00.0": 2, + })) + }) + + It("parses string module IDs (kubevirt-gpu-device-plugin PR#166 format)", func() { + path := writeMapping(`{"0000:19:00.0": "1", "0000:3b:00.0": "2"}`) + mapping, err := fm.LoadPCIModuleMapping(path) + Expect(err).ToNot(HaveOccurred()) + Expect(mapping).To(Equal(map[string]uint32{ + "0000:19:00.0": 1, + "0000:3b:00.0": 2, + })) + }) + + It("normalizes 8-digit-domain uppercase BDFs to sysfs format", func() { + path := writeMapping(`{"00000000:AB:00.0": "8", "00000000:2A:00.0": 4}`) + mapping, err := fm.LoadPCIModuleMapping(path) + Expect(err).ToNot(HaveOccurred()) + Expect(mapping).To(Equal(map[string]uint32{ + "0000:ab:00.0": 8, + "0000:2a:00.0": 4, + })) + }) + + It("returns an error for a missing file", func() { + _, err := fm.LoadPCIModuleMapping(filepath.Join(GinkgoT().TempDir(), "missing.json")) + Expect(err).To(MatchError(ContainSubstring("failed to read PCI module mapping file"))) + }) + + It("returns an error for invalid JSON", func() { + path := writeMapping(`not json`) + _, err := fm.LoadPCIModuleMapping(path) + Expect(err).To(MatchError(ContainSubstring("failed to parse PCI module mapping JSON"))) + }) + + DescribeTable("returns an error for invalid module ID values", + func(content string) { + path := writeMapping(content) + _, err := fm.LoadPCIModuleMapping(path) + Expect(err).To(MatchError(ContainSubstring("invalid module ID"))) + }, + Entry("non-numeric string", `{"0000:19:00.0": "not-a-number"}`), + Entry("negative number", `{"0000:19:00.0": -1}`), + Entry("fractional number", `{"0000:19:00.0": 1.5}`), + Entry("boolean value", `{"0000:19:00.0": true}`), + ) + }) + + Context("NormalizePCIAddress()", func() { + DescribeTable("normalizes to 4-digit lowercase hex domain", + func(input, expected string) { + Expect(fm.NormalizePCIAddress(input)).To(Equal(expected)) + }, + Entry("already normalized", "0000:19:00.0", "0000:19:00.0"), + Entry("uppercase hex", "0000:AB:00.0", "0000:ab:00.0"), + Entry("8-digit domain", "00000000:ab:00.0", "0000:ab:00.0"), + Entry("no domain separator", "invalid", "invalid"), + Entry("non-hex domain", "zzzz:ab:00.0", "zzzz:ab:00.0"), + ) + }) +}) diff --git a/pkg/fabric_manager/nvfm_client_linux.go b/pkg/fabric_manager/nvfm_client_linux.go new file mode 100644 index 00000000..4f3a564c --- /dev/null +++ b/pkg/fabric_manager/nvfm_client_linux.go @@ -0,0 +1,174 @@ +//go:build nvfm && cgo + +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +// CGO bindings to the NVIDIA Fabric Manager SDK (libnvfm). The SDK is +// provided by the nvidia-fabric-manager-devel package, which installs +// /usr/include/nv_fm_agent.h and libnvfm.so. Only compiled with -tags=nvfm +// so that the default build works on machines without the FM SDK. +// +// https://docs.nvidia.com/datacenter/tesla/fabric-manager-user-guide/index.html#fabric-manager-sdk + +package fabric_manager + +/* +#cgo CFLAGS: -I/usr/include +#cgo LDFLAGS: -lnvfm + +#include +#include "nv_fm_agent.h" +*/ +import "C" + +import ( + "context" + "fmt" + "log" + "sync" + "unsafe" +) + +// Built reports whether this binary was compiled with libnvfm support. +const Built = true + +// nvfmClient talks to a running nv-fabricmanager over the FM SDK. The +// connection is established lazily on first use and cached; on any API +// failure the connection is dropped so the next call reconnects. +type nvfmClient struct { + cfg Config + + mu sync.Mutex + handle C.fmHandle_t + connected bool + libInited bool +} + +// NewFMClient returns an FMClient backed by libnvfm. +func NewFMClient(cfg Config) FMClient { + if cfg.Address == "" { + cfg.Address = DefaultFMAddress + } + if cfg.TimeoutMs == 0 { + // fmConnect fails when timeoutMs is 0. + cfg.TimeoutMs = DefaultFMTimeoutMs + } + return &nvfmClient{cfg: cfg} +} + +// connectLocked initializes the FM library and connects to fabric manager. +// Callers must hold c.mu. +func (c *nvfmClient) connectLocked() error { + if c.connected { + return nil + } + + if !c.libInited { + if ret := C.fmLibInit(); ret != C.FM_ST_SUCCESS { + return fmt.Errorf("fmLibInit failed: %d", int32(ret)) + } + c.libInited = true + } + + if len(c.cfg.Address) >= int(C.FM_MAX_STR_LENGTH) { + return fmt.Errorf("fabric manager address %q too long", c.cfg.Address) + } + + var params C.fmConnectParams_t + params.version = C.fmConnectParams_version + params.timeoutMs = C.uint(c.cfg.TimeoutMs) + if c.cfg.AddressType == AddressTypeUnix { + params.addressType = uint32(C.NV_FM_API_ADDR_TYPE_UNIX) + } else { + params.addressType = uint32(C.NV_FM_API_ADDR_TYPE_INET) + } + for i := 0; i < len(c.cfg.Address); i++ { + params.addressInfo[i] = C.char(c.cfg.Address[i]) + } + params.addressInfo[len(c.cfg.Address)] = 0 + + if ret := C.fmConnect(¶ms, &c.handle); ret != C.FM_ST_SUCCESS { + return fmt.Errorf("fmConnect to %s (%s) failed: %d", c.cfg.Address, c.cfg.AddressType, int32(ret)) + } + c.connected = true + log.Printf("fabric_manager: connected to fabric manager at %s (%s)", c.cfg.Address, c.cfg.AddressType) + return nil +} + +// disconnectLocked drops the cached connection. Callers must hold c.mu. +func (c *nvfmClient) disconnectLocked() { + if !c.connected { + return + } + if ret := C.fmDisconnect(c.handle); ret != C.FM_ST_SUCCESS { + log.Printf("fabric_manager: fmDisconnect failed: %d", int32(ret)) + } + c.connected = false +} + +// GetSupportedPartitions returns all fabric partitions supported by the +// fabric manager instance. The FM SDK is synchronous, so ctx is not used to +// cancel an in-flight call; the connection timeout bounds the call instead. +func (c *nvfmClient) GetSupportedPartitions(ctx context.Context) ([]Partition, error) { + c.mu.Lock() + defer c.mu.Unlock() + + if err := c.connectLocked(); err != nil { + return nil, err + } + + // fmFabricPartitionList_t is large (max partitions x max GPUs); allocate + // it on the C heap rather than the Go stack. + list := (*C.fmFabricPartitionList_t)(C.calloc(1, C.sizeof_fmFabricPartitionList_t)) + if list == nil { + return nil, fmt.Errorf("failed to allocate fabric partition list") + } + defer C.free(unsafe.Pointer(list)) + list.version = C.fmFabricPartitionList_version + + if ret := C.fmGetSupportedFabricPartitions(c.handle, list); ret != C.FM_ST_SUCCESS { + // Drop the connection so the next call reconnects. + c.disconnectLocked() + return nil, fmt.Errorf("fmGetSupportedFabricPartitions failed: %d", int32(ret)) + } + + partitions := make([]Partition, 0, int(list.numPartitions)) + for i := 0; i < int(list.numPartitions); i++ { + info := &list.partitionInfo[i] + p := Partition{ + ID: uint32(info.partitionId), + IsActive: info.isActive != 0, + GPUPhysicalIDs: make([]uint32, 0, int(info.numGpus)), + } + for g := 0; g < int(info.numGpus); g++ { + p.GPUPhysicalIDs = append(p.GPUPhysicalIDs, uint32(info.gpuInfo[g].physicalId)) + } + partitions = append(partitions, p) + } + return partitions, nil +} diff --git a/pkg/fabric_manager/nvfm_client_stub.go b/pkg/fabric_manager/nvfm_client_stub.go new file mode 100644 index 00000000..c536afed --- /dev/null +++ b/pkg/fabric_manager/nvfm_client_stub.go @@ -0,0 +1,58 @@ +//go:build !nvfm || !linux || !cgo + +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager + +import ( + "context" + "errors" +) + +// Built reports whether this binary was compiled with libnvfm support. It is +// false in the stub build so callers can refuse to enable fabric manager +// instead of silently advertising a capability that always fails. +const Built = false + +// ErrNotBuilt is returned by the stub client used when the binary was built +// without the nvfm build tag (i.e. without the Fabric Manager SDK). +var ErrNotBuilt = errors.New("fabric_manager: built without nvfm support (rebuild on linux with CGO_ENABLED=1 and -tags=nvfm)") + +type stubClient struct{} + +// NewFMClient returns a stub FMClient that fails all calls with ErrNotBuilt. +// The real libnvfm-backed client is only available when built with +// -tags=nvfm on linux. +func NewFMClient(cfg Config) FMClient { + return stubClient{} +} + +func (stubClient) GetSupportedPartitions(ctx context.Context) ([]Partition, error) { + return nil, ErrNotBuilt +} diff --git a/pkg/fabric_manager/partition_manager.go b/pkg/fabric_manager/partition_manager.go new file mode 100644 index 00000000..44997378 --- /dev/null +++ b/pkg/fabric_manager/partition_manager.go @@ -0,0 +1,206 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager + +import ( + "context" + "fmt" + "log" + "sort" +) + +// PartitionManager selects preferred device sets that exactly match fabric +// manager partitions. It translates plugin device IDs (IOMMU group/fd keys) +// to GPU physical module IDs via PCI BDF addresses, since FM partitions +// identify GPUs by physical module ID. +type PartitionManager struct { + fm FMClient + pciToModule map[string]uint32 + deviceIDToPCI map[string]string +} + +// NewPartitionManager creates a PartitionManager. pciToModule maps normalized +// PCI BDF addresses to GPU physical module IDs (see LoadPCIModuleMapping) and +// deviceIDToPCI maps plugin device IDs to the PCI BDF address of the GPU they +// represent. +func NewPartitionManager(fm FMClient, pciToModule map[string]uint32, deviceIDToPCI map[string]string) *PartitionManager { + normalizedPCIToModule := make(map[string]uint32, len(pciToModule)) + for addr, moduleID := range pciToModule { + normalizedPCIToModule[NormalizePCIAddress(addr)] = moduleID + } + normalizedDeviceIDToPCI := make(map[string]string, len(deviceIDToPCI)) + for id, addr := range deviceIDToPCI { + normalizedDeviceIDToPCI[id] = NormalizePCIAddress(addr) + } + return &PartitionManager{ + fm: fm, + pciToModule: normalizedPCIToModule, + deviceIDToPCI: normalizedDeviceIDToPCI, + } +} + +// physicalID translates a plugin device ID to a GPU physical module ID. +func (pm *PartitionManager) physicalID(deviceID string) (uint32, bool) { + pciAddr, ok := pm.deviceIDToPCI[deviceID] + if !ok { + return 0, false + } + moduleID, ok := pm.pciToModule[pciAddr] + return moduleID, ok +} + +// SelectPreferred picks the devices that form a fabric partition of exactly +// the requested size, from the available device IDs, including all +// mustInclude device IDs. Partitions that are already active are preferred +// over inactive ones; ties are broken deterministically by lowest partition +// ID. It returns nil (with no error) when no partition fits, in which case +// the kubelet falls back to its default allocation. An error is only +// returned for fabric manager communication failures; callers should log it +// and treat it as "no preference" rather than failing the allocation. +func (pm *PartitionManager) SelectPreferred(ctx context.Context, availableDeviceIDs []string, mustInclude []string, size int) ([]string, error) { + if size <= 0 || len(availableDeviceIDs) == 0 { + return nil, nil + } + if len(mustInclude) > size { + log.Printf("fabric_manager: %d must-include devices exceed allocation size %d, no partition preference", len(mustInclude), size) + return nil, nil + } + + // Translate available device IDs to physical module IDs. + availablePhys := make(map[uint32]struct{}, len(availableDeviceIDs)) + physToDevice := make(map[uint32]string, len(availableDeviceIDs)) + for _, deviceID := range availableDeviceIDs { + phys, ok := pm.physicalID(deviceID) + if !ok { + log.Printf("fabric_manager: skipping device %q: no PCI/module mapping", deviceID) + continue + } + availablePhys[phys] = struct{}{} + physToDevice[phys] = deviceID + } + if len(availablePhys) < size { + log.Printf("fabric_manager: only %d/%d available devices mapped to physical IDs (need %d), no partition preference", + len(availablePhys), len(availableDeviceIDs), size) + return nil, nil + } + + // Translate must-include device IDs; if any cannot be mapped we cannot + // guarantee a partition contains it, so express no preference. + mustPhys := make(map[uint32]struct{}, len(mustInclude)) + for _, deviceID := range mustInclude { + phys, ok := pm.physicalID(deviceID) + if !ok { + log.Printf("fabric_manager: must-include device %q has no PCI/module mapping, no partition preference", deviceID) + return nil, nil + } + mustPhys[phys] = struct{}{} + } + + partitions, err := pm.fm.GetSupportedPartitions(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get fabric partitions: %w", err) + } + + // Collect partitions of exactly the requested size whose GPUs are all + // available and that contain all must-include devices. + var candidates []Partition + for _, partition := range partitions { + if len(partition.GPUPhysicalIDs) != size { + continue + } + if !allInSet(partition.GPUPhysicalIDs, availablePhys) { + continue + } + if !setInSlice(mustPhys, partition.GPUPhysicalIDs) { + continue + } + candidates = append(candidates, partition) + } + if len(candidates) == 0 { + log.Printf("fabric_manager: no fabric partition of size %d fits available devices %v (mustInclude %v), no partition preference", + size, availableDeviceIDs, mustInclude) + return nil, nil + } + + // Prefer partitions that are already active, then lowest partition ID. + sort.Slice(candidates, func(i, j int) bool { + if candidates[i].IsActive != candidates[j].IsActive { + return candidates[i].IsActive + } + return candidates[i].ID < candidates[j].ID + }) + best := candidates[0] + log.Printf("fabric_manager: selected partition %d (active: %t, GPUs: %v) from %d candidate(s)", + best.ID, best.IsActive, best.GPUPhysicalIDs, len(candidates)) + + // Build the preferred device list: must-include devices first (request + // order), then the remaining partition members in partition order. + preferred := make([]string, 0, size) + added := make(map[string]struct{}, size) + for _, deviceID := range mustInclude { + if _, exists := added[deviceID]; exists { + continue + } + added[deviceID] = struct{}{} + preferred = append(preferred, deviceID) + } + for _, phys := range best.GPUPhysicalIDs { + deviceID := physToDevice[phys] + if _, exists := added[deviceID]; exists { + continue + } + added[deviceID] = struct{}{} + preferred = append(preferred, deviceID) + } + return preferred, nil +} + +// allInSet reports whether every ID in ids is present in set. +func allInSet(ids []uint32, set map[uint32]struct{}) bool { + for _, id := range ids { + if _, ok := set[id]; !ok { + return false + } + } + return true +} + +// setInSlice reports whether every ID in set is present in ids. +func setInSlice(set map[uint32]struct{}, ids []uint32) bool { + present := make(map[uint32]struct{}, len(ids)) + for _, id := range ids { + present[id] = struct{}{} + } + for id := range set { + if _, ok := present[id]; !ok { + return false + } + } + return true +} diff --git a/pkg/fabric_manager/partition_manager_test.go b/pkg/fabric_manager/partition_manager_test.go new file mode 100644 index 00000000..71d3c342 --- /dev/null +++ b/pkg/fabric_manager/partition_manager_test.go @@ -0,0 +1,211 @@ +/* + * Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions + * are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of NVIDIA CORPORATION nor the names of its + * contributors may be used to endorse or promote products derived + * from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY + * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR + * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR + * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, + * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR + * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY + * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +package fabric_manager_test + +import ( + "context" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + fm "github.com/nvidia/sandbox-device-plugin/pkg/fabric_manager" +) + +// fakeFMClient is a fake FMClient returning a fixed partition list. +type fakeFMClient struct { + partitions []fm.Partition + err error +} + +func (f *fakeFMClient) GetSupportedPartitions(ctx context.Context) ([]fm.Partition, error) { + if f.err != nil { + return nil, f.err + } + return f.partitions, nil +} + +var _ = Describe("PartitionManager", func() { + // H100 HGX topology: 8 GPUs with physical module IDs 1..8. Plugin device + // IDs are iommufd numbers "0".."7"; deviceID "N" maps to physical ID N+1. + deviceIDToPCI := map[string]string{ + "0": "0000:19:00.0", + "1": "0000:3b:00.0", + "2": "0000:4c:00.0", + "3": "0000:5d:00.0", + "4": "0000:9b:00.0", + "5": "0000:bb:00.0", + "6": "0000:cb:00.0", + "7": "0000:db:00.0", + } + pciToModule := map[string]uint32{ + "0000:19:00.0": 1, + "0000:3b:00.0": 2, + "0000:4c:00.0": 3, + "0000:5d:00.0": 4, + "0000:9b:00.0": 5, + "0000:bb:00.0": 6, + "0000:cb:00.0": 7, + "0000:db:00.0": 8, + } + allDevices := []string{"0", "1", "2", "3", "4", "5", "6", "7"} + + // H100 fabric partition table: one 8-GPU, two 4-GPU, four 2-GPU + // "diagonal" pairs and eight singles. Listed in descending-ID order to + // verify selection does not depend on FM list order. + h100Partitions := func() []fm.Partition { + partitions := []fm.Partition{ + {ID: 6, GPUPhysicalIDs: []uint32{6, 8}}, + {ID: 5, GPUPhysicalIDs: []uint32{5, 7}}, + {ID: 4, GPUPhysicalIDs: []uint32{2, 4}}, + {ID: 3, GPUPhysicalIDs: []uint32{1, 3}}, + {ID: 2, GPUPhysicalIDs: []uint32{5, 6, 7, 8}}, + {ID: 1, GPUPhysicalIDs: []uint32{1, 2, 3, 4}}, + {ID: 0, GPUPhysicalIDs: []uint32{1, 2, 3, 4, 5, 6, 7, 8}}, + } + for i := uint32(1); i <= 8; i++ { + partitions = append(partitions, fm.Partition{ID: 6 + i, GPUPhysicalIDs: []uint32{i}}) + } + return partitions + } + + newManager := func(partitions []fm.Partition) *fm.PartitionManager { + return fm.NewPartitionManager(&fakeFMClient{partitions: partitions}, pciToModule, deviceIDToPCI) + } + + Context("SelectPreferred() with the H100 partition table", func() { + DescribeTable("selects a full partition of the requested size", + func(available []string, mustInclude []string, size int, expected []string) { + pm := newManager(h100Partitions()) + preferred, err := pm.SelectPreferred(context.Background(), available, mustInclude, size) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(Equal(expected)) + }, + // size=2 must pick a diagonal NVLink pair, never adjacent {1,2}: + // lowest 2-GPU partition ID is 3 = physical {1,3} = devices {0,2}. + Entry("size 2 picks a diagonal pair, not {1,2}", + allDevices, nil, 2, []string{"0", "2"}), + Entry("size 4 picks the first 4-GPU partition", + allDevices, nil, 4, []string{"0", "1", "2", "3"}), + Entry("size 8 picks the full 8-GPU partition", + allDevices, nil, 8, allDevices), + Entry("size 1 picks a single-GPU partition", + allDevices, nil, 1, []string{"0"}), + // mustInclude device "1" (physical 2) forces partition {2,4}. + Entry("mustInclude steers to the matching pair", + allDevices, []string{"1"}, 2, []string{"1", "3"}), + // Fragmentation: devices 0,1 (physical 1,2) already taken; the + // only fully-available pairs are {5,7} and {6,8}; lowest ID wins. + Entry("fragmentation picks a fully-available pair", + []string{"2", "3", "4", "5", "6", "7"}, nil, 2, []string{"4", "6"}), + Entry("fragmentation picks the fully-available 4-GPU partition", + []string{"2", "3", "4", "5", "6", "7"}, nil, 4, []string{"4", "5", "6", "7"}), + ) + + DescribeTable("returns nil when no partition fits", + func(available []string, mustInclude []string, size int) { + pm := newManager(h100Partitions()) + preferred, err := pm.SelectPreferred(context.Background(), available, mustInclude, size) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(BeNil()) + }, + Entry("no partition of size 3 exists", allDevices, nil, 3), + Entry("no fully-available pair remains", + []string{"0", "1", "4", "5"}, nil, 4), + // mustInclude devices "0","1" (physical 1,2) never share a pair. + Entry("mustInclude devices not in any pair together", + allDevices, []string{"0", "1"}, 2), + Entry("mustInclude exceeds allocation size", allDevices, []string{"0", "2", "4"}, 2), + Entry("size 0 yields no preference", allDevices, nil, 0), + Entry("no available devices", []string{}, nil, 2), + ) + + It("prefers a partition that is already active", func() { + partitions := h100Partitions() + for i := range partitions { + if partitions[i].ID == 6 { // physical {6,8} = devices {5,7} + partitions[i].IsActive = true + } + } + pm := newManager(partitions) + preferred, err := pm.SelectPreferred(context.Background(), allDevices, nil, 2) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(Equal([]string{"5", "7"})) + }) + + It("tie-breaks equal-priority partitions by lowest partition ID", func() { + partitions := h100Partitions() + for i := range partitions { + if partitions[i].ID == 4 || partitions[i].ID == 5 { + partitions[i].IsActive = true + } + } + pm := newManager(partitions) + preferred, err := pm.SelectPreferred(context.Background(), allDevices, nil, 2) + Expect(err).ToNot(HaveOccurred()) + // Both {2,4} (ID 4) and {5,7} (ID 5) are active; ID 4 wins. + Expect(preferred).To(Equal([]string{"1", "3"})) + }) + + It("skips unmappable available devices and still selects a partition", func() { + pm := newManager(h100Partitions()) + available := append([]string{"unknown-device"}, allDevices...) + preferred, err := pm.SelectPreferred(context.Background(), available, nil, 2) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(Equal([]string{"0", "2"})) + }) + + It("returns nil when a must-include device is unmappable", func() { + pm := newManager(h100Partitions()) + preferred, err := pm.SelectPreferred(context.Background(), allDevices, []string{"unknown-device"}, 2) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(BeNil()) + }) + + It("returns an error when fabric manager is unreachable", func() { + pm := fm.NewPartitionManager(&fakeFMClient{err: errors.New("connection refused")}, pciToModule, deviceIDToPCI) + preferred, err := pm.SelectPreferred(context.Background(), allDevices, nil, 2) + Expect(err).To(HaveOccurred()) + Expect(preferred).To(BeNil()) + }) + }) + + Context("NewPartitionManager()", func() { + It("normalizes PCI addresses from both maps", func() { + pm := fm.NewPartitionManager( + &fakeFMClient{partitions: []fm.Partition{{ID: 1, GPUPhysicalIDs: []uint32{1, 2}}}}, + map[string]uint32{"00000000:19:00.0": 1, "0000:3B:00.0": 2}, + map[string]string{"0": "0000:19:00.0", "1": "00000000:3b:00.0"}, + ) + preferred, err := pm.SelectPreferred(context.Background(), []string{"0", "1"}, nil, 2) + Expect(err).ToNot(HaveOccurred()) + Expect(preferred).To(Equal([]string{"0", "1"})) + }) + }) +})