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

Code before:
		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.
	}

	if nvidiaCaps != nil {

Review comment (suggestion):
nit: typo
```suggestion
// injectNVIDIARuntimeHook handles requests for NVIDIA GPUs.
```

Code after:
	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 {
		if _, isNvidiaCap := allNvidiaCaps[c]; isNvidiaCap {
			nvidiaCaps = append(nvidiaCaps, c)
			continue