dtk_image.go

/**
# Copyright (c) 2024, HCUOpt CORPORATION.  All rights reserved.
**/

package image

import (
	"fmt"
	"path/filepath"
	"strings"
        "dcu-container-toolkit/internal/hydcu"
	"strconv"
	"github.com/opencontainers/runtime-spec/specs-go"
	"tags.cncf.io/container-device-interface/pkg/parser"
)

// DTK represents a DTK image that can be used for HCU computing. This wraps
// a map of environment variable to values that can be used to perform lookups
// such as requirements.
type DTK struct {
	env          map[string]string
	mounts       []specs.Mount
	ContainerId  string
}

// NewDTKImageFromSpec creates a DTK image from the input OCI runtime spec.
// The process environment is read (if present) to construct the DTK Image.
func NewDTKImageFromSpec(spec *specs.Spec) (DTK, error) {
	var env []string
	if spec != nil && spec.Process != nil {
		env = spec.Process.Env
	}

	return New(
		WithEnv(env),
		WithMounts(spec.Mounts),
	)
}

// LoadHyhalMethod
func (i DTK) LoadHyhalMethod(envVar string) bool {
	if devs, ok := i.env[envVar]; ok {
		if devs == "copy" {
			return false
		}
	}
	return true
}

func (i DTK) IsLegacy() bool {
	rocm_version := i.env[EnvROCmVersion]
	return len(rocm_version) > 0
}

// Getenv returns the value of the specified environment variable.
// If the environment variable is not specified, an empty string is returned.
func (i DTK) Getenv(key string) string {
	return i.env[key]
}

func Contains[T comparable](slice []T, val T) bool {
	for _, item := range slice {
		if item == val {
			return true
		}
	}
	return false
}

// DevicesFromEnvvars returns the devices requested by the image through environment variables
func (i DTK) DevicesFromEnvvars(envVars ...string) VisibleDevices {
	isHexString := func(s string) bool {
		if len(s) == 0 {
			return false
		}
		for _, c := range s {
			if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) {
				return false
			}
		}
		return true
	}

	// We concantenate all the devices from the specified env.
	var isSet bool
	var devices []string
	requested := make(map[string]bool)
	for _, envVar := range envVars {
		if devs, ok := i.env[envVar]; ok {
			isSet = true
			for _, d := range strings.Split(devs, ",") {
				trimmed := strings.TrimSpace(d)
				if len(trimmed) == 0 || requested[trimmed] {
					continue
				}
				devices = append(devices, trimmed)
				requested[trimmed] = true
			}
		}
	}

	uuidToDCUIdMap, err := hydcu.GetUniqueIdToDeviceIndexMap()
	if err != nil {
		uuidToDCUIdMap = make(map[string][]int)
	}
	for key, value := range i.env {
		if strings.HasPrefix(key, "DOCKER_RESOURCE_") {
			for _, c := range strings.Split(value, ",") {
				if strings.HasPrefix(c, "0x") || strings.HasPrefix(c, "0X") ||
					(len(c) > 8 && isHexString(c)) {
					uuid := strings.ToLower(c)
					if !strings.HasPrefix(uuid, "0x") {
						uuid = "0x" + uuid
					}
					if dcuId, exists := uuidToDCUIdMap[uuid]; exists {
						if !Contains(devices, strconv.Itoa(dcuId[0])) {
							devices = append(devices, strconv.Itoa(dcuId[0]))
						}
					} else {
						uuid = strings.TrimPrefix(uuid, "0x")
						if dcuId, exists := uuidToDCUIdMap[uuid]; exists {
							devices = append(devices, strconv.Itoa(dcuId[0]))
						}
					}
				}
			}
			break
		}
	}

	// Environment variable unset with legacy image: default to "all".
	if !isSet && len(devices) == 0 && i.IsLegacy() {
		return NewVisibleDevices("all")
	}

	// Environment variable unset or empty or "void": return nil
	if len(devices) == 0 || requested["void"] {
		return NewVisibleDevices("void")
	}

	return NewVisibleDevices(devices...)
}

// OnlyFullyQualifiedCDIDevices returns true if all devices requested in the image are requested as CDI devices/
func (i DTK) OnlyFullyQualifiedCDIDevices() bool {
	var hasCDIdevice bool
	for _, device := range i.DevicesFromEnvvars(EnvVarDTKVisibleDevices, EnvVarNvidiaVisibleDevices).List() {
		if !parser.IsQualifiedName(device) {
			return false
		}
		hasCDIdevice = true
	}

	for _, device := range i.DevicesFromMounts() {
		if !strings.HasPrefix(device, "cdi/") {
			return false
		}
		hasCDIdevice = true
	}
	return hasCDIdevice
}

const (
	deviceListAsVolumeMountsRoot = "/var/run/dtk-container-devices"
)

// DevicesFromMounts returns a list of device specified as mounts.
// TODO: This should be merged with getDevicesFromMounts used in the DTK Container Runtime
func (i DTK) DevicesFromMounts() []string {
	root := filepath.Clean(deviceListAsVolumeMountsRoot)
	seen := make(map[string]bool)
	var devices []string
	for _, m := range i.mounts {
		source := filepath.Clean(m.Source)
		// Only consider mounts who's host volume is /dev/null
		if source != "/dev/null" {
			continue
		}

		destination := filepath.Clean(m.Destination)
		if seen[destination] {
			continue
		}
		seen[destination] = true

		// Only consider container mount points that begin with 'root'
		if !strings.HasPrefix(destination, root) {
			continue
		}

		// Grab the full path beyond 'root' and add it to the list of devices
		device := strings.Trim(strings.TrimPrefix(destination, root), "/")
		if len(device) == 0 {
			continue
		}
		devices = append(devices, device)
	}
	return devices
}

// CDIDevicesFromMounts returns a list of CDI devices specified as mounts on the image.
func (i DTK) CDIDevicesFromMounts() []string {
	var devices []string
	for _, mountDevice := range i.DevicesFromMounts() {
		if !strings.HasPrefix(mountDevice, "cdi/") {
			continue
		}
		parts := strings.SplitN(strings.TrimPrefix(mountDevice, "cdi/"), "/", 3)
		if len(parts) != 3 {
			continue
		}
		vendor := parts[0]
		class := parts[1]
		device := parts[2]
		devices = append(devices, fmt.Sprintf("%s/%s=%s", vendor, class, device))
	}
	return devices
}

// VisibleDevicesFromEnvVar returns the set of visible devices requested through
// the DCU_VISIBLE_DEVICES environment variable.
func (i DTK) VisibleDevicesFromEnvVar() []string {
	return i.DevicesFromEnvvars(EnvVarDTKVisibleDevices, EnvVarNvidiaVisibleDevices).List()
}