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

Code before:
	}
	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 = errors.Join(errs, err)
			continue
		}
	}
	return errs
}

// injectNVIDIARuntimeHook handles requestes for NVIDIA GPUs.
// This is done by updating the OCI runtime spec to include the correct value
// for the NVIDIA_VISIBLE_DEVICES environment variable and injecting the
// NVIDIA Container Runtime Hook as a container prestart hook.
func injectNVIDIARuntimeHook(s *specs.Spec, dev *deviceInstance) error {
	deviceIDs, err := getRequestedDevicesIDs(dev.req)
	if err != nil {
		return err
	}
	if len(deviceIDs) == 0 {
		return nil
	}
	s.Process.Env = append(s.Process.Env, "NVIDIA_VISIBLE_DEVICES="+strings.Join(deviceIDs, ","))

	var nvidiaCaps []string
	// req.Capabilities contains device capabilities, some but not all are NVIDIA driver capabilities.
	for _, c := range dev.selectedCaps {
		if _, isNvidiaCap := allNvidiaCaps[c]; isNvidiaCap {
			nvidiaCaps = append(nvidiaCaps, c)
			continue
		}
		// TODO: nvidia.WithRequiredCUDAVersion
		// for now we let the prestart hook verify cuda versions but errors are not pretty.

Review comment (suggestion):
IIUC the intent is to return multiple errors for each of the updater, but currently the errors.Join will nest each error in the previous one.

Should it be something like below instead?

```suggestion
	var errs []error
	for _, u := range us {
		if u == nil {
			continue
		}
		if err := u(s, dev); err != nil {
			errs = append(errs, err)
			continue
		}
	}
	return errors.Join(errs...)
```

Code after:
	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 {
			log.G(context.TODO()).WithError(errors.Join(errs...)).Warning("Ignoring previous errors updating spec")
		}
		return nil
	}
	return errors.Join(errs...)
}

// injectNVIDIARuntimeHook handles requests for NVIDIA GPUs.
// This is done by updating the OCI runtime spec to include the correct value
// for the NVIDIA_VISIBLE_DEVICES environment variable and injecting the
// NVIDIA Container Runtime Hook as a container prestart hook.
func injectNVIDIARuntimeHook(s *specs.Spec, dev *deviceInstance) error {
	deviceIDs, err := getRequestedDevicesIDs(dev.req)
	if err != nil {
		return err
	}
	if len(deviceIDs) == 0 {
		return nil
	}
	s.Process.Env = append(s.Process.Env, "NVIDIA_VISIBLE_DEVICES="+strings.Join(deviceIDs, ","))

	var nvidiaCaps []string
	// req.Capabilities contains device capabilities, some but not all are NVIDIA driver capabilities.
	for _, c := range dev.selectedCaps {