fetch_ml/internal/worker/gpu.go

package worker

import (
	"os"
	"strconv"
	"strings"
)

// _gpuVisibleDevicesString constructs the visible devices string from config
func _gpuVisibleDevicesString(cfg *Config, fallback string) string {
	if cfg == nil {
		return strings.TrimSpace(fallback)
	}
	if len(cfg.GPUVisibleDeviceIDs) > 0 {
		parts := make([]string, 0, len(cfg.GPUVisibleDeviceIDs))
		for _, id := range cfg.GPUVisibleDeviceIDs {
			id = strings.TrimSpace(id)
			if id == "" {
				continue
			}
			parts = append(parts, id)
		}
		return strings.Join(parts, ",")
	}
	if len(cfg.GPUVisibleDevices) == 0 {
		return strings.TrimSpace(fallback)
	}
	parts := make([]string, 0, len(cfg.GPUVisibleDevices))
	for _, v := range cfg.GPUVisibleDevices {
		if v < 0 {
			continue
		}
		parts = append(parts, strconv.Itoa(v))
	}
	return strings.Join(parts, ",")
}

// _filterExistingDevicePaths filters device paths that actually exist
func _filterExistingDevicePaths(paths []string) []string {
	if len(paths) == 0 {
		return nil
	}
	seen := make(map[string]struct{}, len(paths))
	out := make([]string, 0, len(paths))
	for _, p := range paths {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		if _, ok := seen[p]; ok {
			continue
		}
		if _, err := os.Stat(p); err != nil {
			continue
		}
		seen[p] = struct{}{}
		out = append(out, p)
	}
	return out
}

// _gpuVisibleEnvVarName returns the appropriate env var for GPU visibility
func _gpuVisibleEnvVarName(cfg *Config) string {
	if cfg == nil {
		return "CUDA_VISIBLE_DEVICES"
	}
	switch strings.ToLower(strings.TrimSpace(cfg.GPUVendor)) {
	case "amd":
		return "HIP_VISIBLE_DEVICES"
	case string(GPUTypeApple), string(GPUTypeNone):
		return ""
	default:
		return "CUDA_VISIBLE_DEVICES"
	}
}