llmgw/internal/agent/provision/install_test.go

package provision

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
)

var ubuntu = discovery.Host{OSID: "ubuntu", OSVersion: "24.04", Arch: "amd64", CPUFlags: []string{"avx2"}}

// release is a fake GitHub: one release, its assets and their digests.
type release struct {
	assets  map[string][]byte
	digests map[string]string
	gets    int
	mu      sync.Mutex
}

func tarGz(t *testing.T, files map[string]string, links map[string]string) []byte {
	t.Helper()
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	for name, body := range files {
		_ = tw.WriteHeader(&tar.Header{Name: name, Mode: 0o755, Size: int64(len(body)), Typeflag: tar.TypeReg})
		_, _ = tw.Write([]byte(body))
	}
	for name, target := range links {
		_ = tw.WriteHeader(&tar.Header{Name: name, Linkname: target, Typeflag: tar.TypeSymlink})
	}
	_ = tw.Close()
	_ = gz.Close()
	return buf.Bytes()
}

func (r *release) serve(t *testing.T) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		r.mu.Lock()
		r.gets++
		r.mu.Unlock()
		if strings.HasPrefix(req.URL.Path, "/repos/") {
			var out struct {
				Assets []map[string]string `json:"assets"`
			}
			for name := range r.assets {
				out.Assets = append(out.Assets, map[string]string{"name": name, "digest": r.digests[name]})
			}
			_ = json.NewEncoder(w).Encode(out)
			return
		}
		body, ok := r.assets[filepath.Base(req.URL.Path)]
		if !ok {
			http.NotFound(w, req)
			return
		}
		_, _ = w.Write(body)
	}))
	t.Cleanup(srv.Close)
	return srv
}

func newRelease(assets map[string][]byte) *release {
	r := &release{assets: assets, digests: map[string]string{}}
	for name, b := range assets {
		sum := sha256.Sum256(b)
		r.digests[name] = "sha256:" + hex.EncodeToString(sum[:])
	}
	return r
}

// fakeSystem answers the package tools, ldd and the engine itself.
type fakeSystem struct {
	mu        sync.Mutex
	installed map[string]bool
	tools     map[string]bool
	lddOut    string
	calls     []string
	// locked is how many more package commands find the lock held.
	locked int
}

var errExit = errors.New("exit 1")

func (f *fakeSystem) run(_ context.Context, _ []string, name string, args ...string) ([]byte, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls = append(f.calls, name+" "+strings.Join(args, " "))
	switch {
	case name == "sh" && strings.Contains(args[1], "dpkg-query"),
		name == "rpm" && args[0] == "-q":
		if f.installed[args[len(args)-1]] {
			return nil, nil
		}
		return nil, errExit
	case name == "sh" && strings.Contains(args[1], "command -v"):
		if f.tools[args[2]] {
			return nil, nil
		}
		return nil, errExit
	case name == "apt-get" || name == "dnf":
		if f.locked > 0 {
			f.locked--
			return []byte("E: Unable to lock directory /var/lib/apt/lists/\n"), errExit
		}
		for i, a := range args {
			if a == "install" {
				for _, p := range args[i+1:] {
					if !strings.HasPrefix(p, "-") {
						f.installed[p] = true
					}
				}
			}
		}
		return nil, nil
	case name == "ldd":
		return []byte(f.lddOut), nil
	case strings.HasSuffix(name, "llama-server"):
		return []byte("version: 10991 (abc)\n"), nil
	}
	return nil, errors.New("unexpected command " + name)
}

func installer(t *testing.T, srv *httptest.Server, sys *fakeSystem, host discovery.Host) (*Installer, *[]Status) {
	var seen []Status
	var mu sync.Mutex
	return &Installer{
		Root: t.TempDir(), Host: host, Run: sys.run,
		APIBase: srv.URL, DownloadBase: srv.URL, HTTP: srv.Client(),
		OnStatus: func(s Status) { mu.Lock(); seen = append(seen, s); mu.Unlock() },
	}, &seen
}

func llamaRelease(t *testing.T) *release {
	return newRelease(map[string][]byte{
		"llama-b1-bin-ubuntu-x64.tar.gz": tarGz(t,
			map[string]string{"llama-b1/llama-server": "#!/bin/sh\n", "llama-b1/libllama.so.0": "x"},
			map[string]string{"llama-b1/libllama.so": "libllama.so.0"}),
	})
}

func TestAnEngineIsInstalledFromNothing(t *testing.T) {
	// A node rebuilt from a bare image: no engine, no OpenMP runtime. What it
	// ends with has to have come from the publisher, checked, and been run.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}}
	in, seen := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if len(got) != 1 || got[0].State != Ready {
		t.Fatalf("not ready: %+v", got)
	}
	if !sys.installed["libgomp1"] {
		t.Error("the OpenMP runtime the build links against was not installed")
	}
	if !strings.HasPrefix(got[0].Binary, in.Root) {
		t.Errorf("binary %s is not the one installed under %s", got[0].Binary, in.Root)
	}
	if _, err := os.Lstat(filepath.Join(filepath.Dir(got[0].Binary), "libllama.so")); err != nil {
		t.Errorf("the library symlink was not unpacked: %v", err)
	}
	var states []string
	for _, s := range *seen {
		states = append(states, s.State+": "+s.Detail)
	}
	if !strings.Contains(strings.Join(states, "\n"), "downloading") {
		t.Errorf("a download was not reported while it ran:\n%s", strings.Join(states, "\n"))
	}
}

