UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-04-02 12:00:18 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Add Index type to nvcaps

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2025-03-12 12:43:07 +02:00
parent f3b730c805
commit d4b331fbbb
No known key found for this signature in database
1 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
internal/nvcaps/nvcaps.go
View File

@ -36,8 +36,12 @@ const (
nvcapsDevicePath = "/dev/nvidia-caps"
)
// An Index represents a gpu, ci, or gi index.
// We use uint32 as this typically maps to a device minor number.
type Index uint32
// MigMinor represents the minor number of a MIG device
type MigMinor uint32
type MigMinor Index
// MigCap represents the path to a MIG cap file.
// These are listed in /proc/driver/nvidia-caps/mig-minors and have one of the
@ -53,30 +57,30 @@ type MigCaps map[MigCap]MigMinor
// NewGPUInstanceCap creates a MigCap for the specified MIG GPU instance.
// A GPU instance is uniquely defined by the GPU minor number and GI instance ID.
func NewGPUInstanceCap(gpu, gi int) MigCap {
func NewGPUInstanceCap[T uint32 | int | Index](gpu, gi T) MigCap {
return MigCap(fmt.Sprintf("gpu%d/gi%d/access", gpu, gi))
}
// NewComputeInstanceCap creates a MigCap for the specified MIG Compute instance.
// A GPU instance is uniquely defined by the GPU minor number, GI instance ID, and CI instance ID.
func NewComputeInstanceCap(gpu, gi, ci int) MigCap {
func NewComputeInstanceCap[T uint32 | int | Index](gpu, gi, ci T) MigCap {
return MigCap(fmt.Sprintf("gpu%d/gi%d/ci%d/access", gpu, gi, ci))
}
// FilterForGPU limits the MIG Caps to those associated with a particular GPU.
func (m MigCaps) FilterForGPU(gpu int) MigCaps {
func (m MigCaps) FilterForGPU(gpu Index) MigCaps {
if m == nil {
return nil
}
filtered := make(MigCaps)
for gi := 0; ; gi++ {
for gi := Index(0); ; gi++ {
giCap := NewGPUInstanceCap(gpu, gi)
giMinor, exist := m[giCap]
if !exist {
break
}
filtered[giCap] = giMinor
for ci := 0; ; ci++ {
for ci := Index(0); ; ci++ {
ciCap := NewComputeInstanceCap(gpu, gi, ci)
ciMinor, exist := m[ciCap]
if !exist {