diff --git a/internal/rm/device_map.go b/internal/rm/device_map.go index b4fe345c9e..2100facb3e 100644 --- a/internal/rm/device_map.go +++ b/internal/rm/device_map.go @@ -41,17 +41,18 @@ type deviceMapBuilder struct { type DeviceMap map[spec.ResourceName]Devices // NewDeviceMap creates a device map for the specified NVML library and config. -func NewDeviceMap(infolib info.Interface, devicelib device.Interface, config *spec.Config) (DeviceMap, error) { +func NewDeviceMap(devicelib device.Interface, config *spec.Config, platform info.Platform) (DeviceMap, error) { + newGPUDevice := newNvmlGPUDevice + if platform == info.PlatformWSL { + newGPUDevice = newWslAllGPUsDevice + } + b := deviceMapBuilder{ Interface: devicelib, migStrategy: config.Flags.MigStrategy, resources: &config.Resources, replicatedResources: config.Sharing.ReplicatedResources(), - newGPUDevice: newNvmlGPUDevice, - } - - if infolib.ResolvePlatform() == info.PlatformWSL { - b.newGPUDevice = newWslGPUDevice + newGPUDevice: newGPUDevice, } return b.build() diff --git a/internal/rm/device_map_test.go b/internal/rm/device_map_test.go index 82e4270387..08a883ab35 100644 --- a/internal/rm/device_map_test.go +++ b/internal/rm/device_map_test.go @@ -25,6 +25,30 @@ import ( spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) +func TestWslDeviceMapHasSingleAllDevice(t *testing.T) { + // Simulate building a GPU device map with 3 GPUs on WSL. + // Because newWslAllGPUsDevice always returns index/UUID "all", the map + // should collapse to exactly one device entry per resource. + devices := make(DeviceMap) + resourceName := spec.ResourceName("nvidia.com/gpu") + + for i := 0; i < 3; i++ { + index, info := newWslAllGPUsDevice(i, nil) + err := devices.setEntry(resourceName, index, info) + require.NoError(t, err) + } + + gpuDevices, ok := devices[resourceName] + require.True(t, ok) + require.Len(t, gpuDevices, 1) + + dev, ok := gpuDevices["all"] + require.True(t, ok) + require.Equal(t, "all", dev.ID) + require.Equal(t, "all", dev.Index) + require.Equal(t, []string{"/dev/dxg"}, dev.Paths) +} + func TestDeviceMapInsert(t *testing.T) { device0 := Device{Device: pluginapi.Device{ID: "0"}} device0withIndex := Device{Device: pluginapi.Device{ID: "0"}, Index: "index"} diff --git a/internal/rm/nvml_devices.go b/internal/rm/nvml_devices.go index 11987b3088..14b69bda65 100644 --- a/internal/rm/nvml_devices.go +++ b/internal/rm/nvml_devices.go @@ -49,9 +49,8 @@ func newNvmlGPUDevice(i int, gpu nvml.Device) (string, deviceInfo) { return index, nvmlDevice{gpu} } -func newWslGPUDevice(i int, gpu nvml.Device) (string, deviceInfo) { - index := fmt.Sprintf("%v", i) - return index, wslDevice{gpu} +func newWslAllGPUsDevice(_ int, _ nvml.Device) (string, deviceInfo) { + return "all", wslAllGPUsDevice{} } func newMigDevice(i int, j int, mig nvml.Device) (string, nvmlMigDevice) { diff --git a/internal/rm/nvml_manager.go b/internal/rm/nvml_manager.go index fac9234291..f0ca775d1b 100644 --- a/internal/rm/nvml_manager.go +++ b/internal/rm/nvml_manager.go @@ -48,7 +48,9 @@ func NewNVMLResourceManagers(infolib info.Interface, nvmllib nvml.Interface, dev } }() - deviceMap, err := NewDeviceMap(infolib, devicelib, config) + platform := infolib.ResolvePlatform() + + deviceMap, err := NewDeviceMap(devicelib, config, platform) if err != nil { return nil, fmt.Errorf("error building device map: %v", err) } @@ -58,15 +60,25 @@ func NewNVMLResourceManagers(infolib info.Interface, nvmllib nvml.Interface, dev if len(devices) == 0 { continue } - r := &nvmlResourceManager{ - resourceManager: resourceManager{ - config: config, - resource: resourceName, - devices: devices, - }, - nvml: nvmllib, + + resources := resourceManager{ + config: config, + resource: resourceName, + devices: devices, } - rms = append(rms, r) + + var rm ResourceManager + switch platform { + case info.PlatformWSL: + rm = &resources + default: + rm = &nvmlResourceManager{ + resourceManager: resources, + nvml: nvmllib, + } + } + + rms = append(rms, rm) } return rms, nil @@ -87,7 +99,7 @@ func (r *nvmlResourceManager) GetDevicePaths(ids []string) []string { "/dev/nvidia-modeset", } - return append(paths, r.Devices().Subset(ids).GetPaths()...) + return append(paths, r.resourceManager.GetDevicePaths(ids)...) } // CheckHealth performs health checks on a set of devices, writing to the 'unhealthy' channel with any unhealthy devices diff --git a/internal/rm/rm.go b/internal/rm/rm.go index 33f44b9d85..742762a351 100644 --- a/internal/rm/rm.go +++ b/internal/rm/rm.go @@ -48,6 +48,23 @@ type ResourceManager interface { ValidateRequest(AnnotatedIDs) error } +var _ ResourceManager = (*resourceManager)(nil) + +// CheckHealth is disabled on the base resourceManager. +func (r *resourceManager) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error { + return nil +} + +// GetDevicePaths returns the paths for the devices associated with the resource manager. +func (r *resourceManager) GetDevicePaths(ids []string) []string { + return r.Devices().Subset(ids).GetPaths() +} + +// GetPreferredAllocation runs an allocation algorithm over the inputs. +func (r *resourceManager) GetPreferredAllocation(available, required []string, size int) ([]string, error) { + return r.distributedAlloc(available, required, size) +} + // Resource gets the resource name associated with the ResourceManager func (r *resourceManager) Resource() spec.ResourceName { return r.resource diff --git a/internal/rm/tegra_manager.go b/internal/rm/tegra_manager.go index 65ca2022f8..92bfbdf7a5 100644 --- a/internal/rm/tegra_manager.go +++ b/internal/rm/tegra_manager.go @@ -22,12 +22,6 @@ import ( spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) -type tegraResourceManager struct { - resourceManager -} - -var _ ResourceManager = (*tegraResourceManager)(nil) - // NewTegraResourceManagers returns a set of ResourceManagers for tegra resources func NewTegraResourceManagers(config *spec.Config) ([]ResourceManager, error) { deviceMap, err := buildTegraDeviceMap(config) @@ -45,32 +39,13 @@ func NewTegraResourceManagers(config *spec.Config) ([]ResourceManager, error) { if len(devices) == 0 { continue } - r := &tegraResourceManager{ - resourceManager: resourceManager{ - config: config, - resource: resourceName, - devices: devices, - }, - } - if len(devices) != 0 { - rms = append(rms, r) + r := &resourceManager{ + config: config, + resource: resourceName, + devices: devices, } + rms = append(rms, r) } return rms, nil } - -// GetPreferredAllocation returns a standard allocation for the Tegra resource manager. -func (r *tegraResourceManager) GetPreferredAllocation(available, required []string, size int) ([]string, error) { - return r.distributedAlloc(available, required, size) -} - -// GetDevicePaths returns an empty slice for the tegraResourceManager -func (r *tegraResourceManager) GetDevicePaths(ids []string) []string { - return nil -} - -// CheckHealth is disabled for the tegraResourceManager -func (r *tegraResourceManager) CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error { - return nil -} diff --git a/internal/rm/wsl_devices.go b/internal/rm/wsl_devices.go index b8319409f7..1d442acc11 100644 --- a/internal/rm/wsl_devices.go +++ b/internal/rm/wsl_devices.go @@ -16,31 +16,31 @@ package rm -type wslDevice nvmlDevice +type wslAllGPUsDevice struct{} -var _ deviceInfo = (*wslDevice)(nil) +var _ deviceInfo = (*wslAllGPUsDevice)(nil) -// GetUUID returns the UUID of the device -func (d wslDevice) GetUUID() (string, error) { - return nvmlDevice(d).GetUUID() +// GetUUID returns "all" to represent all GPUs accessible via /dev/dxg on WSL. +func (d wslAllGPUsDevice) GetUUID() (string, error) { + return "all", nil } -// GetPaths returns the paths for a tegra device. -func (d wslDevice) GetPaths() ([]string, error) { +// GetPaths returns the WSL GPU device path. +func (d wslAllGPUsDevice) GetPaths() ([]string, error) { return []string{"/dev/dxg"}, nil } -// GetNumaNode returns the NUMA node associated with the GPU device -func (d wslDevice) GetNumaNode() (bool, int, error) { - return nvmlDevice(d).GetNumaNode() +// GetNumaNode returns no NUMA node association for WSL devices. +func (d wslAllGPUsDevice) GetNumaNode() (bool, int, error) { + return false, 0, nil } -// GetTotalMemory returns the total memory available on the device. -func (d wslDevice) GetTotalMemory() (uint64, error) { - return nvmlDevice(d).GetTotalMemory() +// GetTotalMemory returns 0 as memory info is not available for WSL devices. +func (d wslAllGPUsDevice) GetTotalMemory() (uint64, error) { + return 0, nil } -// GetComputeCapability returns the CUDA compute capability for the device. -func (d wslDevice) GetComputeCapability() (string, error) { - return nvmlDevice(d).GetComputeCapability() +// GetComputeCapability returns an empty string as compute capability is not available for WSL devices. +func (d wslAllGPUsDevice) GetComputeCapability() (string, error) { + return "", nil }