llmgw/internal/agent/provision/offline_test.go

package provision

import (
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"
	"testing"

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

// offlineMirror is a cluster's mirror holding a bundle's engine files, and
// nothing that answers for GitHub.
type offlineMirror struct {
	files map[string][]byte
	mu    sync.Mutex
	paths []string
}

func (m *offlineMirror) serve(t *testing.T) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		m.mu.Lock()
		m.paths = append(m.paths, r.URL.Path)
		m.mu.Unlock()
		b, ok := m.files[r.URL.Path]
		if !ok {
			http.NotFound(w, r)
			return
		}
		_, _ = w.Write(b)
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestAnOfflineNodeInstallsFromTheMirrorWithChecksumsFromTheControlPlane(t *testing.T) {
	rel := llamaRelease(t)
	asset := "llama-b1-bin-ubuntu-x64.tar.gz"
	m := &offlineMirror{files: map[string][]byte{"/_llmgw/engines/llamacpp/b1/" + asset: rel.assets[asset]}}
	srv := m.serve(t)
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
	// GitHub is the mirror too, so a request for it would be seen and
	// answered 404: an offline install must not ask.
	in, _ := installer(t, srv, sys, ubuntu)

	want := Want{Engine: "llamacpp", Version: "b1", Offline: &Offline{
		Base:   srv.URL + "/_llmgw/engines/llamacpp/b1",
		SHA256: map[string]string{asset: strings.TrimPrefix(rel.digests[asset], "sha256:")},
	}}
	got := in.Converge(t.Context(), []Want{want})
	if got[0].State != Ready {
		t.Fatalf("not installed from the mirror: %+v", got[0])
	}
	for _, p := range m.paths {
		if strings.HasPrefix(p, "/repos/") || strings.Contains(p, "/releases/download/") {
			t.Errorf("an offline install asked for %s", p)
		}
	}

	// A file on the mirror that is not what the bundle recorded is refused,
	// whatever the mirror says of it.
	m.files["/_llmgw/engines/llamacpp/b2/llama-b2-bin-ubuntu-x64.tar.gz"] = []byte("replaced on the mirror")
	got = in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b2", Offline: &Offline{
		Base:   srv.URL + "/_llmgw/engines/llamacpp/b2",
		SHA256: map[string]string{"llama-b2-bin-ubuntu-x64.tar.gz": strings.Repeat("a", 64)},
	}}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "checksum") {
		t.Errorf("a file not matching the bundle's checksum was not refused: %+v", got[0])
	}

	// One the bundle has no checksum for is not fetched at all, and the
	// operator is told which bundle would have it.
	got = in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b3", Offline: &Offline{
		Base: srv.URL + "/_llmgw/engines/llamacpp/b3",
	}}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "bundle create --engine llamacpp=b3") {
		t.Errorf("a version missing from the bundle did not say how to add it: %+v", got[0])
	}
}

func TestAnOfflineNodeIsToldPythonEnginesComeFromPyPI(t *testing.T) {
	gpu := ubuntu
	gpu.NVIDIADriver, gpu.MinComputeCapability = "580.65", 8.9
	in := &Installer{Root: t.TempDir(), Host: gpu, Run: (&fakeSystem{}).run}
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0", Offline: &Offline{Base: "http://mirror"}}})
	if got[0].State != Unsupported || !strings.Contains(got[0].Detail, "PyPI") {
		t.Fatalf("vLLM offline: %+v", got[0])
	}

	// A node it is not meant for still says nothing is wrong.
	in.Host = ubuntu
	got = in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0", Offline: &Offline{Base: "http://mirror"}}})
	if got[0].State != Skipped {
		t.Errorf("vLLM offline on a CPU node: %+v", got[0])
	}
}

func TestABundleCarriesEveryBuildTheNodesMightChoose(t *testing.T) {
	variants := func(arch string) []string {
		bs, err := Builds("llamacpp", "b1", arch)
		if err != nil {
			t.Fatal(err)
		}
		var out []string
		for _, b := range bs {
			out = append(out, b.Variant)
		}
		return out
	}
	if got := variants("amd64"); !slices.Equal(got, []string{"cpu", "cuda-13.3", "cuda-12.8", "vulkan"}) {
		t.Errorf("amd64 builds: %v", got)
	}
	if got := variants("arm64"); !slices.Equal(got, []string{"cpu", "cuda-13.3", "vulkan"}) {
		t.Errorf("arm64 builds: %v", got)
	}

	// Each one a node would choose is among them, file for file.
	bundled := map[string][]string{}
	bs, _ := Builds("llamacpp", "b1", "amd64")
	for _, b := range bs {
		bundled[b.Variant] = b.Assets
	}
	for _, h := range []discovery.Host{
		ubuntu,
		{Arch: "amd64", NVIDIADriver: "550.1", MinComputeCapability: 8.6},
		{Arch: "amd64", NVIDIADriver: "470.1", MinComputeCapability: 6.1},
	} {
		chosen, err := llamaCppCandidates("b1", "", h)
		if err != nil {
			t.Fatal(err)
		}
		for _, b := range chosen {
			if !slices.Equal(bundled[b.Variant], b.Assets) {
				t.Errorf("%+v chooses %s %v, which the bundle carries as %v", h, b.Variant, b.Assets, bundled[b.Variant])
			}
		}
	}

	if _, err := Builds("vllm", "0.30.0", "amd64"); err == nil || !strings.Contains(err.Error(), "PyPI") {
		t.Errorf("a Python engine was offered for a bundle: %v", err)
	}
	if _, err := Builds("llamacpp", "../b1", "amd64"); err == nil {
		t.Error("a version that is a path was accepted")
	}
}