diff --git a/internal/agent/engines/llamacpp.go b/internal/agent/engines/llamacpp.go index fb7bf22..ed5cc29 100644 --- a/internal/agent/engines/llamacpp.go +++ b/internal/agent/engines/llamacpp.go @@ -277,6 +277,9 @@ func (l *LlamaCpp) bin() string { return "llama-server" } +// Binary is the llama-server this engine runs. +func (l *LlamaCpp) Binary() string { return l.bin() } + // FallsBackToCPU says that llama.cpp keeps running when its GPU backend does // not load: --n-gpu-layers asks for layers on a GPU it then has none of, and // the server comes up and answers on the CPU, at a tenth of the speed, with no diff --git a/internal/cli/enginecheck.go b/internal/cli/enginecheck.go index bed0d0b..ebc506b 100644 --- a/internal/cli/enginecheck.go +++ b/internal/cli/enginecheck.go @@ -4,10 +4,15 @@ import ( "context" "fmt" "io" + "net" + "net/http" + "os" "os/exec" "path/filepath" "sort" + "strconv" "strings" + "sync" "time" "github.com/spf13/cobra" @@ -22,6 +27,7 @@ func newEngineCheckCmd() *cobra.Command { nodes int tp int perNode int + model string ) cmd := &cobra.Command{ @@ -31,6 +37,10 @@ func newEngineCheckCmd() *cobra.Command { "engines rename and remove flags between versions. When that happens the\n" + "failure arrives minutes into a model load, in an engine's own error\n" + "message, on a node nobody was watching.\n\n" + + "For llama.cpp on a machine with NVIDIA GPUs it then loads a small GGUF\n" + + "(--model, or the smallest in the weights cache) with every layer on the GPU\n" + + "and asks for one token: a build whose CUDA backend does not load computes on\n" + + "the CPU, and that is where it is caught.\n\n" + "This asks the engine itself. It runs the engine's --help and checks every\n" + "flag llmgw would pass against it, which needs no GPU — a CLI prints its\n" + "help on any machine. Run it after upgrading an engine.\n\n" + @@ -49,6 +59,11 @@ func newEngineCheckCmd() *cobra.Command { // rarely on PATH: a check that cannot find the engine reports // nothing rather than reporting a problem. cfg := engines.Config{} + // The build the agent installed, when none is named: the newest + // under /opt/llmgw/engines, which is where every node keeps it. + if bin == "" && name == "llamacpp" { + bin = installedLlama() + } if bin != "" { switch name { case "llamacpp": @@ -135,6 +150,11 @@ func newEngineCheckCmd() *cobra.Command { var deviceErr error if fb, ok := eng.(engines.CPUFallback); ok && help == "" { deviceErr = checkDevices(cmd.Context(), name, fb, out) + // Seeing a GPU is not computing on it: a small model loaded + // on it and asked for a token is. + if b, ok := eng.(interface{ Binary() string }); ok && deviceErr == nil && model != "none" { + deviceErr = computeOnGPU(cmd.Context(), b.Binary(), model, out) + } } fmt.Fprintf(out, "%s: %d flags checked", name, checked) @@ -160,7 +180,9 @@ func newEngineCheckCmd() *cobra.Command { f := cmd.Flags() f.StringVar(&help, "help-output", "", "a saved --help to check against, instead of running the engine") - f.StringVar(&bin, "bin", "", "where the engine lives, when it is not on PATH") + f.StringVar(&bin, "bin", "", "where the engine lives (default: the build the agent installed, else PATH)") + f.StringVar(&model, "model", "", "a GGUF file to load on the GPU and ask for a token (default: the smallest "+ + "in this node's weights cache; none skips it)") f.IntVar(&nodes, "nodes", 2, "shape the command for a replica of this many nodes") f.IntVar(&tp, "tp", 8, "tensor parallel degree") f.IntVar(&perNode, "gpus-per-node", 4, "GPUs each node contributes") @@ -286,3 +308,128 @@ func checkDevices(ctx context.Context, name string, fb engines.CPUFallback, out fmt.Fprintf(out, "computes on: %s\n", strings.Join(devs, "; ")) return nil } + +// installedLlama is the newest llama-server the agent installed here. +func installedLlama() string { + m, _ := filepath.Glob("/opt/llmgw/engines/llamacpp/*/llama-server") + if len(m) == 0 { + return "" + } + sort.Strings(m) + return m[len(m)-1] +} + +// weightDirs are where a node keeps models it has fetched. +var weightDirs = []string{"/var/lib/llmgw-weights", "/var/lib/llmgw-models", "/srv/models"} + +// smallestGGUF is the smallest GGUF file under dirs, for a check that only +// needs something to load. +func smallestGGUF(dirs []string) string { + best, size := "", int64(0) + for _, d := range dirs { + _ = filepath.WalkDir(d, func(p string, e os.DirEntry, err error) error { + if err != nil || e.IsDir() || !strings.HasSuffix(strings.ToLower(p), ".gguf") || + strings.Contains(strings.ToLower(filepath.Base(p)), "mmproj") { + return nil + } + if info, err := e.Info(); err == nil && (best == "" || info.Size() < size) { + best, size = p, info.Size() + } + return nil + }) + } + return best +} + +// computeOnGPU loads a model with every layer on the GPU, asks for one +// token, and reads from llama.cpp's own output where the layers went. +func computeOnGPU(ctx context.Context, bin, model string, out io.Writer) error { + if _, err := exec.LookPath("nvidia-smi"); err != nil { + return nil + } + if model == "" { + if model = smallestGGUF(weightDirs); model == "" { + fmt.Fprintln(out, "no GGUF in this node's weights cache to compute with; --model PATH runs one") + return nil + } + } + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return err + } + port := ln.Addr().(*net.TCPAddr).Port + _ = ln.Close() + ctx, cancel := context.WithTimeout(ctx, 3*time.Minute) + defer cancel() + var logs syncBuffer + c := exec.CommandContext(ctx, bin, "-m", model, "--host", "127.0.0.1", "--port", strconv.Itoa(port), //nolint:gosec // the engine's own binary + "--n-gpu-layers", "999", "--ctx-size", "512", "--parallel", "1") + c.Stdout, c.Stderr = &logs, &logs + if err := c.Start(); err != nil { + return fmt.Errorf("starting %s: %w", bin, err) + } + defer func() { _ = c.Process.Kill(); _ = c.Wait() }() + base := "http://127.0.0.1:" + strconv.Itoa(port) + start := time.Now() + for { + resp, err := http.Get(base + "/health") //nolint:gosec,noctx // loopback, bounded by the loop + if err == nil { + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + break + } + } + if ctx.Err() != nil { + return fmt.Errorf("%s did not come up with %s: %s", bin, filepath.Base(model), tailLines(logs.String(), 5)) + } + time.Sleep(500 * time.Millisecond) + } + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, base+"/completion", + strings.NewReader(`{"prompt":"Hello","n_predict":1}`)) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return fmt.Errorf("asking it for a token: %w", err) + } + _ = resp.Body.Close() + took := time.Since(start) + comp, said := engines.ReadLlamaCompute(strings.Split(logs.String(), "\n")) + if !said { + fmt.Fprintf(out, "computed a token with %s, but llama.cpp did not say where\n", filepath.Base(model)) + return nil + } + if comp.OnTheCPU() { + return fmt.Errorf("%s computed on the CPU though every layer was asked for on the GPU (%d of %d offloaded): "+ + "its CUDA backend did not load; ldd on its libggml-cuda.so names the missing library", + filepath.Base(model), comp.Offloaded, comp.Total) + } + fmt.Fprintf(out, "computed a token on %s: %s, %d of %d layers offloaded, in %s\n", comp.Backend, + strings.Join(comp.Devices, "; "), comp.Offloaded, comp.Total, took.Round(100*time.Millisecond)) + return nil +} + +func tailLines(s string, n int) string { + l := strings.Split(strings.TrimSpace(s), "\n") + if len(l) > n { + l = l[len(l)-n:] + } + return strings.Join(l, " | ") +} + +// syncBuffer is a buffer two writers share. +type syncBuffer struct { + mu sync.Mutex + b strings.Builder +} + +func (s *syncBuffer) Write(p []byte) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.b.Write(p) +} + +func (s *syncBuffer) String() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.b.String() +} diff --git a/internal/cli/enginecheck_gpu_test.go b/internal/cli/enginecheck_gpu_test.go new file mode 100644 index 0000000..61fec52 --- /dev/null +++ b/internal/cli/enginecheck_gpu_test.go @@ -0,0 +1,66 @@ +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()) + } +}