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>
67 lines
2.1 KiB
Go
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())
|
|
}
|
|
}
|