diff --git a/mcv/Makefile b/mcv/Makefile index 8d42a8396..7e2a277b5 100644 --- a/mcv/Makefile +++ b/mcv/Makefile @@ -184,6 +184,19 @@ image-amd: ## [dev-only] Build AMD container image (with ROCm libraries) -f images/Containerfile \ .. +.PHONY: image-gaudi +image-gaudi: ## [dev-only] Build Intel Gaudi container image (with hl-smi) + @if [ -z "$(CONTAINER_RUNTIME)" ]; then \ + echo "Error: No container runtime found. Please install docker or podman"; \ + exit 1; \ + fi + @echo "Building gaudi specific MCV image (Gaudi) for linux/$(GOARCH) using $(CONTAINER_RUNTIME)..." + $(CONTAINER_RUNTIME) build --platform linux/$(GOARCH) \ + --target mcv-gaudi \ + -t $(IMAGE_REGISTRY)/$(IMAGE_REPOSITORY)/$(IMAGE_NAME):gaudi \ + -f images/Containerfile \ + .. + .PHONY: image-nvidia image-nvidia: ## [dev-only] Build NVIDIA container image (with CUDA+NVML) @if [ -z "$(CONTAINER_RUNTIME)" ]; then \ diff --git a/mcv/docs/unified-mcv-container.md b/mcv/docs/unified-mcv-container.md index 654145dc7..ab853cfde 100644 --- a/mcv/docs/unified-mcv-container.md +++ b/mcv/docs/unified-mcv-container.md @@ -163,10 +163,38 @@ docker run --rm --gpus all quay.io/gkm/mcv:unified \ **Intel Gaudi** (nodes are `0666` — device access only, no group needed): ```bash +# Podman — --device accepts the whole /dev/accel directory. podman run --rm --device /dev/accel quay.io/gkm/mcv:gaudi \ --check-compat --image quay.io/myorg/cache:v1 ``` +To verify Gaudi device access from the container without a cache image, run +`--gpu-info` (queries `hl-smi` only — no cache dir or image build): + +```bash +# Docker — --device does NOT accept a directory, so enumerate the nodes. +# seccomp/apparmor are unconfined because mcv calls buildah.InitReexec() at +# startup (for every subcommand), which sets up an unprivileged user namespace +# that Ubuntu 24.04's default container profile blocks. +docker run --rm \ + --user "$(id -u):$(id -g)" \ + --security-opt seccomp=unconfined \ + --security-opt apparmor=unconfined \ + $(printf -- '--device=%s ' /dev/accel/accel*) \ + quay.io/gkm/mcv:gaudi \ + --gpu-info +``` + +Expected output lists the detected fleet, e.g. `GPU Type: HL-325L` (gaudi3) with +one ID per device. The `newuidmap/newgidmap … Falling back to single mapping` +warnings from `InitReexec` are benign for `--gpu-info`. + +> **Gaudi preflight is backend-only.** Habana Synapse recipe caches don't encode +> the target architecture or warp size in their recipe filenames, so MCV matches +> Gaudi caches on backend (`hpu`) alone — not on a specific Gaudi generation. A +> `gaudi2` cache will pass `--check-compat` on a `gaudi3` host and vice versa; +> validate generation compatibility out of band if it matters for your models. + **AMD** (render/DRI nodes are group-owned — add the group): ```bash diff --git a/mcv/images/Containerfile b/mcv/images/Containerfile index 2a626c096..aeb1e5344 100644 --- a/mcv/images/Containerfile +++ b/mcv/images/Containerfile @@ -164,6 +164,101 @@ LABEL variant="amd" # COPY and ENTRYPOINT are inherited from mcv-minimal base stage +# ============================================================================ +# GAUDI TARGET: For Intel Gaudi GPU validation with hl-smi +# Includes Habana Labs tools for Gaudi device detection +# Build: docker build --target mcv-gaudi -t quay.io/gkm/mcv:gaudi -f mcv/images/Containerfile . +# Usage: mcv --check-compat --image foo (on Intel Gaudi systems) +# mcv --extract --image foo (with Gaudi GPU preflight check) +# ============================================================================ +FROM public.ecr.aws/docker/library/ubuntu:24.04 AS mcv-gaudi + +ARG TARGETARCH + +# Gaudi only publishes amd64 packages; fail fast on other architectures. +RUN [ "$TARGETARCH" = "amd64" ] || \ + { echo "ERROR: mcv-gaudi requires amd64 - Habana does not support ${TARGETARCH}"; exit 1; } + +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgpgme11t64 \ + libbtrfs0 \ + libffi8 \ + libc6 \ + ca-certificates \ + buildah \ + netavark aardvark-dns \ + hwdata \ + python3 \ + python3-setuptools \ + python3-wheel \ + wget \ + gnupg2 \ + pciutils \ + && rm -rf /var/lib/apt/lists/* + +RUN mkdir -p /etc/containers && \ + printf '[storage]\ndriver="vfs"\nrunroot="/home/appuser/.local/share/containers/runroot"\ngraphroot="/home/appuser/.local/share/containers/storage"\n' \ + > /etc/containers/storage.conf + +# Install hl-smi via Habana apt repo (pattern from HabanaAI/Setup_and_Install) +# Download key to a temp file and verify its fingerprint before trusting it to +# authenticate packages from the same host. HABANA_SIGNING_FP must be set to the +# canonical fingerprint published in Intel/Habana's official documentation +# (https://docs.habana.ai) — update it whenever the signing key rotates. +ARG HABANA_SIGNING_FP="6D4D7C0F52A263F383D782791E676CE836A2DE65" +RUN wget -q -O /tmp/habana-key.asc https://vault.habana.ai/artifactory/api/gpg/key/public && \ + gpg --dearmor < /tmp/habana-key.asc > /tmp/habana-key.gpg && \ + actual=$(gpg --no-default-keyring --keyring /tmp/habana-key.gpg --fingerprint 2>/dev/null \ + | awk '/^ /{gsub(/ /,"",$0); print}' | head -1) && \ + expected=$(printf '%s' "${HABANA_SIGNING_FP}" | tr -d ' :') && \ + [ "$actual" = "$expected" ] || { echo "Habana GPG key fingerprint mismatch: got $actual expected $expected"; exit 1; } && \ + mv /tmp/habana-key.gpg /usr/share/keyrings/habana-artifactory.gpg && \ + rm /tmp/habana-key.asc && \ + chmod 644 /usr/share/keyrings/habana-artifactory.gpg && \ + echo "deb [signed-by=/usr/share/keyrings/habana-artifactory.gpg] https://vault.habana.ai/artifactory/debian noble main" \ + > /etc/apt/sources.list.d/habana.list && \ + apt-get update && \ + apt-get install -y --no-install-recommends habanalabs-firmware-tools && \ + rm -f /etc/apt/sources.list.d/habana.list && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +COPY --from=builder /usr/src/mcv/_output/bin/linux_${TARGETARCH}/mcv /mcv +COPY mcv/images/entrypoint.sh /entrypoint.sh + +RUN chmod +x /entrypoint.sh + +# Allow non-root user to run commands. +# Drop any pre-existing ubuntu user/group that claims UID/GID 1000 (present in +# ubuntu:24.04 and nvcr.io/nvidia/cuda:*-ubuntu24.04 base images). Without this, +# useradd -u 1000 fails and the silent fallback creates appuser at UID 1001 — +# K8s runAsUser: 1000 then enters the leftover ubuntu account, not appuser. +RUN userdel -r ubuntu 2>/dev/null; groupdel ubuntu 2>/dev/null; \ + groupadd -g 1000 appgroup && \ + useradd -u 1000 -g 1000 -m -s /bin/bash appuser +RUN test "$(id -u appuser)" = "1000" +RUN chown appuser:1000 /mcv +RUN chown appuser:1000 /entrypoint.sh +# Pre-create buildah's per-user storage so storage.GetStore (running as UID 1000) +# uses these appuser-owned runroot/graphroot paths instead of trying to create +# root-owned /run/containers or /var/lib/containers at runtime. Also create +# ~/.config/containers: when Docker runs the image as the named appuser it sets +# HOME=/home/appuser, and buildah stats $HOME/.config during setup — it errors +# if that directory is absent. +RUN mkdir -p /home/appuser/.local/share/containers/storage \ + /home/appuser/.local/share/containers/runroot \ + /home/appuser/.config/containers && \ + chown -R appuser:1000 /home/appuser/.local /home/appuser/.config +WORKDIR /app +RUN chown -R appuser:1000 /app +USER appuser + +LABEL description="MCV Gaudi - includes Habana Labs tools for Intel Gaudi GPU detection and validation" +LABEL variant="gaudi" + +ENTRYPOINT ["/entrypoint.sh"] +CMD ["/mcv"] + # ============================================================================ # NVIDIA TARGET: For NVIDIA GPU validation with CUDA/NVML support # Includes CUDA runtime and NVML libraries for GPU detection @@ -244,6 +339,9 @@ ARG AMDGPU_VERSION=7.0.1.70001 ARG OPT_ROCM_VERSION=7.0.1 # SHA-256 of amdgpu-install_7.0.1.70001-1_all.deb from repo.radeon.com — update when bumping AMDGPU_VERSION ARG AMDGPU_INSTALLER_SHA256=f4cec24612039c03271e6ab494bc1e18cb5647d59188755aa8e31b6d74bb06df +# HABANA_SIGNING_FP must be set to the canonical fingerprint published in Intel/Habana's official documentation +# (https://docs.habana.ai) — update it whenever the signing key rotates. +ARG HABANA_SIGNING_FP="6D4D7C0F52A263F383D782791E676CE836A2DE65" # Install base runtime dependencies # CUDA base already provides: libnvidia-ml.so.1 (NVML library) @@ -278,6 +376,20 @@ RUN if [ "$TARGETARCH" = "amd64" ]; then \ amd-smi-lib rocm-smi-lib libdrm2 && \ ln -s /opt/rocm-${OPT_ROCM_VERSION}/bin/amd-smi /usr/bin/amd-smi && \ ln -s /opt/rocm-${OPT_ROCM_VERSION}/bin/rocm-smi /usr/bin/rocm-smi && \ + wget -q -O /tmp/habana-key.asc https://vault.habana.ai/artifactory/api/gpg/key/public && \ + gpg --dearmor < /tmp/habana-key.asc > /tmp/habana-key.gpg && \ + actual=$(gpg --no-default-keyring --keyring /tmp/habana-key.gpg --fingerprint 2>/dev/null \ + | awk '/^ /{gsub(/ /,"",$0); print}' | head -1) && \ + expected=$(printf '%s' "${HABANA_SIGNING_FP}" | tr -d ' :') && \ + [ "$actual" = "$expected" ] || { echo "Habana GPG key fingerprint mismatch: got $actual expected $expected"; exit 1; } && \ + mv /tmp/habana-key.gpg /usr/share/keyrings/habana-artifactory.gpg && \ + rm /tmp/habana-key.asc && \ + chmod 644 /usr/share/keyrings/habana-artifactory.gpg && \ + echo "deb [signed-by=/usr/share/keyrings/habana-artifactory.gpg] https://vault.habana.ai/artifactory/debian noble main" \ + > /etc/apt/sources.list.d/habana.list && \ + apt-get update && \ + apt-get install -y --no-install-recommends habanalabs-firmware-tools && \ + rm -f /etc/apt/sources.list.d/habana.list && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* /tmp/amdgpu-install.deb; \ fi diff --git a/mcv/pkg/accelerator/devices/amd.go b/mcv/pkg/accelerator/devices/amd.go index 4d623757b..15e264767 100644 --- a/mcv/pkg/accelerator/devices/amd.go +++ b/mcv/pkg/accelerator/devices/amd.go @@ -12,11 +12,15 @@ import ( "time" "github.com/redhat-et/GKM/mcv/pkg/config" + "github.com/redhat-et/GKM/mcv/pkg/constants" "github.com/redhat-et/GKM/mcv/pkg/utils" logging "github.com/sirupsen/logrus" ) -const amdHwType = config.GPU +const ( + amdHwType = config.GPU + gfxArchMI210 = "gfx90a" // Aldebaran/MI200 [Instinct MI210] GFX architecture +) var ( amdAccImpl = gpuAMD{} @@ -330,7 +334,7 @@ func (r *gpuAMD) Init() error { Arch: TranslateGPUToArch(info.Board.ProductName), WarpSize: 64, MemoryTotalMB: memTotal, - Backend: hipBackend, + Backend: constants.BackendHIP, ID: gpuID, }, Summary: DeviceSummary{ diff --git a/mcv/pkg/accelerator/devices/device.go b/mcv/pkg/accelerator/devices/device.go index 48c7383bf..ee242a8a2 100644 --- a/mcv/pkg/accelerator/devices/device.go +++ b/mcv/pkg/accelerator/devices/device.go @@ -35,11 +35,7 @@ const ( AMD NVML ROCM - - // GPU architecture and backend constants - gfxArchMI210 = "gfx90a" - hipBackend = "hip" - stubbedAMDName = "STUBBED AMD" + GAUDI ) var ( @@ -76,7 +72,7 @@ type CachedDevice struct { } func (d DeviceType) String() string { - return [...]string{"MOCK", "AMD", "NVML", "ROCM"}[d] + return [...]string{"MOCK", "AMD", "NVML", "ROCM", "GAUDI"}[d] } type Device interface { @@ -160,6 +156,7 @@ func registerDevices(r *Registry) { amdCheck(r) rocmCheck(r) nvmlCheck(r) + gaudiCheck(r) } } @@ -301,6 +298,8 @@ func Startup(a string, registry *Registry) Device { device = &gpuNvml{} case ROCM: device = &gpuROCm{} + case GAUDI: + device = &gpuGaudi{} default: logging.Errorf("Unsupported device type %s", cachedDevice.DeviceType.String()) return nil diff --git a/mcv/pkg/accelerator/devices/gaudi.go b/mcv/pkg/accelerator/devices/gaudi.go new file mode 100644 index 000000000..c0870a817 --- /dev/null +++ b/mcv/pkg/accelerator/devices/gaudi.go @@ -0,0 +1,309 @@ +package devices + +import ( + "bufio" + "context" + "fmt" + "os/exec" + "strconv" + "strings" + "time" + + "github.com/redhat-et/GKM/mcv/pkg/config" + "github.com/redhat-et/GKM/mcv/pkg/constants" + "github.com/redhat-et/GKM/mcv/pkg/utils" + logging "github.com/sirupsen/logrus" +) + +const gaudiHwType = config.GPU + +var ( + gaudiAccImpl = gpuGaudi{} + gaudiType DeviceType +) + +const ( + productHL325L = "HL-325L" + productHL325 = "HL-325" + productHL225 = "HL-225" + productHL225B = "HL-225B" + archGaudi3 = "gaudi3" + archGaudi2 = "gaudi2" +) + +// productToArch maps hl-smi product names to architecture strings. +var productToArch = map[string]string{ + productHL325L: archGaudi3, + productHL325: archGaudi3, + productHL225: archGaudi2, + productHL225B: archGaudi2, +} + +type gpuGaudi struct { + name string + deviceType DeviceType + hwType string + tritonInfo []TritonGPUInfo + summaries []DeviceSummary + devices map[int]GPUDevice +} + +type hlsmiDevice struct { + Index int + Name string + BusID string + DriverVersion string + UUID string + ModuleID int + MemoryTotal int64 + MemoryFree int64 + MemoryUsed int64 +} + +func (g *gpuGaudi) SetName(name string) { + g.name = name +} + +func (g *gpuGaudi) SetDeviceType(deviceType DeviceType) { + g.deviceType = deviceType +} + +func (g *gpuGaudi) SetHwType(hwType string) { + g.hwType = hwType +} + +func (g *gpuGaudi) SetTritonInfo(info []TritonGPUInfo) { + g.tritonInfo = info + g.devices = make(map[int]GPUDevice, len(info)) + for _, ti := range info { + g.devices[ti.ID] = GPUDevice{ + ID: ti.ID, + TritonInfo: ti, + } + } +} + +func (g *gpuGaudi) SetSummaries(summaries []DeviceSummary) { + g.summaries = summaries + if g.devices != nil { + for _, summary := range summaries { + var gpuID int + if _, err := fmt.Sscanf(summary.ID, "%d", &gpuID); err == nil { + if dev, exists := g.devices[gpuID]; exists { + dev.Summary = summary + g.devices[gpuID] = dev + } + } + } + } +} + +func gaudiCheck(r *Registry) { + if !utils.HasApp("hl-smi") { + logging.Debug("hl-smi not found, skipping Gaudi detection") + return + } + + gaudiType = GAUDI + if err := addDeviceInterface(r, gaudiType, gaudiHwType, gaudiDeviceStartup); err == nil { + logging.Debugf("Using %s to obtain GPU info", gaudiAccImpl.Name()) + } else { + logging.Debugf("Error registering Gaudi: %v", err) + } +} + +func gaudiDeviceStartup() Device { + a := gaudiAccImpl + if err := a.InitLib(); err != nil { + logging.Debugf("Error initializing %s: %v", gaudiType.String(), err) + return nil + } + if err := a.Init(); err != nil { + logging.Errorf("failed to Init device: %v", err) + return nil + } + logging.Debugf("Using %s to obtain GPU info", gaudiType.String()) + return &a +} + +func (g *gpuGaudi) Name() string { + return gaudiType.String() +} + +func (g *gpuGaudi) DevType() DeviceType { + return gaudiType +} + +func (g *gpuGaudi) HwType() string { + return gaudiHwType +} + +func (g *gpuGaudi) InitLib() error { + return nil +} + +func (g *gpuGaudi) Init() error { + devices, err := queryHLSMI() + if err != nil { + return fmt.Errorf("failed to query hl-smi: %w", err) + } + + logging.Debugf("Detected %d Gaudi device(s)", len(devices)) + g.devices = make(map[int]GPUDevice, len(devices)) + + for _, dev := range devices { + arch := "unknown" + if a, ok := productToArch[dev.Name]; ok { + arch = a + } + + tritonInfo := TritonGPUInfo{ + ID: dev.Index, + Name: dev.Name, + UUID: dev.UUID, + Backend: constants.BackendHPU, + Arch: arch, + WarpSize: 0, + MemoryTotalMB: uint64(dev.MemoryTotal), + } + + summary := DeviceSummary{ + ID: strconv.Itoa(dev.Index), + ProductName: dev.Name, + DriverVersion: dev.DriverVersion, + } + + g.devices[dev.Index] = GPUDevice{ + ID: dev.Index, + TritonInfo: tritonInfo, + Summary: summary, + } + + logging.Debugf("Gaudi device %d: %s arch=%s mem=%d MiB bus=%s", + dev.Index, dev.Name, arch, dev.MemoryTotal, dev.BusID) + } + + return nil +} + +func queryHLSMI() ([]hlsmiDevice, error) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + cmd := exec.CommandContext(ctx, "hl-smi", "-Q", + "index,name,bus_id,driver_version,uuid,module_id,memory.total,memory.free,memory.used", + "-f", "csv") + + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to execute hl-smi: %w", err) + } + + return parseHLSMICSV(string(output)) +} + +// parseHLSMICSV parses CSV output from hl-smi. +// First line is the header, subsequent lines are device data. +// Numeric values may include unit suffixes (e.g., "131072 MiB"). +func parseHLSMICSV(output string) ([]hlsmiDevice, error) { + var devices []hlsmiDevice + + scanner := bufio.NewScanner(strings.NewReader(output)) + + if !scanner.Scan() { + return nil, fmt.Errorf("empty hl-smi output") + } + + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + + fields := strings.Split(line, ", ") + if len(fields) < 9 { + logging.Debugf("Skipping malformed hl-smi line: %s", line) + continue + } + + index, err := strconv.Atoi(strings.TrimSpace(fields[0])) + if err != nil { + logging.Debugf("Failed to parse device index %q: %v", fields[0], err) + continue + } + + moduleID, _ := strconv.Atoi(strings.TrimSpace(fields[5])) + + devices = append(devices, hlsmiDevice{ + Index: index, + Name: strings.TrimSpace(fields[1]), + BusID: strings.TrimSpace(fields[2]), + DriverVersion: strings.TrimSpace(fields[3]), + UUID: strings.TrimSpace(fields[4]), + ModuleID: moduleID, + MemoryTotal: parseMiBValue(fields[6]), + MemoryFree: parseMiBValue(fields[7]), + MemoryUsed: parseMiBValue(fields[8]), + }) + } + + if len(devices) == 0 { + return nil, fmt.Errorf("no devices found in hl-smi output") + } + + return devices, nil +} + +func parseMiBValue(s string) int64 { + s = strings.TrimSpace(s) + s = strings.TrimSuffix(s, " MiB") + s = strings.TrimSuffix(s, " W") + v, _ := strconv.ParseInt(strings.TrimSpace(s), 10, 64) + return v +} + +func (g *gpuGaudi) Shutdown() bool { + return true +} + +func (g *gpuGaudi) GetGPUInfo(gpuID int) (TritonGPUInfo, error) { + dev, exists := g.devices[gpuID] + if !exists { + return TritonGPUInfo{}, fmt.Errorf("GPU device %d not found", gpuID) + } + return dev.TritonInfo, nil +} + +func (g *gpuGaudi) GetAllGPUInfo() ([]TritonGPUInfo, error) { + if len(g.tritonInfo) > 0 { + return g.tritonInfo, nil + } + + var allTritonInfo []TritonGPUInfo + for id := range g.devices { + allTritonInfo = append(allTritonInfo, g.devices[id].TritonInfo) + } + g.tritonInfo = allTritonInfo + return allTritonInfo, nil +} + +func (g *gpuGaudi) GetSummary(gpuID int) (DeviceSummary, error) { + dev, exists := g.devices[gpuID] + if !exists { + return DeviceSummary{}, fmt.Errorf("GPU device %d not found", gpuID) + } + return dev.Summary, nil +} + +func (g *gpuGaudi) GetAllSummaries() ([]DeviceSummary, error) { + if len(g.summaries) > 0 { + return g.summaries, nil + } + + var allSummaries []DeviceSummary + for id := range g.devices { + allSummaries = append(allSummaries, g.devices[id].Summary) + } + g.summaries = allSummaries + return allSummaries, nil +} diff --git a/mcv/pkg/accelerator/devices/gaudi_test.go b/mcv/pkg/accelerator/devices/gaudi_test.go new file mode 100644 index 000000000..7fd20290e --- /dev/null +++ b/mcv/pkg/accelerator/devices/gaudi_test.go @@ -0,0 +1,176 @@ +package devices + +import ( + "testing" +) + +func TestParseHLSMICSV(t *testing.T) { + t.Run("parses full 8-device output", func(t *testing.T) { + csv := `index, name, bus_id, driver_version, uuid, module_id, memory.total [MiB], memory.free [MiB], memory.used [MiB] +0, HL-325L, 0000:3b:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAN071-12-03-01, 1, 131072 MiB, 130400 MiB, 672 MiB +1, HL-325L, 0000:4c:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAN071-05-08-05, 3, 131072 MiB, 130400 MiB, 672 MiB +2, HL-325L, 0000:5d:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAS681-15-03-02, 2, 131072 MiB, 130400 MiB, 672 MiB +3, HL-325L, 0000:9b:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAN071-04-07-03, 6, 131072 MiB, 130400 MiB, 672 MiB +4, HL-325L, 0000:19:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAN071-08-07-01, 0, 131072 MiB, 130400 MiB, 672 MiB +5, HL-325L, 0000:bb:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAH066-20-03-04, 7, 131072 MiB, 130400 MiB, 672 MiB +6, HL-325L, 0000:cb:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAN075-04-06-07, 5, 131072 MiB, 130400 MiB, 672 MiB +7, HL-325L, 0000:db:00.0, 1.24.1-b336d5e, 01P4-HL3090A0-18-UAH079-13-03-03, 4, 131072 MiB, 130400 MiB, 672 MiB +` + devices, err := parseHLSMICSV(csv) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(devices) != 8 { + t.Fatalf("expected 8 devices, got %d", len(devices)) + } + + d := devices[0] + if d.Index != 0 { + t.Errorf("expected index 0, got %d", d.Index) + } + if d.Name != productHL325L { + t.Errorf("expected name HL-325L, got %s", d.Name) + } + if d.BusID != "0000:3b:00.0" { + t.Errorf("expected bus ID 0000:3b:00.0, got %s", d.BusID) + } + if d.DriverVersion != "1.24.1-b336d5e" { + t.Errorf("expected driver 1.24.1-b336d5e, got %s", d.DriverVersion) + } + if d.UUID != "01P4-HL3090A0-18-UAN071-12-03-01" { + t.Errorf("expected UUID 01P4-HL3090A0-18-UAN071-12-03-01, got %s", d.UUID) + } + if d.ModuleID != 1 { + t.Errorf("expected module ID 1, got %d", d.ModuleID) + } + if d.MemoryTotal != 131072 { + t.Errorf("expected memory total 131072, got %d", d.MemoryTotal) + } + if d.MemoryFree != 130400 { + t.Errorf("expected memory free 130400, got %d", d.MemoryFree) + } + if d.MemoryUsed != 672 { + t.Errorf("expected memory used 672, got %d", d.MemoryUsed) + } + + last := devices[7] + if last.Index != 7 { + t.Errorf("expected last device index 7, got %d", last.Index) + } + if last.BusID != "0000:db:00.0" { + t.Errorf("expected last bus ID 0000:db:00.0, got %s", last.BusID) + } + }) + + t.Run("parses single device", func(t *testing.T) { + csv := `index, name, bus_id, driver_version, uuid, module_id, memory.total [MiB], memory.free [MiB], memory.used [MiB] +0, HL-225, 0000:3b:00.0, 1.20.0-abc123, some-uuid-here, 0, 98304 MiB, 97000 MiB, 1304 MiB +` + devices, err := parseHLSMICSV(csv) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(devices) != 1 { + t.Fatalf("expected 1 device, got %d", len(devices)) + } + if devices[0].Name != productHL225 { + t.Errorf("expected name HL-225, got %s", devices[0].Name) + } + if devices[0].MemoryTotal != 98304 { + t.Errorf("expected memory total 98304, got %d", devices[0].MemoryTotal) + } + }) + + t.Run("returns error on empty output", func(t *testing.T) { + _, err := parseHLSMICSV("") + if err == nil { + t.Fatal("expected error for empty output") + } + }) + + t.Run("returns error on header only", func(t *testing.T) { + csv := `index, name, bus_id, driver_version, uuid, module_id, memory.total [MiB], memory.free [MiB], memory.used [MiB] +` + _, err := parseHLSMICSV(csv) + if err == nil { + t.Fatal("expected error for header-only output") + } + }) + + t.Run("skips malformed lines", func(t *testing.T) { + csv := `index, name, bus_id, driver_version, uuid, module_id, memory.total [MiB], memory.free [MiB], memory.used [MiB] +not-a-number, HL-325L, 0000:3b:00.0, 1.24.1, uuid1, 0, 131072 MiB, 130400 MiB, 672 MiB +0, HL-325L, 0000:4c:00.0, 1.24.1, uuid2, 1, 131072 MiB, 130400 MiB, 672 MiB +too, few, fields +1, HL-325L, 0000:5d:00.0, 1.24.1, uuid3, 2, 131072 MiB, 130400 MiB, 672 MiB +` + devices, err := parseHLSMICSV(csv) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(devices) != 2 { + t.Fatalf("expected 2 valid devices, got %d", len(devices)) + } + if devices[0].Index != 0 { + t.Errorf("expected first valid device index 0, got %d", devices[0].Index) + } + if devices[1].Index != 1 { + t.Errorf("expected second valid device index 1, got %d", devices[1].Index) + } + }) +} + +func TestParseMiBValue(t *testing.T) { + tests := []struct { + input string + expected int64 + }{ + {"131072 MiB", 131072}, + {"130400 MiB", 130400}, + {"672 MiB", 672}, + {"239 W", 239}, + {"0 MiB", 0}, + {"131072", 131072}, + {" 131072 MiB ", 131072}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := parseMiBValue(tt.input) + if got != tt.expected { + t.Errorf("parseMiBValue(%q) = %d, want %d", tt.input, got, tt.expected) + } + }) + } +} + +func TestProductToArch(t *testing.T) { + tests := []struct { + product string + expected string + }{ + {productHL325L, archGaudi3}, + {productHL325, archGaudi3}, + {productHL225, archGaudi2}, + {productHL225B, archGaudi2}, + } + + for _, tt := range tests { + t.Run(tt.product, func(t *testing.T) { + got, ok := productToArch[tt.product] + if !ok { + t.Fatalf("product %q not found in productToArch", tt.product) + } + if got != tt.expected { + t.Errorf("productToArch[%q] = %q, want %q", tt.product, got, tt.expected) + } + }) + } + + t.Run("unknown product returns false", func(t *testing.T) { + _, ok := productToArch["UNKNOWN-DEVICE"] + if ok { + t.Error("expected unknown product to not be in map") + } + }) +} diff --git a/mcv/pkg/accelerator/devices/nvml.go b/mcv/pkg/accelerator/devices/nvml.go index 7524f910a..c737749ab 100644 --- a/mcv/pkg/accelerator/devices/nvml.go +++ b/mcv/pkg/accelerator/devices/nvml.go @@ -24,6 +24,7 @@ import ( logging "github.com/sirupsen/logrus" "github.com/redhat-et/GKM/mcv/pkg/config" + "github.com/redhat-et/GKM/mcv/pkg/constants" ) const ( @@ -253,7 +254,7 @@ func getNVMLTritonGPUInfo(device nvml.Device) (TritonGPUInfo, error) { WarpSize: warpSize, MemoryTotalMB: mem.Total / (1024 * 1024), PTXVersion: ptxVersion, - Backend: "cuda", + Backend: constants.BackendCUDA, }, nil } diff --git a/mcv/pkg/accelerator/devices/rocm.go b/mcv/pkg/accelerator/devices/rocm.go index c5658520b..70f590e34 100644 --- a/mcv/pkg/accelerator/devices/rocm.go +++ b/mcv/pkg/accelerator/devices/rocm.go @@ -10,6 +10,7 @@ import ( "time" "github.com/redhat-et/GKM/mcv/pkg/config" + "github.com/redhat-et/GKM/mcv/pkg/constants" "github.com/redhat-et/GKM/mcv/pkg/utils" logging "github.com/sirupsen/logrus" ) @@ -187,7 +188,7 @@ func (r *gpuROCm) Init() error { Arch: info.GFXVersion, WarpSize: 64, MemoryTotalMB: memTotal / (1024 * 1024), - Backend: hipBackend, + Backend: constants.BackendHIP, ID: gpuID, }, Summary: DeviceSummary{ diff --git a/mcv/pkg/accelerator/devices/static.go b/mcv/pkg/accelerator/devices/static.go index d007adc1a..0edbc4429 100644 --- a/mcv/pkg/accelerator/devices/static.go +++ b/mcv/pkg/accelerator/devices/static.go @@ -77,49 +77,5 @@ func staticDeviceStartup() Device { } func NewStubbedDeviceCache() *DeviceCache { - return &DeviceCache{ - Devices: map[string]CachedDevice{ - "gpu": { - Name: stubbedAMDName, - DeviceType: 1, // DeviceType for GPU, adjust if you have a constant - HwType: "gpu", - TritonInfo: []TritonGPUInfo{ - { - Name: "card0", - UUID: "daff740f-0000-1000-8062-0165038984ec", - ComputeCapability: "", - Arch: gfxArchMI210, - WarpSize: 64, - MemoryTotalMB: 65520, - PTXVersion: 0, - Backend: hipBackend, - ID: 0, - }, - { - Name: "card1", - UUID: "acff740f-0000-1000-806b-c6ef57f28db1", - ComputeCapability: "", - Arch: gfxArchMI210, - WarpSize: 64, - MemoryTotalMB: 65520, - PTXVersion: 0, - Backend: hipBackend, - ID: 1, - }, - }, - Summaries: []DeviceSummary{ - { - ID: "0", - DriverVersion: "6.12.10-100.fc40.x86_64", - ProductName: "STUBBED Aldebaran/MI200 [Instinct MI210]", - }, - { - ID: "1", - DriverVersion: "6.12.10-100.fc40.x86_64", - ProductName: "STUBBED Aldebaran/MI200 [Instinct MI210]", - }, - }, - }, - }, - } + return stubbedDeviceCache(activeStubProfile()) } diff --git a/mcv/pkg/accelerator/devices/stub.go b/mcv/pkg/accelerator/devices/stub.go new file mode 100644 index 000000000..551870923 --- /dev/null +++ b/mcv/pkg/accelerator/devices/stub.go @@ -0,0 +1,163 @@ +package devices + +import ( + "fmt" + "os" + + "github.com/jaypipes/ghw" + "github.com/jaypipes/ghw/pkg/accelerator" + "github.com/jaypipes/ghw/pkg/pci" + "github.com/jaypipes/pcidb" + logging "github.com/sirupsen/logrus" + + "github.com/redhat-et/GKM/mcv/pkg/config" + "github.com/redhat-et/GKM/mcv/pkg/constants" +) + +// EnvStubProfile selects which GPU type the stub simulates. +// Set to "amd" (default), "gaudi", or "nvidia". +// Example: MCV_STUB_PROFILE=amd ./mcv +const EnvStubProfile = "MCV_STUB_PROFILE" + +// stubCardSpec holds per-card data for a single simulated accelerator. +type stubCardSpec struct { + Name string + UUID string + Arch string + WarpSize int + MemoryMB uint64 + Backend string +} + +// stubProfile groups all data needed to simulate a GPU in stub mode. +// Add a new profile here to support an additional GPU type; no other +// files need to change. +type stubProfile struct { + DeviceName string + DeviceType DeviceType + HwType string + DriverVersion string + PciVendorName string + PciVendorID string + Cards []stubCardSpec +} + +// Pre-built profiles — one per GPU family. Add entries here to extend. +var ( + gaudiStubProfile = &stubProfile{ + DeviceName: "STUBBED Habana", + DeviceType: GAUDI, + HwType: config.GPU, + DriverVersion: "1.24.1-b336d5e", + PciVendorName: "STUBBED Habana", + PciVendorID: "1da3", // Habana Labs (Intel) PCI vendor ID + Cards: []stubCardSpec{ + {Name: productHL325L, UUID: "01P4-HL3090A0-18-UAN071-12-03-01", Arch: archGaudi3, WarpSize: 0, MemoryMB: 131072, Backend: constants.BackendHPU}, + {Name: productHL325L, UUID: "01P4-HL3090A0-18-UAN071-05-08-05", Arch: archGaudi3, WarpSize: 0, MemoryMB: 131072, Backend: constants.BackendHPU}, + }, + } + + amdStubProfile = &stubProfile{ + DeviceName: "STUBBED AMD", + DeviceType: AMD, + HwType: config.GPU, + DriverVersion: "6.12.10-100.fc40.x86_64", + PciVendorName: "STUBBED AMD", + PciVendorID: "1002", // AMD PCI vendor ID + Cards: []stubCardSpec{ + {Name: "card0", UUID: "daff740f-0000-1000-8062-0165038984ec", Arch: gfxArchMI210, WarpSize: 64, MemoryMB: 65520, Backend: constants.BackendHIP}, + {Name: "card1", UUID: "acff740f-0000-1000-806b-c6ef57f28db1", Arch: gfxArchMI210, WarpSize: 64, MemoryMB: 65520, Backend: constants.BackendHIP}, + }, + } + + // nvidiaStubProfile is a placeholder — fill in real values when needed. + nvidiaStubProfile = &stubProfile{ + DeviceName: "STUBBED NVIDIA", + DeviceType: NVML, + HwType: config.GPU, + DriverVersion: "550.54.15", + PciVendorName: "STUBBED NVIDIA", + PciVendorID: "10de", // NVIDIA PCI vendor ID + Cards: []stubCardSpec{ + {Name: "Tesla A100-SXM4-80GB", UUID: "GPU-00000000-0000-0000-0000-000000000001", Arch: "sm_80", WarpSize: 32, MemoryMB: 81920, Backend: constants.BackendCUDA}, + {Name: "Tesla A100-SXM4-80GB", UUID: "GPU-00000000-0000-0000-0000-000000000002", Arch: "sm_80", WarpSize: 32, MemoryMB: 81920, Backend: constants.BackendCUDA}, + }, + } +) + +// activeStubProfile returns the stub profile indicated by MCV_STUB_PROFILE. +// Defaults to Gaudi when the variable is unset. Logs a warning and falls back +// to Gaudi for any unrecognized value. +func activeStubProfile() *stubProfile { + val := os.Getenv(EnvStubProfile) + switch val { + case "", "amd": + return amdStubProfile + case "gaudi": + return gaudiStubProfile + case "nvidia": + return nvidiaStubProfile + default: + logging.Warnf("Unknown MCV_STUB_PROFILE=%q; falling back to gaudi. Valid values: gaudi, amd, nvidia", val) + return gaudiStubProfile + } +} + +// stubbedDeviceCache builds a DeviceCache from a stub profile. +func stubbedDeviceCache(p *stubProfile) *DeviceCache { + tritonInfo := make([]TritonGPUInfo, len(p.Cards)) + summaries := make([]DeviceSummary, len(p.Cards)) + for i, c := range p.Cards { + tritonInfo[i] = TritonGPUInfo{ + Name: c.Name, + UUID: c.UUID, + Arch: c.Arch, + WarpSize: c.WarpSize, + MemoryTotalMB: c.MemoryMB, + Backend: c.Backend, + ID: i, + } + summaries[i] = DeviceSummary{ + ID: fmt.Sprintf("%d", i), + DriverVersion: p.DriverVersion, + ProductName: c.Name, + } + } + return &DeviceCache{ + Devices: map[string]CachedDevice{ + config.GPU: { + Name: p.DeviceName, + DeviceType: p.DeviceType, + HwType: p.HwType, + TritonInfo: tritonInfo, + Summaries: summaries, + }, + }, + } +} + +// stubbedAcceleratorInfo builds a ghw.AcceleratorInfo from a stub profile. +func stubbedAcceleratorInfo(p *stubProfile) *ghw.AcceleratorInfo { + devices := make([]*accelerator.AcceleratorDevice, len(p.Cards)) + for i, c := range p.Cards { + devices[i] = &accelerator.AcceleratorDevice{ + Address: fmt.Sprintf("0000:00:%02x.0", i+1), + PCIDevice: &pci.Device{ + Vendor: &pcidb.Vendor{ + Name: p.PciVendorName, + ID: p.PciVendorID, + }, + Product: &pcidb.Product{ + Name: "STUBBED " + c.Name, + ID: "STUBBED " + c.Name, + }, + Driver: "dummy", + Class: &pcidb.Class{ + Name: "controller", + ID: "0300", + }, + }, + } + } + return &ghw.AcceleratorInfo{Devices: devices} +} diff --git a/mcv/pkg/accelerator/devices/utils.go b/mcv/pkg/accelerator/devices/utils.go index 20b9b85ee..4dd7cc2e2 100644 --- a/mcv/pkg/accelerator/devices/utils.go +++ b/mcv/pkg/accelerator/devices/utils.go @@ -4,9 +4,6 @@ import ( "fmt" "github.com/jaypipes/ghw" - "github.com/jaypipes/ghw/pkg/accelerator" - "github.com/jaypipes/ghw/pkg/pci" - "github.com/jaypipes/pcidb" "github.com/redhat-et/GKM/mcv/pkg/config" logging "github.com/sirupsen/logrus" ) @@ -26,52 +23,12 @@ func GetProductName(id int) (name string, err error) { } // DetectAccelerators detects hardware accelerators and enables GPU logic if supported hardware is found. -// If stub mode is enabled, it simulates the presence of an AMD Aldebaran MI200 GPU. +// If stub mode is enabled, it simulates the accelerator selected by MCV_STUB_PROFILE (default: amd). // If no hardware accelerators are found, it returns nil without an error. func DetectAccelerators() (accInfo *ghw.AcceleratorInfo) { if config.IsStubEnabled() { logging.Debug("Stub mode configured, simulating accelerator device") - accInfo = &ghw.AcceleratorInfo{ - Devices: []*accelerator.AcceleratorDevice{ - { - Address: "0000:00:01.0", - PCIDevice: &pci.Device{ - Vendor: &pcidb.Vendor{ - Name: stubbedAMDName, - ID: "1002", - }, - Product: &pcidb.Product{ - Name: stubbedAMDName, - ID: "STUBBED Aldebaran/MI200", - }, - Driver: "dummy", - Class: &pcidb.Class{ - Name: "controller", - ID: "0300", - }, - }, - }, - { - Address: "0000:00:02.0", - PCIDevice: &pci.Device{ - Vendor: &pcidb.Vendor{ - Name: stubbedAMDName, - ID: "1002", - }, - Product: &pcidb.Product{ - Name: "STUBBED Product", - ID: "STUBBED Aldebaran/MI200", - }, - Driver: "dummy", - Class: &pcidb.Class{ - Name: "controller", - ID: "0300", - }, - }, - }, - }, - } - return accInfo + return stubbedAcceleratorInfo(activeStubProfile()) } acc, err := ghw.Accelerator() diff --git a/mcv/pkg/cache/cache.go b/mcv/pkg/cache/cache.go index 08d779b64..3595f19ce 100644 --- a/mcv/pkg/cache/cache.go +++ b/mcv/pkg/cache/cache.go @@ -39,6 +39,8 @@ func DetectCaches(root string) []Cache { caches = append(caches, vllm) } else if triton := DetectTritonCache(root); triton != nil { caches = append(caches, triton) + } else if habana := DetectHabanaCache(root); habana != nil { + caches = append(caches, habana) } return caches @@ -146,7 +148,7 @@ func CacheTypes(caches []Cache) []string { // GetTagsFromCaches returns the manifest and cache directory tags for the available cache type func GetTagsFromCaches(caches []Cache) (manifestTag, cacheTag string, err error) { for _, c := range caches { - if c.Name() == constants.VLLM || c.Name() == constants.Triton { + if c.Name() == constants.VLLM || c.Name() == constants.Triton || c.Name() == constants.Habana { return c.ManifestTag(), c.CacheTag(), nil } } @@ -171,6 +173,8 @@ func ExtractCacheDirectory(r io.Reader, cacheType string) (extractedDirs []strin return ExtractTritonCacheDirectory(r) case constants.VLLM: return ExtractVLLMCacheDirectory(r) + case constants.Habana: + return ExtractHabanaCacheDirectory(r) default: return nil, 0, fmt.Errorf("unsupported cache type: %s", cacheType) } diff --git a/mcv/pkg/cache/habana.go b/mcv/pkg/cache/habana.go new file mode 100644 index 000000000..79150f13b --- /dev/null +++ b/mcv/pkg/cache/habana.go @@ -0,0 +1,232 @@ +package cache + +import ( + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + + "github.com/redhat-et/GKM/mcv/pkg/cacheplan" + "github.com/redhat-et/GKM/mcv/pkg/constants" + logging "github.com/sirupsen/logrus" +) + +// habanaRecipeCacheSizeMB returns the cache size limit to stamp into OCI labels. +// It reads EnvHabanaRecipeCacheSizeMB (injected by the KServe sidecar injector) +// so the OCI label matches exactly what the predictor's PT_HPU_RECIPE_CACHE_CONFIG +// was set to at capture time. +func habanaRecipeCacheSizeMB() int { + if val := os.Getenv(cacheplan.EnvHabanaRecipeCacheSizeMB); val != "" { + if n, err := strconv.Atoi(val); err == nil && n > 0 { + return n + } + logging.Warnf("Invalid %s=%q, using default %d MB", cacheplan.EnvHabanaRecipeCacheSizeMB, val, cacheplan.DefaultHabanaRecipeCacheSizeMB) + } + return cacheplan.DefaultHabanaRecipeCacheSizeMB +} + +const ( + cacheHabanaImagePrefix = "cache.habana.image" + cacheHabanaImageEntryCount = cacheHabanaImagePrefix + "/entry-count" + cacheHabanaImageCacheSize = cacheHabanaImagePrefix + "/cache-size-bytes" + cacheHabanaImageSummary = cacheHabanaImagePrefix + "/summary" + + habanaBackend = constants.BackendHPU + + // habanaCacheRootPath is the directory a KServe serving container mounts the + // Habana recipe cache at. It is baked into the cache-root-env label so the + // consumer knows both the env var and the mount location. + habanaCacheRootPath = constants.KServeHome + "/" + constants.HabanaCache +) + +// recipeFileRegex matches Habana recipe cache filenames: +// +// {graph_hash}_{device_fingerprint}_syn{synapse_version}.recipe +// {graph_hash}_{device_fingerprint}_syn{synapse_version}.metadata +var recipeFileRegex = regexp.MustCompile( + `^(\d+)_([a-f0-9]+)_syn([\d.]+[a-f0-9]*)\.recipe$`, +) + +// HabanaCache represents a Habana Synapse recipe cache directory. +type HabanaCache struct { + rootPath string + tmpPath string + allMetadata []HabanaRecipeMetadata +} + +// HabanaRecipeMetadata holds parsed info from a single recipe file pair. +type HabanaRecipeMetadata struct { + GraphHash string `json:"graphHash"` + DeviceID string `json:"deviceId"` + SynapseVersion string `json:"synapseVersion"` + RecipeSize int64 `json:"recipeSize"` + MetadataSize int64 `json:"metadataSize"` +} + +// DetectHabanaCache scans cacheDir for Habana recipe files. +// Returns nil if none are found. +func DetectHabanaCache(cacheDir string) *HabanaCache { + entries, err := os.ReadDir(cacheDir) + if err != nil { + logging.WithError(err).Debugf("Cannot read directory: %s", cacheDir) + return nil + } + + var metadata []HabanaRecipeMetadata + for _, entry := range entries { + if entry.IsDir() { + continue + } + m := recipeFileRegex.FindStringSubmatch(entry.Name()) + if m == nil { + continue + } + + recipeInfo, err := entry.Info() + if err != nil { + logging.WithError(err).Warnf("Failed to stat recipe file: %s", entry.Name()) + continue + } + + // Look for the matching .metadata file + metaName := strings.TrimSuffix(entry.Name(), ".recipe") + ".metadata" + var metaSize int64 + if metaInfo, err := os.Stat(filepath.Join(cacheDir, metaName)); err == nil { + metaSize = metaInfo.Size() + } + + metadata = append(metadata, HabanaRecipeMetadata{ + GraphHash: m[1], + DeviceID: m[2], + SynapseVersion: m[3], + RecipeSize: recipeInfo.Size(), + MetadataSize: metaSize, + }) + } + + if len(metadata) == 0 { + logging.Debugf("No Habana recipe cache found in: %s", cacheDir) + return nil + } + + logging.Infof("Detected Habana recipe cache: %d recipes in %s", len(metadata), cacheDir) + return &HabanaCache{ + rootPath: cacheDir, + allMetadata: metadata, + } +} + +func (h *HabanaCache) Name() string { return constants.Habana } + +func (h *HabanaCache) EntryCount() int { return len(h.allMetadata) } + +func (h *HabanaCache) CacheSizeBytes() int64 { + dir := h.rootPath + if h.tmpPath != "" { + dir = h.tmpPath + } + size, _ := getTotalDirSize(dir) + return size +} + +func (h *HabanaCache) Summary() string { + summary, err := buildHabanaSummary(h.allMetadata) + if err != nil { + logging.WithError(err).Error("failed to build Habana summary") + return "" + } + data, err := json.Marshal(summary) + if err != nil { + logging.WithError(err).Error("failed to marshal Habana summary") + return "" + } + return string(data) +} + +func (h *HabanaCache) Metadata() []CacheEntry { + entries := make([]CacheEntry, 0, len(h.allMetadata)) + for _, m := range h.allMetadata { + entries = append(entries, m) + } + return entries +} + +func (h *HabanaCache) Labels() map[string]string { + // Build the cache-root-env label from the shared cacheplan helper so the + // label MCV stamps here is byte-identical to the env cacheplan derives on the + // consume side (PT_HPU_RECIPE_CACHE_CONFIG=,false,). + // The size is read from EnvHabanaRecipeCacheSizeMB, injected by the KServe + // sidecar injector so it matches whatever was set in the predictor container. + sizeMB := habanaRecipeCacheSizeMB() + rootEnv, err := cacheplan.RootEnvLabel(constants.Habana, habanaCacheRootPath, sizeMB) + if err != nil { + // This cannot happen for a known cache type, but fall back to the bare + // env name rather than emitting an empty label. + logging.WithError(err).Error("failed to build Habana cache-root-env label") + rootEnv = constants.HabanaRecipeCacheEnv + "=" + habanaCacheRootPath + } + + labels := map[string]string{ + cacheHabanaImageEntryCount: strconv.Itoa(h.EntryCount()), + cacheHabanaImageCacheSize: strconv.FormatInt(h.CacheSizeBytes(), 10), + cacheHabanaImageSummary: h.Summary(), + + cacheplan.LabelFramework: constants.VLLM, + cacheplan.LabelCacheType: constants.CacheTypeHabanaRecipe, + cacheplan.LabelCacheRootEnv: rootEnv, + cacheplan.LabelCacheMountSubpath: ".", + } + return labels +} + +func (h *HabanaCache) ManifestTag() string { + return fmt.Sprintf("./%s", constants.MCVHabanaManifestDir) +} + +func (h *HabanaCache) CacheTag() string { + return fmt.Sprintf("./%s", constants.MCVHabanaCacheDir) +} + +func (h *HabanaCache) SetTmpPath(path string) { + if path != "" { + h.tmpPath = path + } +} + +func buildHabanaSummary(metadata []HabanaRecipeMetadata) (*Summary, error) { + if len(metadata) == 0 { + return nil, fmt.Errorf("no Habana metadata to summarize") + } + + // Deduplicate by device fingerprint — each unique device ID + // maps to an architecture via the Gaudi device layer (gaudi2, gaudi3). + // We don't know the arch from recipe filenames alone, so we leave + // arch empty here; the preflight check matches on backend only. + seen := make(map[string]bool) + var targets []SummaryTargetInfo + for _, m := range metadata { + if seen[m.DeviceID] { + continue + } + seen[m.DeviceID] = true + targets = append(targets, SummaryTargetInfo{ + Backend: habanaBackend, + WarpSize: 0, + }) + } + return &Summary{Targets: targets}, nil +} + +func ExtractHabanaCacheDirectory(r io.Reader) (dirs []string, bytesWritten int64, err error) { + return extractCacheAndManifestDirectory( + r, + constants.MCVHabanaCacheDir, + "io.habana.manifest/", + constants.ExtractCacheDir, + constants.ExtractManifestDir, + ) +} diff --git a/mcv/pkg/cache/habana_test.go b/mcv/pkg/cache/habana_test.go new file mode 100644 index 000000000..71bd58e59 --- /dev/null +++ b/mcv/pkg/cache/habana_test.go @@ -0,0 +1,213 @@ +package cache + +import ( + "os" + "path/filepath" + "testing" +) + +const ( + testDeviceID = "a16fb581eee" + testSynapseVersion = "1.24.1.6210fda" + testSynapseShort = "1.24.0" +) + +func TestRecipeFileRegex(t *testing.T) { + tests := []struct { + name string + filename string + wantHash string + wantDev string + wantVer string + wantOK bool + }{ + { + name: "valid recipe file", + filename: "10034702590219030139_a16fb581eee_syn1.24.1.6210fda.recipe", + wantHash: "10034702590219030139", + wantDev: testDeviceID, + wantVer: testSynapseVersion, + wantOK: true, + }, + { + name: "metadata file does not match recipe regex", + filename: "10034702590219030139_a16fb581eee_syn1.24.1.6210fda.metadata", + wantOK: false, + }, + { + name: "debug dir does not match", + filename: "10034702590219030139_a16fb581eee_syn1.24.1.6210fda.recipe_debug_files", + wantOK: false, + }, + { + name: "random file", + filename: "README.md", + wantOK: false, + }, + { + name: "short hash", + filename: "123_abc_syn2.0.recipe", + wantHash: "123", + wantDev: "abc", + wantVer: "2.0", + wantOK: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + m := recipeFileRegex.FindStringSubmatch(tt.filename) + if !tt.wantOK { + if m != nil { + t.Errorf("expected no match for %q, got %v", tt.filename, m) + } + return + } + if m == nil { + t.Fatalf("expected match for %q, got nil", tt.filename) + } + if m[1] != tt.wantHash { + t.Errorf("hash: got %q, want %q", m[1], tt.wantHash) + } + if m[2] != tt.wantDev { + t.Errorf("device: got %q, want %q", m[2], tt.wantDev) + } + if m[3] != tt.wantVer { + t.Errorf("version: got %q, want %q", m[3], tt.wantVer) + } + }) + } +} + +func TestDetectHabanaCache(t *testing.T) { + t.Run("detects valid cache", func(t *testing.T) { + dir := t.TempDir() + // Create two recipe triplets + files := []struct { + name string + size int + }{ + {"111_aaa_syn" + testSynapseShort + ".recipe", 1024}, + {"111_aaa_syn" + testSynapseShort + ".metadata", 128}, + {"222_aaa_syn" + testSynapseShort + ".recipe", 2048}, + {"222_aaa_syn" + testSynapseShort + ".metadata", 256}, + } + for _, f := range files { + if err := os.WriteFile(filepath.Join(dir, f.name), make([]byte, f.size), 0644); err != nil { + t.Fatal(err) + } + } + // Also create a debug dir (should be ignored) + os.Mkdir(filepath.Join(dir, "111_aaa_syn"+testSynapseShort+".recipe_debug_files"), 0755) + + cache := DetectHabanaCache(dir) + if cache == nil { + t.Fatal("expected cache to be detected") + } + if cache.EntryCount() != 2 { + t.Errorf("entry count: got %d, want 2", cache.EntryCount()) + } + if cache.Name() != "habana" { + t.Errorf("name: got %q, want %q", cache.Name(), "habana") + } + // Verify metadata + for _, m := range cache.allMetadata { + if m.DeviceID != "aaa" { + t.Errorf("deviceID: got %q, want %q", m.DeviceID, "aaa") + } + if m.SynapseVersion != testSynapseShort { + t.Errorf("version: got %q, want %q", m.SynapseVersion, testSynapseShort) + } + } + }) + + t.Run("returns nil for empty dir", func(t *testing.T) { + dir := t.TempDir() + cache := DetectHabanaCache(dir) + if cache != nil { + t.Error("expected nil for empty directory") + } + }) + + t.Run("returns nil for non-habana files", func(t *testing.T) { + dir := t.TempDir() + os.WriteFile(filepath.Join(dir, "somefile.txt"), []byte("hello"), 0644) + cache := DetectHabanaCache(dir) + if cache != nil { + t.Error("expected nil for non-habana files") + } + }) + + t.Run("returns nil for nonexistent dir", func(t *testing.T) { + cache := DetectHabanaCache("/nonexistent/path") + if cache != nil { + t.Error("expected nil for nonexistent directory") + } + }) +} + +func TestHabanaLabels(t *testing.T) { + h := &HabanaCache{ + rootPath: t.TempDir(), + allMetadata: []HabanaRecipeMetadata{ + {GraphHash: "111", DeviceID: testDeviceID, SynapseVersion: testSynapseVersion, RecipeSize: 1024}, + }, + } + + labels := h.Labels() + + // The cache-root-env label must carry the full PT_HPU_RECIPE_CACHE_CONFIG + // value (name, dir, and the ",false,8192" tunables), not just the env name. + want := "PT_HPU_RECIPE_CACHE_CONFIG=/home/kserve/.cache/habana,false,8192" + if got := labels["io.kserve.km/cache-root-env"]; got != want { + t.Errorf("cache-root-env: got %q, want %q", got, want) + } + if got := labels["io.kserve.km/cache-type"]; got != "habana-recipe" { + t.Errorf("cache-type: got %q, want %q", got, "habana-recipe") + } + if got := labels["io.kserve.km/cache-mount-subpath"]; got != "." { + t.Errorf("cache-mount-subpath: got %q, want %q", got, ".") + } +} + +func TestBuildHabanaSummary(t *testing.T) { + metadata := []HabanaRecipeMetadata{ + {GraphHash: "111", DeviceID: testDeviceID, SynapseVersion: testSynapseVersion, RecipeSize: 1024}, + {GraphHash: "222", DeviceID: testDeviceID, SynapseVersion: testSynapseVersion, RecipeSize: 2048}, + } + + summary, err := buildHabanaSummary(metadata) + if err != nil { + t.Fatal(err) + } + if len(summary.Targets) != 1 { + t.Fatalf("targets: got %d, want 1", len(summary.Targets)) + } + if summary.Targets[0].Backend != "hpu" { + t.Errorf("backend: got %q, want %q", summary.Targets[0].Backend, "hpu") + } + if summary.Targets[0].WarpSize != 0 { + t.Errorf("warpSize: got %d, want 0", summary.Targets[0].WarpSize) + } + + t.Run("empty metadata returns error", func(t *testing.T) { + _, err := buildHabanaSummary(nil) + if err == nil { + t.Error("expected error for nil metadata") + } + }) + + t.Run("multiple device IDs produce multiple targets", func(t *testing.T) { + meta := []HabanaRecipeMetadata{ + {GraphHash: "111", DeviceID: "aaa", SynapseVersion: testSynapseShort}, + {GraphHash: "222", DeviceID: "bbb", SynapseVersion: testSynapseShort}, + } + s, err := buildHabanaSummary(meta) + if err != nil { + t.Fatal(err) + } + if len(s.Targets) != 2 { + t.Errorf("targets: got %d, want 2", len(s.Targets)) + } + }) +} diff --git a/mcv/pkg/cache/triton.go b/mcv/pkg/cache/triton.go index d1b49b0fc..d11488794 100644 --- a/mcv/pkg/cache/triton.go +++ b/mcv/pkg/cache/triton.go @@ -229,7 +229,7 @@ func CompareTritonCacheToGPU(cacheData *TritonCacheData, acc accelerator.Acceler warpMatch := cacheData.Target.WarpSize == gpu.WarpSize ptxMatch := true - if gpu.Backend == "cuda" && cacheData.PtxVersion != nil { + if gpu.Backend == constants.BackendCUDA && cacheData.PtxVersion != nil { ptxMatch = *cacheData.PtxVersion == gpu.PTXVersion if !ptxMatch { logging.WithFields(logging.Fields{ diff --git a/mcv/pkg/cache/vllm.go b/mcv/pkg/cache/vllm.go index 9ffb10452..0b8e54d13 100644 --- a/mcv/pkg/cache/vllm.go +++ b/mcv/pkg/cache/vllm.go @@ -41,10 +41,12 @@ const ( BinaryCacheFormat = "binary" AOTCompileCacheFormat = "aot_compile" TritonCacheFormat = "triton" - CUDABackend = "cuda" - ROCmBackend = "rocm" - HIPBackend = "hip" - UnknownBackend = "UnknownBackend" + // Exported aliases of constants.Backend* — kept for callers that already + // reference cache.CUDABackend / cache.HIPBackend / cache.ROCmBackend. + CUDABackend = constants.BackendCUDA + ROCmBackend = constants.BackendROCm + HIPBackend = constants.BackendHIP + UnknownBackend = "UnknownBackend" // torchAOTCompileDirName is the extra directory vLLM introduces above // the per-model hash dir when VLLM_USE_AOT_COMPILE is enabled. diff --git a/mcv/pkg/cacheplan/cacheplan.go b/mcv/pkg/cacheplan/cacheplan.go new file mode 100644 index 000000000..edc29570b --- /dev/null +++ b/mcv/pkg/cacheplan/cacheplan.go @@ -0,0 +1,304 @@ +// Package cacheplan is the single source of truth for how a kernel cache is +// interpreted on both sides of the Kernel Manager (KM) flow: +// +// - Producer (capture): given a framework/cache type and the directory the +// framework should write to, ProducerEnv returns the environment variables +// that make the framework populate that directory. Some drivers (notably the +// Habana/Synapse recipe cache) only write to disk when their env var is set, +// so the capture side must inject this up front. +// +// - Consumer (serving): given the OCI image labels MCV stamped at build time, +// Derive returns a typed CachePlan describing the env vars to set, the bare +// directory to mount at, the payload subpath/prefix within the image, and +// whether the cache needs a writable copy. +// +// The package is pure: it performs no I/O and depends on no Kubernetes types. +// It takes already-fetched labels and returns plain data, so it can be shared +// verbatim between MCV (which authors the labels) and KServe (which actuates the +// resulting plan onto a pod spec). The OCI label schema remains the persisted +// compatibility contract between the two repositories. +package cacheplan + +import ( + "errors" + "fmt" + "strings" + + "github.com/redhat-et/GKM/mcv/pkg/constants" +) + +// Generic KServe Kernel Manager labels stamped by each cache class's Labels() +// method. These are the compatibility contract between MCV and KServe. +const ( + LabelCacheType = constants.KMPrefix + "/cache-type" + LabelCacheRootEnv = constants.KMPrefix + "/cache-root-env" + LabelCacheMountSubpath = constants.KMPrefix + "/cache-mount-subpath" + LabelCacheHash = constants.KMPrefix + "/cache-hash" + LabelFramework = constants.KMPrefix + "/framework" +) + +// Per-class summary labels, used to infer the cache type for older images that +// predate the io.kserve.km/cache-type label. +const ( + summaryLabelTriton = "cache.triton.image/summary" + summaryLabelVLLM = "cache.vllm.image/summary" + summaryLabelHabana = "cache.habana.image/summary" +) + +// Habana/Synapse recipe cache tunables encoded into PT_HPU_RECIPE_CACHE_CONFIG +// (format: ",,"). delete=false means populate-if-empty / +// reuse, which lets a pre-seeded cache be reused instead of rebuilt and cuts +// warmup time substantially. +const ( + habanaRecipeDelete = "false" + + // DefaultHabanaRecipeCacheSizeMB is the max disk the Synapse driver will use + // for the recipe cache (8 GB). Callers that read this from KernelCacheCapture.Spec + // should use that value instead and fall back to this default when unset. + DefaultHabanaRecipeCacheSizeMB = 8192 + + // EnvHabanaRecipeCacheSizeMB is the env var the KServe sidecar injector sets on + // the MCV capture container so that habana.Labels() stamps the same size into the + // OCI image label as was set in the predictor's PT_HPU_RECIPE_CACHE_CONFIG. + EnvHabanaRecipeCacheSizeMB = "HABANA_RECIPE_CACHE_SIZE_MB" +) + +// EnvVar is a name/value environment variable pair. It deliberately mirrors the +// shape of corev1.EnvVar without depending on the Kubernetes API. +type EnvVar struct { + Name string + Value string +} + +// CachePlan is the typed interpretation of a kernel cache image, derived from +// its OCI labels. KServe consumes these fields to build the serving pod spec. +type CachePlan struct { + // CacheType is the canonical cache type identifier + // (constants.CacheTypeVLLMTorchCompile, constants.CacheTypeHabanaRecipe). + CacheType string + + // Env is the set of environment variables to set on the serving container + // verbatim (e.g. PT_HPU_RECIPE_CACHE_CONFIG=/dir,false,8192). + Env []EnvVar + + // MountDir is the bare directory the cache should be mounted at, with any + // tunable suffix (such as Habana's ",false,8192") stripped. + MountDir string + + // SubPath is the subpath within the OCI payload to expose. + SubPath string + + // PayloadPrefix is the top-level directory MCV writes the cache payload + // under inside the OCI image (e.g. "io.vllm.cache", "io.habana.cache"). + PayloadPrefix string + + // RequiresWritable is true when the framework needs write access to the + // cache directory, so a read-only mount alone is insufficient and the + // consumer must provide a writable copy (e.g. Habana's flat recipe dir). + RequiresWritable bool +} + +// UnsupportedCacheTypeError is returned when the labels describe a cache type +// this version of cacheplan does not know how to interpret (e.g. an image built +// by a newer MCV). Consumers should treat it as a signal to fall back to legacy +// handling rather than a hard failure. +type UnsupportedCacheTypeError struct { + CacheType string +} + +func (e *UnsupportedCacheTypeError) Error() string { + if e.CacheType == "" { + return "cacheplan: unsupported or unknown cache type" + } + return fmt.Sprintf("cacheplan: unsupported cache type %q", e.CacheType) +} + +// IsUnsupportedCacheType reports whether err is (or wraps) an +// UnsupportedCacheTypeError. +func IsUnsupportedCacheType(err error) bool { + var e *UnsupportedCacheTypeError + return errors.As(err, &e) +} + +// Derive decodes OCI image labels into a typed CachePlan. It dispatches on the +// io.kserve.km/cache-type label, falling back to the per-class summary labels +// for older images. It returns an *UnsupportedCacheTypeError for cache types +// with no serving plan (e.g. bare Triton, or types added by a newer MCV). +func Derive(labels map[string]string) (CachePlan, error) { + if len(labels) == 0 { + return CachePlan{}, fmt.Errorf("cacheplan: no labels provided") + } + + switch detectCacheType(labels) { + case constants.CacheTypeVLLMTorchCompile: + return deriveVLLM(labels) + case constants.CacheTypeHabanaRecipe: + return deriveHabana(labels) + case constants.Triton: + // Bare Triton images carry no mounting metadata; there is no serving + // plan for them today. + return CachePlan{}, &UnsupportedCacheTypeError{CacheType: constants.Triton} + default: + return CachePlan{}, &UnsupportedCacheTypeError{CacheType: labels[LabelCacheType]} + } +} + +// ProducerEnv returns the environment variables that make the given framework +// write its kernel cache to path. This is the capture-side counterpart of +// Derive: the capture pod has no labels to read (nothing has been built yet), so +// the env must be synthesized from the framework/cache type, target path, and +// (for Habana) the maximum cache size in MB. +// +// For non-Habana cache types the sizeMB parameter is ignored. +// Pass DefaultHabanaRecipeCacheSizeMB when no explicit size is configured. +func ProducerEnv(cacheType, path string, sizeMB int) ([]EnvVar, error) { + if path == "" { + return nil, fmt.Errorf("cacheplan: empty cache path") + } + + switch normalizeCacheType(cacheType) { + case constants.CacheTypeHabanaRecipe: + return []EnvVar{{ + Name: constants.HabanaRecipeCacheEnv, + Value: HabanaRecipeEnvValue(path, sizeMB), + }}, nil + case constants.CacheTypeVLLMTorchCompile: + return []EnvVar{{ + Name: constants.VLLMCacheRoot, + Value: path, + }}, nil + default: + return nil, &UnsupportedCacheTypeError{CacheType: cacheType} + } +} + +// RootEnvLabel returns the "NAME=VALUE" string for the io.kserve.km/cache-root-env +// label of a producer cache. It is used by the cache classes' Labels() encoders +// so that the label MCV stamps and the env cacheplan derives stay identical by +// construction. +// +// For Habana caches, sizeMB is the max cache size in MB (see DefaultHabanaRecipeCacheSizeMB). +// For other cache types the sizeMB parameter is ignored. +func RootEnvLabel(cacheType, path string, sizeMB int) (string, error) { + env, err := ProducerEnv(cacheType, path, sizeMB) + if err != nil { + return "", err + } + return env[0].Name + "=" + env[0].Value, nil +} + +// HabanaRecipeEnvValue formats the PT_HPU_RECIPE_CACHE_CONFIG value +// (",false,") for the given cache directory and size limit. +func HabanaRecipeEnvValue(path string, sizeMB int) string { + return fmt.Sprintf("%s,%s,%d", path, habanaRecipeDelete, sizeMB) +} + +// detectCacheType returns the canonical cache type identifier for the labels, +// preferring the explicit io.kserve.km/cache-type label and falling back to the +// per-class summary labels. +func detectCacheType(labels map[string]string) string { + if t := normalizeCacheType(labels[LabelCacheType]); t != "" { + return t + } + + switch { + case has(labels, summaryLabelHabana): + return constants.CacheTypeHabanaRecipe + case has(labels, summaryLabelVLLM): + return constants.CacheTypeVLLMTorchCompile + case has(labels, summaryLabelTriton): + return constants.Triton + } + return "" +} + +func deriveVLLM(labels map[string]string) (CachePlan, error) { + env, mountDir, err := parseRootEnv(labels[LabelCacheRootEnv]) + if err != nil { + return CachePlan{}, err + } + if env.Name != constants.VLLMCacheRoot { + return CachePlan{}, fmt.Errorf("cacheplan: unexpected env name %q in vLLM cache-root-env (want %s)", + env.Name, constants.VLLMCacheRoot) + } + if mountDir == "" { + return CachePlan{}, fmt.Errorf("cacheplan: empty mount directory in vLLM cache-root-env") + } + return CachePlan{ + CacheType: constants.CacheTypeVLLMTorchCompile, + Env: []EnvVar{env}, + MountDir: mountDir, + SubPath: labels[LabelCacheMountSubpath], + PayloadPrefix: constants.MCVVLLMCacheDir, + RequiresWritable: false, + }, nil +} + +func deriveHabana(labels map[string]string) (CachePlan, error) { + env, mountDir, err := parseRootEnv(labels[LabelCacheRootEnv]) + if err != nil { + return CachePlan{}, err + } + if env.Name != constants.HabanaRecipeCacheEnv { + return CachePlan{}, fmt.Errorf("cacheplan: unexpected env name %q in Habana cache-root-env (want %s)", + env.Name, constants.HabanaRecipeCacheEnv) + } + if mountDir == "" { + return CachePlan{}, fmt.Errorf("cacheplan: empty mount directory in Habana cache-root-env") + } + subPath := labels[LabelCacheMountSubpath] + if subPath == "" { + subPath = "." + } + return CachePlan{ + CacheType: constants.CacheTypeHabanaRecipe, + Env: []EnvVar{env}, + MountDir: mountDir, + SubPath: subPath, + PayloadPrefix: constants.MCVHabanaCacheDir, + RequiresWritable: true, + }, nil +} + +// parseRootEnv splits a "NAME=VALUE" cache-root-env label into an EnvVar and the +// bare mount directory. The mount directory is the value up to the first comma, +// so tunable suffixes (Habana's ",false,8192") are stripped while plain values +// (vLLM's "/home/kserve/.cache/vllm") pass through unchanged. +func parseRootEnv(rootEnv string) (EnvVar, string, error) { + if rootEnv == "" { + return EnvVar{}, "", fmt.Errorf("cacheplan: missing %s label", LabelCacheRootEnv) + } + name, value, ok := strings.Cut(rootEnv, "=") + if !ok || name == "" { + return EnvVar{}, "", fmt.Errorf("cacheplan: invalid cache-root-env %q: expected NAME=VALUE", rootEnv) + } + + mountDir := value + if i := strings.IndexByte(value, ','); i >= 0 { + mountDir = value[:i] + } + return EnvVar{Name: name, Value: value}, mountDir, nil +} + +// normalizeCacheType maps the various framework/cache-type aliases callers may +// use onto the canonical cache type identifiers. Returns "" for the empty +// string and passes unknown values through unchanged. +func normalizeCacheType(t string) string { + switch strings.ToLower(strings.TrimSpace(t)) { + case "": + return "" + case constants.VLLM, constants.CacheTypeVLLMTorchCompile: + return constants.CacheTypeVLLMTorchCompile + case constants.Habana, "gaudi", constants.CacheTypeHabanaRecipe: + return constants.CacheTypeHabanaRecipe + case constants.Triton: + return constants.Triton + default: + return t + } +} + +func has(labels map[string]string, key string) bool { + _, ok := labels[key] + return ok +} diff --git a/mcv/pkg/cacheplan/cacheplan_test.go b/mcv/pkg/cacheplan/cacheplan_test.go new file mode 100644 index 000000000..a9446a558 --- /dev/null +++ b/mcv/pkg/cacheplan/cacheplan_test.go @@ -0,0 +1,198 @@ +package cacheplan + +import ( + "testing" + + "github.com/redhat-et/GKM/mcv/pkg/constants" +) + +const ( + testVLLMCacheDir = constants.KServeHome + "/" + constants.VLLMCache // testVLLMCacheDir + testHabanaCacheDir = constants.KServeHome + "/" + constants.HabanaCache // testHabanaCacheDir + testHabanaCacheEnvValue = testHabanaCacheDir + ",false,8192" +) + +func TestDeriveVLLM(t *testing.T) { + labels := map[string]string{ + LabelCacheType: constants.CacheTypeVLLMTorchCompile, + LabelCacheRootEnv: constants.VLLMCacheRoot + "=" + testVLLMCacheDir, + LabelCacheMountSubpath: "torch_compile_cache", + LabelCacheHash: "abc123", + } + + plan, err := Derive(labels) + if err != nil { + t.Fatalf("Derive returned error: %v", err) + } + if plan.CacheType != constants.CacheTypeVLLMTorchCompile { + t.Errorf("CacheType: got %q, want %q", plan.CacheType, constants.CacheTypeVLLMTorchCompile) + } + if len(plan.Env) != 1 || plan.Env[0].Name != constants.VLLMCacheRoot || plan.Env[0].Value != testVLLMCacheDir { + t.Errorf("Env: got %+v", plan.Env) + } + if plan.MountDir != testVLLMCacheDir { + t.Errorf("MountDir: got %q", plan.MountDir) + } + if plan.SubPath != "torch_compile_cache" { + t.Errorf("SubPath: got %q", plan.SubPath) + } + if plan.PayloadPrefix != constants.MCVVLLMCacheDir { + t.Errorf("PayloadPrefix: got %q, want %q", plan.PayloadPrefix, constants.MCVVLLMCacheDir) + } + if plan.RequiresWritable { + t.Error("RequiresWritable: got true, want false for vLLM") + } +} + +func TestDeriveHabana(t *testing.T) { + labels := map[string]string{ + LabelCacheType: constants.CacheTypeHabanaRecipe, + LabelCacheRootEnv: "PT_HPU_RECIPE_CACHE_CONFIG=/home/kserve/.cache/habana,false,8192", + LabelCacheMountSubpath: ".", + } + + plan, err := Derive(labels) + if err != nil { + t.Fatalf("Derive returned error: %v", err) + } + if plan.CacheType != constants.CacheTypeHabanaRecipe { + t.Errorf("CacheType: got %q, want %q", plan.CacheType, constants.CacheTypeHabanaRecipe) + } + // Env must be preserved verbatim, including the ",false,8192" tunables. + if len(plan.Env) != 1 || + plan.Env[0].Name != constants.HabanaRecipeCacheEnv || + plan.Env[0].Value != testHabanaCacheEnvValue { + t.Errorf("Env: got %+v", plan.Env) + } + // MountDir must strip the tunable suffix. + if plan.MountDir != testHabanaCacheDir { + t.Errorf("MountDir: got %q, want %q", plan.MountDir, testHabanaCacheDir) + } + if plan.PayloadPrefix != constants.MCVHabanaCacheDir { + t.Errorf("PayloadPrefix: got %q, want %q", plan.PayloadPrefix, constants.MCVHabanaCacheDir) + } + if !plan.RequiresWritable { + t.Error("RequiresWritable: got false, want true for Habana") + } +} + +// TestDeriveInfersFromSummary verifies the fallback path for older images that +// carry the per-class summary label but no io.kserve.km/cache-type label. +func TestDeriveInfersFromSummary(t *testing.T) { + labels := map[string]string{ + summaryLabelHabana: `{"targets":[{"backend":"hpu"}]}`, + LabelCacheRootEnv: "PT_HPU_RECIPE_CACHE_CONFIG=/home/kserve/.cache/habana,false,8192", + } + plan, err := Derive(labels) + if err != nil { + t.Fatalf("Derive returned error: %v", err) + } + if plan.CacheType != constants.CacheTypeHabanaRecipe { + t.Errorf("CacheType: got %q, want %q", plan.CacheType, constants.CacheTypeHabanaRecipe) + } +} + +func TestDeriveUnsupported(t *testing.T) { + cases := map[string]map[string]string{ + "bare triton": {summaryLabelTriton: `{"targets":[]}`}, + "unknown type": {LabelCacheType: "future-cache-type"}, + } + for name, labels := range cases { + t.Run(name, func(t *testing.T) { + _, err := Derive(labels) + if err == nil { + t.Fatal("expected error, got nil") + } + if !IsUnsupportedCacheType(err) { + t.Errorf("expected UnsupportedCacheTypeError, got %v", err) + } + }) + } +} + +func TestDeriveErrors(t *testing.T) { + t.Run("nil labels", func(t *testing.T) { + if _, err := Derive(nil); err == nil { + t.Error("expected error for nil labels") + } + }) + t.Run("missing cache-root-env", func(t *testing.T) { + _, err := Derive(map[string]string{LabelCacheType: constants.CacheTypeVLLMTorchCompile}) + if err == nil { + t.Error("expected error for missing cache-root-env") + } + }) + t.Run("malformed cache-root-env", func(t *testing.T) { + _, err := Derive(map[string]string{ + LabelCacheType: constants.CacheTypeVLLMTorchCompile, + LabelCacheRootEnv: "NOTANENVVAR", + }) + if err == nil { + t.Error("expected error for malformed cache-root-env") + } + }) +} + +func TestProducerEnv(t *testing.T) { + tests := []struct { + name string + cacheType string + path string + wantName string + wantValue string + }{ + {"habana alias", "habana", testHabanaCacheDir, constants.HabanaRecipeCacheEnv, testHabanaCacheEnvValue}, + {"gaudi alias", "gaudi", "/data/habana", constants.HabanaRecipeCacheEnv, "/data/habana,false,8192"}, + {"habana-recipe id", constants.CacheTypeHabanaRecipe, "/c", constants.HabanaRecipeCacheEnv, "/c,false,8192"}, + {"vllm alias", "vllm", testVLLMCacheDir, constants.VLLMCacheRoot, testVLLMCacheDir}, + {"torch-compile id", constants.CacheTypeVLLMTorchCompile, "/v", constants.VLLMCacheRoot, "/v"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + env, err := ProducerEnv(tt.cacheType, tt.path, DefaultHabanaRecipeCacheSizeMB) + if err != nil { + t.Fatalf("ProducerEnv error: %v", err) + } + if len(env) != 1 || env[0].Name != tt.wantName || env[0].Value != tt.wantValue { + t.Errorf("got %+v, want {%s %s}", env, tt.wantName, tt.wantValue) + } + }) + } +} + +func TestProducerEnvErrors(t *testing.T) { + t.Run("empty path", func(t *testing.T) { + if _, err := ProducerEnv("vllm", "", DefaultHabanaRecipeCacheSizeMB); err == nil { + t.Error("expected error for empty path") + } + }) + t.Run("unsupported type", func(t *testing.T) { + _, err := ProducerEnv("triton", "/c", DefaultHabanaRecipeCacheSizeMB) + if err == nil || !IsUnsupportedCacheType(err) { + t.Errorf("expected UnsupportedCacheTypeError, got %v", err) + } + }) +} + +// TestRootEnvLabelRoundTrip proves the produce/consume symmetry: the label +// RootEnvLabel writes decodes back to the same env and mount dir via Derive. +func TestRootEnvLabelRoundTrip(t *testing.T) { + rootEnv, err := RootEnvLabel(constants.Habana, testHabanaCacheDir, DefaultHabanaRecipeCacheSizeMB) + if err != nil { + t.Fatalf("RootEnvLabel error: %v", err) + } + + plan, err := Derive(map[string]string{ + LabelCacheType: constants.CacheTypeHabanaRecipe, + LabelCacheRootEnv: rootEnv, + }) + if err != nil { + t.Fatalf("Derive error: %v", err) + } + if plan.Env[0].Value != testHabanaCacheEnvValue { + t.Errorf("round-trip env value: got %q", plan.Env[0].Value) + } + if plan.MountDir != testHabanaCacheDir { + t.Errorf("round-trip mount dir: got %q", plan.MountDir) + } +} diff --git a/mcv/pkg/constants/constants.go b/mcv/pkg/constants/constants.go index 6691da1b6..115d80da8 100644 --- a/mcv/pkg/constants/constants.go +++ b/mcv/pkg/constants/constants.go @@ -11,6 +11,7 @@ import ( const ( VLLM = "vllm" Triton = "triton" + Habana = "habana" MCVBuildDir = "/tmp/.mcv" CacheDir = "cache" ManifestDir = "manifest" @@ -18,11 +19,16 @@ const ( VLLMHOME = "/home/vllm" KServeHome = "/home/kserve" VLLMCache = ".cache/vllm" + // HabanaCache is the Habana recipe cache directory relative to KServeHome, + // i.e. where the recipe cache is mounted inside a KServe serving container. + HabanaCache = ".cache/habana" MCVTritonCacheDir = "io.triton.cache/" MCVTritonManifestDir = "io.triton.manifest" MCVVLLMCacheDir = "io.vllm.cache" MCVVLLMManifestDir = "io.vllm.manifest" + MCVHabanaCacheDir = "io.habana.cache" + MCVHabanaManifestDir = "io.habana.manifest" EnvTritonCacheDir = "TRITON_CACHE_DIR" DefaultCacheFilePath = "/tmp/device_cache.json" @@ -35,6 +41,17 @@ const ( // Cache type identifiers CacheTypeVLLMTorchCompile = "torch-compile" + CacheTypeHabanaRecipe = "habana-recipe" + + // Habana recipe cache env var + HabanaRecipeCacheEnv = "PT_HPU_RECIPE_CACHE_CONFIG" + + // Accelerator backend identifiers — canonical strings used in TritonGPUInfo.Backend, + // OCI image labels, and cache-type detection. + BackendCUDA = "cuda" + BackendHIP = "hip" + BackendROCm = "rocm" + BackendHPU = "hpu" ) // Configurable runtime paths @@ -43,6 +60,7 @@ var ( ExtractCacheDir string ExtractManifestDir string VLLMCacheDir string + HabanaCacheDir string HasTritonCache bool HasVLLMCache bool LogLevels = []string{"debug", "info", "warning", "error"} // accepted log levels @@ -73,4 +91,6 @@ func init() { if _, err := os.Stat(VLLMCacheDir); err == nil { HasVLLMCache = true } + + HabanaCacheDir = filepath.Join(home, ".cache", "habana") } diff --git a/mcv/pkg/fetcher/imgfetcher.go b/mcv/pkg/fetcher/imgfetcher.go index 70aaa95da..c89450cf2 100644 --- a/mcv/pkg/fetcher/imgfetcher.go +++ b/mcv/pkg/fetcher/imgfetcher.go @@ -208,6 +208,8 @@ func (e *cacheExtractor) ExtractCache(img v1.Image) error { constants.ExtractCacheDir = constants.TritonCacheDir case constants.VLLM: constants.ExtractCacheDir = constants.VLLMCacheDir + case constants.Habana: + constants.ExtractCacheDir = constants.HabanaCacheDir default: return fmt.Errorf("unsupported cache type: %s", cacheType) } @@ -405,6 +407,8 @@ func validateExtractedCacheSize(labels map[string]string, cacheType string, extr labelKey = "cache.triton.image/cache-size-bytes" case constants.VLLM: labelKey = "cache.vllm.image/cache-size-bytes" + case constants.Habana: + labelKey = "cache.habana.image/cache-size-bytes" default: return fmt.Errorf("unsupported cache type: %s", cacheType) } diff --git a/mcv/pkg/preflightcheck/triton.go b/mcv/pkg/preflightcheck/triton.go index 7293efa05..5caea4b11 100644 --- a/mcv/pkg/preflightcheck/triton.go +++ b/mcv/pkg/preflightcheck/triton.go @@ -71,7 +71,7 @@ func CompareTritonEntriesToGPU(entries []cache.TritonCacheMetadata, devInfo []de warpMatches := entry.WarpSize == gpuInfo.WarpSize ptxMatches := true - if entry.Backend == "cuda" { + if entry.Backend == cache.CUDABackend { ptxMatches = entry.PtxVersion == gpuInfo.PTXVersion } diff --git a/mcv/pkg/preflightcheck/utils.go b/mcv/pkg/preflightcheck/utils.go index eb5699fed..ad53a6086 100644 --- a/mcv/pkg/preflightcheck/utils.go +++ b/mcv/pkg/preflightcheck/utils.go @@ -18,7 +18,7 @@ import ( // normalizeArchForComparison normalizes architecture strings for comparison // Strips sm_ prefix from CUDA architectures to handle both "75" and "sm_75" formats func normalizeArchForComparison(backend, arch string) string { - if backend == "cuda" { + if backend == constants.BackendCUDA { return strings.TrimPrefix(arch, "sm_") } return arch @@ -41,7 +41,9 @@ func CompareCacheSummaryLabelToGPU(img v1.Image, labels map[string]string, devIn summaryStr, ok := labels["cache.triton.image/summary"] if !ok { if summaryStr, ok = labels["cache.vllm.image/summary"]; !ok { - return nil, nil, errors.New("image missing cache summary label") + if summaryStr, ok = labels["cache.habana.image/summary"]; !ok { + return nil, nil, errors.New("image missing cache summary label") + } } } @@ -62,6 +64,15 @@ func CompareCacheSummaryLabelToGPU(img v1.Image, labels map[string]string, devIn isMatch := false for _, target := range summary.Targets { backendMatches := target.Backend == gpu.Backend + + // HPU (Habana) caches are backend-matched only — recipe files + // don't encode arch or warp size. + if backendMatches && target.Backend == constants.BackendHPU { + logging.Debugf("Habana backend match: target=%s, gpu=%s", target.Backend, gpu.Backend) + isMatch = true + break + } + // Normalize architectures for comparison (handles "75" vs "sm_75" for CUDA) normalizedTargetArch := normalizeArchForComparison(target.Backend, target.Arch) normalizedGPUArch := normalizeArchForComparison(gpu.Backend, gpu.Arch) @@ -105,6 +116,9 @@ func DetectCacheTypeFromLabels(labels map[string]string) (string, error) { if _, ok := labels["cache.vllm.image/summary"]; ok { return constants.VLLM, nil } + if _, ok := labels["cache.habana.image/summary"]; ok { + return constants.Habana, nil + } return "", fmt.Errorf("unknown cache type from labels") } @@ -118,6 +132,10 @@ func CompareCacheManifestToGPU(manifestPath, cacheType string, devInfo []devices return CompareTritonCacheManifestToGPU(manifestPath, devInfo) case constants.VLLM: return CompareVLLMCacheManifestToGPU(manifestPath, devInfo) + case constants.Habana: + // Habana recipe caches are backend-matched at the summary level; + // no per-recipe manifest comparison is needed. + return nil default: return fmt.Errorf("unsupported cache type: %s", cacheType) }