Files
llmgw/internal/cli/enginecheck_gpu_test.go
T
Ilker ManapandClaude Opus 5.5 e4e15200d6 engine-check computes a token on the GPU, and finds the installed llama.cpp
Seeing a device is not computing on it. For llama.cpp on a machine with
NVIDIA GPUs, engine-check now loads a small GGUF (--model, or the smallest in
the weights cache) with every layer on the GPU, asks for one token, and
reads from llama.cpp's own output where the layers went: a build that
computes on the CPU fails, naming the fix. With no --bin it checks the build
the agent installed under /opt/llmgw/engines.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-07 13:04:38 +02:00

67 lines
2.1 KiB
Go

package cli
import (
"bytes"
"fmt"
"net/http"
"os"
"path/filepath"
"strings"
"testing"
)
// TestMain lets the test binary stand in for llama-server: run with
// LLMGW_FAKE_LLAMA set, it prints what llama.cpp prints and answers.
func TestMain(m *testing.M) {
if mode := os.Getenv("LLMGW_FAKE_LLAMA"); mode != "" {
fakeLlama(mode)
return
}
os.Exit(m.Run())
}
func fakeLlama(mode string) {
port := ""
for i, a := range os.Args {
if a == "--port" && i+1 < len(os.Args) {
port = os.Args[i+1]
}
}
if mode == "gpu" {
fmt.Println("llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free")
fmt.Println("load_tensors: offloaded 25/25 layers to GPU")
} else {
fmt.Println("warning: no usable GPU found, --gpu-layers option will be ignored")
fmt.Println("load_tensors: offloaded 0/25 layers to GPU")
}
http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })
http.HandleFunc("/completion", func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(`{"content":"!"}`)) })
_ = http.ListenAndServe("127.0.0.1:"+port, nil) //nolint:gosec // a test's stand-in
}
func TestEngineCheckComputesATokenOnTheGPU(t *testing.T) {
// nvidia-smi present, so the check runs.
dir := t.TempDir()
smi := filepath.Join(dir, "nvidia-smi")
if err := os.WriteFile(smi, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
t.Fatal(err)
}
t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
model := filepath.Join(dir, "tiny.gguf")
_ = os.WriteFile(model, []byte("GGUF"), 0o644)
self, _ := os.Executable()
t.Setenv("LLMGW_FAKE_LLAMA", "gpu")
var out bytes.Buffer
if err := computeOnGPU(t.Context(), self, model, &out); err != nil || !strings.Contains(out.String(), "on CUDA") ||
!strings.Contains(out.String(), "25 of 25") {
t.Fatalf("on the GPU: %v %s", err, out.String())
}
t.Setenv("LLMGW_FAKE_LLAMA", "cpu")
out.Reset()
if err := computeOnGPU(t.Context(), self, model, &out); err == nil || !strings.Contains(err.Error(), "computed on the CPU") {
t.Fatalf("a build that falls back to the CPU: %v %s", err, out.String())
}
}