Repository: moby/moby
File: daemon/devices_nvidia_linux.go (Go)

Code before:
		return nil
	}

	nvidiaDrivers := make(map[string]*deviceDriver)
	// Register a driver specific to the nvidia-container-runtime-hook if present.
	// This has no capabilities associated with it so as to not inadvertently
	// match requests.
	if _, ok := hasNVIDIAExecutables[nvidiaContainerRuntimeHookExecutableName]; ok {
		nvidiaDrivers["nvidia.runtime-hook"] = &deviceDriver{
			capset:     nil,
			updateSpec: setNvidiaGPUs,
		}
	}

	// Register a driver specific to CDI if present.
	// This has no capabilities associated with it so as to not inadvertently
	// match requests.
	if _, ok := hasNVIDIAExecutables[nvidiaCDIHookExecutableName]; ok {
		nvidiaDrivers["nvidia.cdi"] = &deviceDriver{
			capset: nil,
			updateSpec: (&cdiDeviceInjector{
				defaultCDIDeviceKind: "nvidia.com/gpu",
			}).injectDevices,
		}
	}

	var composite specUpdaters
	// We construct a composite handler that prefers nvidia.cdi over nvidia.runtime-hook
	for _, driverName := range []string{"nvidia.cdi", "nvidia.runtime-hook"} {
		d := nvidiaDrivers[driverName]
		if d == nil || d.updateSpec == nil {
			continue
		}
		composite = append(composite, d.updateSpec)
	}

	// We associate all NVIDIA capabilities with this composite driver.
	capset := capabilities.Set{"gpu": struct{}{}, "nvidia": struct{}{}}
	for c := range allNvidiaCaps {
		capset[c] = struct{}{}
	}
	nvidiaDrivers["nvidia"] = &deviceDriver{
		capset:     capset,
		updateSpec: composite.firstSuccessful,
	}

	return nvidiaDrivers
}

// specUpdaters refer to a list of functions used updated an OCI spec for a
// given device instance.

Review comment (suggestion):
Minor nit; I guess we can skip the `hasNVIDIAExecutables` intermediate map here, and directly use the `nvidiaDrivers` map?


```go
	nvidiaDrivers := make(map[string]*deviceDriver)
	if _, err := exec.LookPath(nvidiaCDIHookExecutableName); err == nil {
		// Register a driver specific to CDI if present.
		// This has no capabilities associated to not inadvertently match requests.
		nvidiaDrivers["nvidia.cdi"] = &deviceDriver{
			updateSpec: (&cdiDeviceInjector{
				defaultCDIDeviceKind: "nvidia.com/gpu",
			}).injectDevices,
		}
	}

	if _, err := exec.LookPath(nvidiaContainerRuntimeHookExecutableName); err == nil {
		// Register a driver specific to the nvidia-container-runtime-hook if present.
		// This has no capabilities associated to not inadvertently match requests.
		nvidiaDrivers["nvidia.runtime-hook"] = &deviceDriver{
			updateSpec: setNvidiaGPUs,
		}
	}

	if len(nvidiaDrivers) == 0 {
		return nil
	}
```


Perhaps even the `composite` could be rolled up into the above, but it might be clearer to construct that below to express the "order is important here".


Code after:
			updateSpec: (&cdiDeviceInjector{
				defaultCDIDeviceKind: "nvidia.com/gpu",
			}).injectDevices,
		}
		nvidiaDrivers["nvidia.cdi"] = cdiDeviceDriver
		composite = append(composite, cdiDeviceDriver.updateSpec)
	}

	if _, err := exec.LookPath(nvidiaContainerRuntimeHookExecutableName); err == nil {
		// Register a driver specific to the nvidia-container-runtime-hook if present.
		// This has no capabilities associated to not inadvertently match requests.
		runtimeHookDeviceDriver := &deviceDriver{
			updateSpec: injectNVIDIARuntimeHook,
		}
		nvidiaDrivers["nvidia.runtime-hook"] = runtimeHookDeviceDriver
		composite = append(composite, runtimeHookDeviceDriver.updateSpec)
	}

	if len(nvidiaDrivers) == 0 {
		return nil
	}

	// We associate all NVIDIA capabilities with the composite updater
	capset := capabilities.Set{"gpu": struct{}{}, "nvidia": struct{}{}}
	for c := range allNvidiaCaps {
		capset[c] = struct{}{}
	}
	nvidiaDrivers["nvidia"] = &deviceDriver{
		capset:     capset,
		updateSpec: composite.updateSpec,
	}

	return nvidiaDrivers
}

// specUpdaters refer to a list of functions used updated an OCI spec for a
// given device instance.
type firstSuccessfulUpdater []func(*specs.Spec, *deviceInstance) error

// updateSpec returns on the first successful spec update.
func (us firstSuccessfulUpdater) updateSpec(s *specs.Spec, dev *deviceInstance) error {
	var errs []error
	for _, u := range us {
		if u == nil {
			continue
		}
		if err := u(s, dev); err != nil {
			errs = append(errs, err)
			continue
		}
		if len(errs) > 0 {