Repository: containerd/containerd
File: cmd/ctr/commands/run/run_unix.go (Go)

Code before:
func validNamespace(ns string) bool {
	linuxNs := specs.LinuxNamespaceType(ns)
	switch linuxNs {
	case specs.PIDNamespace,
		specs.NetworkNamespace,
		specs.UTSNamespace,
		specs.MountNamespace,
		specs.UserNamespace,
		specs.IPCNamespace,
		specs.CgroupNamespace:
		return true
	default:
		return false
	}
}

func getNetNSPath(_ context.Context, task containerd.Task) (string, error) {
	return fmt.Sprintf("/proc/%d/ns/net", task.Pid()), nil
}

// vendorLister is an interface for listing GPU vendors from CDI.
type vendorLister interface {
	ListVendors() []string
}

// detectGPUVendor detects available GPU vendors from CDI spec files.
// It returns the vendor string for the first known GPU vendor found,
// or an empty string if no GPU vendor is detected.
func detectGPUVendor(ctx context.Context, lister vendorLister) (string, error) {
	if lister == nil {
		return "", errors.New("vendor lister is nil")
	}

	availableVendors := lister.ListVendors()
	knownVendors := []string{"nvidia.com", "amd.com"}

	// Check if a known GPU vendor is available
	// NVIDIA is checked first followed by AMD
	for _, known := range knownVendors {
		for _, available := range availableVendors {
			if available == known {
				log.G(ctx).Debugf("Detected GPU vendor from CDI specs: %s", known)
				return known, nil
			}
		}
	}

	return "", fmt.Errorf("no known GPU vendor detected in CDI specs, only AMD and NVIDIA are supported")
}

// gpuDeviceNames converts GPU indices to qualified CDI device names.

Review comment (suggestion):
this list needs to be configurable I think.. Fair to have a default list but we support many platforms and they should probably all be there or at least have an easy way to add.

Code after:

func validNamespace(ns string) bool {
	linuxNs := specs.LinuxNamespaceType(ns)
	switch linuxNs {
	case specs.PIDNamespace,
		specs.NetworkNamespace,
		specs.UTSNamespace,
		specs.MountNamespace,
		specs.UserNamespace,
		specs.IPCNamespace,
		specs.CgroupNamespace:
		return true
	default:
		return false
	}
}

func getNetNSPath(_ context.Context, task containerd.Task) (string, error) {
	return fmt.Sprintf("/proc/%d/ns/net", task.Pid()), nil
}

// detectGPUVendor detects a known GPU vendor from a list of available vendors.
// It returns the vendor string for the first known GPU vendor found in availableVendors,
// or an error if no known GPU vendor is detected.
func detectGPUVendor(ctx context.Context, availableVendors []string) (string, error) {
	knownVendors := []string{"nvidia.com", "amd.com"}
	for _, known := range knownVendors {
		if slices.Contains(availableVendors, known) {
			log.G(ctx).Debugf("Detected GPU vendor from CDI specs: %s", known)
			return known, nil
		}
	}
	return "", fmt.Errorf("no known GPU vendor detected in CDI specs, only AMD and NVIDIA are supported")
}

// gpuDeviceNames converts GPU indices to qualified CDI device names for the given vendor.
func gpuDeviceNames(ctx context.Context, vendor string, gpuIDs ...int) ([]string, error) {
	if len(gpuIDs) == 0 {
		return nil, nil
	}
	if vendor == "" {
		return nil, fmt.Errorf("gpu vendor can't be empty")
	}
	devices := make([]string, 0, len(gpuIDs))
	for _, id := range gpuIDs {
		devices = append(devices, fmt.Sprintf("%s/gpu=%d", vendor, id))
	}
	log.G(ctx).Debugf("Resolved GPU device names: %v", devices)
	return devices, nil
}