func TestAnInstalledEngineIsCheckedNotFetchedAgain(t *testing.T) {
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)
	w := []Want{{Engine: "llamacpp", Version: "b1"}}
	in.Converge(t.Context(), w)
	before := rel.gets

	got := in.Converge(t.Context(), w)
	if got[0].State != Ready {
		t.Fatalf("second pass: %+v", got[0])
	}
	if rel.gets != before {
		t.Errorf("a ready engine was downloaded again (%d requests)", rel.gets-before)
	}
}

func TestAnEngineThatCannotLoadALibraryIsNotReady(t *testing.T) {
	// Found on a node before: llama-server present, executable, and unable to
	// start for want of a library — while the node reported ready.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true},
		lddOut: "\tlibcuda.so.1 => not found\n"}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "libcuda.so.1") {
		t.Fatalf("expected a failure naming the library: %+v", got[0])
	}
	if _, err := os.Stat(filepath.Join(in.Root, "llamacpp", "b1-cpu", marker)); err == nil {
		t.Error("a build that does not run was marked installed")
	}
}

func TestADownloadThatDoesNotMatchItsChecksumIsRefused(t *testing.T) {
	rel := llamaRelease(t)
	for name := range rel.digests {
		rel.digests[name] = "sha256:" + strings.Repeat("0", 64)
	}
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "checksum") {
		t.Fatalf("a tampered download was not refused: %+v", got[0])
	}
}

func TestAnArchiveCannotWriteOutsideItsDirectory(t *testing.T) {
	rel := newRelease(map[string][]byte{
		"llama-b1-bin-ubuntu-x64.tar.gz": tarGz(t,
			map[string]string{"../../escaped": "x", "llama-server": ""}, nil),
	})
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed {
		t.Fatalf("an archive escaping its directory was installed: %+v", got[0])
	}
}

func TestTheBuildFollowsTheHardware(t *testing.T) {
	gpu := ubuntu
	gpu.NVIDIADriver = "570.86.15"
	b, err := chooseLlamaCpp("b1", gpu)
	if err != nil || !strings.HasPrefix(b.Variant, "cuda") || len(b.Assets) != 2 {
		t.Errorf("an NVIDIA node got %+v, %v; expected the CUDA build and its runtime", b, err)
	}
	if _, err := chooseLlamaCpp("b1", discovery.Host{OSID: "ubuntu", Arch: "riscv64"}); err == nil {
		t.Error("a build was chosen for an architecture llama.cpp does not publish for")
	}

	slow := ubuntu
	slow.CPUFlags = nil
	st := (&Installer{Host: slow}).ready(Status{}, Build{CPUFlags: []string{"avx2"}}, "/x/llama-server", "")
	if !strings.Contains(st.Detail, "avx2") {
		t.Errorf("a CPU without AVX2 was not called out: %q", st.Detail)
	}
}

func TestAnEngineTheAgentCannotInstallSaysSo(t *testing.T) {
	in := &Installer{Root: t.TempDir(), Host: ubuntu}
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.29.0"}})
	if got[0].State != Unsupported {
		t.Errorf("got %+v", got[0])
	}
}

func TestARedHatFamilyNodeInstallsWithItsOwnPackageManager(t *testing.T) {
	// The build is not Ubuntu's in any way that matters; the package that
	// carries OpenMP's runtime is named differently, and installed by dnf.
	rocky := discovery.Host{OSID: "rocky", OSVersion: "10.0", OSLike: []string{"rhel", "centos", "fedora"},
		Arch: "amd64", CPUFlags: []string{"avx2"}}
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}, tools: map[string]bool{"dnf": true}}
	in, _ := installer(t, rel.serve(t), sys, rocky)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	if !sys.installed["libgomp"] || sys.installed["libgomp1"] {
		t.Errorf("installed %v; expected the rpm name, libgomp", sys.installed)
	}
}

func TestASystemTooOldForTheBuildSaysWhatItLacks(t *testing.T) {
	// RHEL 9: glibc is new enough, libstdc++ is one GCC release short. ldd
	// finds the library and reports the symbol, and "libstdc++ not found"
	// would send somebody looking for a package that is installed.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true},
		lddOut: "\t/lib64/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /x/libllama.so)\n" +
			"\t/lib64/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /x/llama-server)\n"}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "libstdc++.so.6 is too old (needs GLIBCXX_3.4.30)") {
		t.Fatalf("detail %q", got[0].Detail)
	}
	if strings.Count(got[0].Detail, "GLIBCXX") != 1 {
		t.Errorf("the same shortfall was listed more than once: %q", got[0].Detail)
	}
}

func TestAPackageManagerBusyWithItsOwnUpdatesIsWaitedFor(t *testing.T) {
	// A new node's first minutes: cloud-init and unattended-upgrades hold
	// apt's lists lock, which DPkg::Lock::Timeout does not cover.
	lockRetry = time.Millisecond
	t.Cleanup(func() { lockRetry = 10 * time.Second })

	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}, locked: 3}
	in, seen := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Ready {
		t.Fatalf("gave up on a lock another apt held: %+v", got[0])
	}
	waited := false
	for _, s := range *seen {
		waited = waited || strings.Contains(s.Detail, "waiting for another apt-get")
	}
	if !waited {
		t.Error("the wait was not reported")
	}
}