Skip to content

Commit

Permalink
Add ability to get MIG device info from the NVML API
Browse files Browse the repository at this point in the history
Previously, we relied solely on trying to parse the UUID of the MIG
device (since the API didn't support getting an NVML device handle from
a UUID). Moving forward, however, the UUIDs won't have the correct
structure to parse them. That said, functionality is being added in it's
place to actually get the MIG device handle from the UUID (obsoleting
the old method).

This patch ensures that we use the API to get this info (if available)
and only fall back to parsing the UUID if not (for backwards
compatibiltiy).

Signed-off-by: Kevin Klues <[email protected]>
  • Loading branch information
klueska committed Sep 21, 2020
1 parent 145fcea commit 9531ba1
Showing 1 changed file with 35 additions and 1 deletion.
36 changes: 35 additions & 1 deletion bindings/go/nvml/nvml.go
Original file line number Diff line number Diff line change
Expand Up @@ -733,7 +733,41 @@ func (d *Device) GetMigParentDeviceLite() (*Device, error) {
return NewDeviceLite(*index)
}

func ParseMigDeviceUUID(mig string) (string, uint, uint, error) {
func ParseMigDeviceUUID(uuid string) (string, uint, uint, error) {
migHandle, err := deviceGetHandleByUUID(uuid)
if err == nil {
return getMIGDeviceInfo(migHandle)
}
return parseMigDeviceUUID(uuid)
}

func getMIGDeviceInfo(migHandle handle) (string, uint, uint, error) {
parentHandle, err := migHandle.deviceGetDeviceHandleFromMigDeviceHandle()
if err != nil {
return "", 0, 0, err
}

parentUUID, err := parentHandle.deviceGetUUID()
if err != nil {
return "", 0, 0, err
}

migDevice := Device{handle: migHandle}

gi, err := migDevice.GetGPUInstanceId()
if err != nil {
return "", 0, 0, err
}

ci, err := migDevice.GetComputeInstanceId()
if err != nil {
return "", 0, 0, err
}

return *parentUUID, uint(gi), uint(ci), err
}

func parseMigDeviceUUID(mig string) (string, uint, uint, error) {
tokens := strings.SplitN(mig, "-", 2)
if len(tokens) != 2 || tokens[0] != "MIG" {
return "", 0, 0, fmt.Errorf("Unable to parse UUID as MIG device")
Expand Down

0 comments on commit 9531ba1

Please sign in to comment.