llmgw/internal/control/adminapi/api_test.go

package adminapi

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// nullSink discards published endpoints; these tests are about the API.
type nullSink struct{}

func (nullSink) SetEndpoints(context.Context, []endpoint.Info) {}

// nullDispatcher accepts every command.
type nullDispatcher struct{}

func (nullDispatcher) StartShard(context.Context, reconciler.ShardCommand) error { return nil }
func (nullDispatcher) StopShard(context.Context, string, string) error           { return nil }
func (nullDispatcher) DrainShard(context.Context, string, string) error          { return nil }

func setup(t *testing.T, models []catalog.Model, nodes ...inventory.Node) (*API, *control.Controller) {
	t.Helper()
	cat, err := catalog.New(models)
	if err != nil {
		t.Fatal(err)
	}
	cfg := control.DefaultConfig()
	// No startup settling window: these tests build the whole fleet before
	// anything runs, which is the state settling waits for. Leaving it on would
	// make every test here assert that placement had not happened yet.
	cfg.Reconciler.SettleFor = -1
	ctl := control.New(cfg, catalog.NewLive(cat), nullDispatcher{}, nullSink{}, slog.New(slog.DiscardHandler))
	for _, n := range nodes {
		ctl.Inventory().Upsert(n)
	}
	return New(ctl), ctl
}

func gpuNode(id string, n int, fabrics ...string) inventory.Node {
	nd := inventory.Node{ID: id, State: inventory.Ready, AgentVersion: "v1"}
	for i := range n {
		nd.GPUs = append(nd.GPUs, discovery.GPU{
			Index: i, Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559,
		})
	}
	for _, f := range fabrics {
		nd.Fabrics = append(nd.Fabrics, inventory.FabricAttachment{Fabric: f, HCAs: []string{"mlx5_0"}})
	}
	return nd
}

func testModels() []catalog.Model {
	return []catalog.Model{{
		Name: "llama", Aliases: []string{"l"}, Source: "/m",
		Scaling:  catalog.Scaling{Min: 2, Max: 4},
		Profiles: []catalog.Profile{{Name: "tp8", GPUClass: "h100-80g", TP: 8}},
	}}
}

func call(t *testing.T, a *API, method, path string, body string, out any) int {
	t.Helper()
	mux := http.NewServeMux()
	a.Register(mux)

	var r *http.Request
	if body != "" {
		r = httptest.NewRequest(method, path, strings.NewReader(body))
		r.Header.Set("Content-Type", "application/json")
	} else {
		r = httptest.NewRequest(method, path, nil)
	}
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, r)

	if out != nil && rec.Code == http.StatusOK {
		if err := json.Unmarshal(rec.Body.Bytes(), out); err != nil {
			t.Fatalf("decoding %s: %v (%s)", path, err, rec.Body.String())
		}
	}
	return rec.Code
}

func TestListNodes(t *testing.T) {
	a, _ := setup(t, testModels(), gpuNode("n1", 8, "ib-a"))

	var nodes []NodeView
	if code := call(t, a, "GET", "/api/v1/nodes", "", &nodes); code != 200 {
		t.Fatalf("status %d", code)
	}
	if len(nodes) != 1 {
		t.Fatalf("got %d nodes", len(nodes))
	}
	n := nodes[0]
	if n.ID != "n1" || n.State != "ready" {
		t.Errorf("%+v", n)
	}
	if n.GPUTotal != 8 || n.GPUFree != 8 {
		t.Errorf("gpus %d/%d", n.GPUFree, n.GPUTotal)
	}
	if len(n.Classes) != 1 || n.Classes[0] != "h100-80g" {
		t.Errorf("classes=%v", n.Classes)
	}
	if len(n.Fabrics) != 1 || n.Fabrics[0] != "ib-a" {
		t.Errorf("fabrics=%v", n.Fabrics)
	}
}

func TestUnhealthyGPUsAreSurfaced(t *testing.T) {
	// Usually the answer when a node looks idle but nothing schedules onto it.
	n := gpuNode("n1", 4)
	n.GPUs[1].Health.Problems = []string{"3 uncorrected ECC errors"}
	a, _ := setup(t, testModels(), n)

	var nodes []NodeView
	call(t, a, "GET", "/api/v1/nodes", "", &nodes)

	if len(nodes[0].Unhealthy) != 1 {
		t.Fatalf("unhealthy=%v", nodes[0].Unhealthy)
	}
	if !strings.Contains(nodes[0].Unhealthy[0], "ECC") {
		t.Errorf("the reason was lost: %q", nodes[0].Unhealthy[0])
	}
	// And such a GPU is not counted as free.
	if nodes[0].GPUFree != 3 {
		t.Errorf("free=%d, an unhealthy GPU was counted", nodes[0].GPUFree)
	}
}

func TestListModels(t *testing.T) {
	a, _ := setup(t, testModels(), gpuNode("n1", 8))

	var models []ModelView
	call(t, a, "GET", "/api/v1/models", "", &models)

	if len(models) != 1 {
		t.Fatalf("got %d models", len(models))
	}
	m := models[0]
	if m.Name != "llama" || m.Min != 2 || m.Max != 4 {
		t.Errorf("%+v", m)
	}
	if len(m.Aliases) != 1 || m.Aliases[0] != "l" {
		t.Errorf("aliases=%v", m.Aliases)
	}
	if len(m.Profiles) != 1 || m.Profiles[0] != "tp8" {
		t.Errorf("profiles=%v", m.Profiles)
	}
}

func TestExplainPlacementForUnplaceableModel(t *testing.T) {
	// Four GPUs, a profile that needs eight.
	a, ctl := setup(t, testModels(), gpuNode("n1", 4))
	ctl.Reconciler().Reconcile(t.Context())

	var v ExplainView
	if code := call(t, a, "GET", "/api/v1/explain-placement?model=llama", "", &v); code != 200 {
		t.Fatalf("status %d", code)
	}
	if v.Placed {
		t.Error("an unplaceable model was reported as placed")
	}
	if len(v.Rejections) == 0 {
		t.Fatal("no rejections were recorded; the operator learns nothing")
	}
	found := false
	for _, r := range v.Rejections {
		if strings.Contains(r.Reason, "needs 8") {
			found = true
		}
	}
	if !found {
		t.Errorf("the reason does not say what was missing: %+v", v.Rejections)
	}
}

func TestExplainPlacementRequiresAKnownModel(t *testing.T) {
	a, _ := setup(t, testModels(), gpuNode("n1", 8))

	if code := call(t, a, "GET", "/api/v1/explain-placement", "", nil); code != 400 {
		t.Errorf("missing model gave %d, expected 400", code)
	}
	if code := call(t, a, "GET", "/api/v1/explain-placement?model=nope", "", nil); code != 404 {
		t.Errorf("unknown model gave %d, expected 404", code)
	}
}

func TestSetNodeState(t *testing.T) {
	a, ctl := setup(t, testModels(), gpuNode("n1", 8))

	for _, c := range []struct {
		body string
		want inventory.NodeState
	}{
		{`{"state":"cordon"}`, inventory.Cordoned},
		{`{"state":"drain"}`, inventory.Draining},
		{`{"state":"uncordon"}`, inventory.Ready},
		{`{"state":"cordoned"}`, inventory.Cordoned},
	} {
		if code := call(t, a, "POST", "/api/v1/nodes/n1/state", c.body, nil); code != 200 {
			t.Fatalf("%s gave %d", c.body, code)
		}
		n, _ := ctl.Inventory().Get("n1")
		if n.State != c.want {
			t.Errorf("%s left the node %s, expected %s", c.body, n.State, c.want)
		}
	}
}

func TestSetNodeStateRejectsNonsense(t *testing.T) {
	a, _ := setup(t, testModels(), gpuNode("n1", 8))

	if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `{"state":"banana"}`, nil); code != 400 {
		t.Errorf("an unknown state gave %d, expected 400", code)
	}
	if code := call(t, a, "POST", "/api/v1/nodes/nope/state", `{"state":"cordon"}`, nil); code != 404 {
		t.Errorf("an unknown node gave %d, expected 404", code)
	}
	if code := call(t, a, "POST", "/api/v1/nodes/n1/state", `not json`, nil); code != 400 {
		t.Errorf("malformed body gave %d, expected 400", code)
	}
}

func TestEmptyClusterReturnsArraysNotNull(t *testing.T) {
	// A JSON null would make every consumer special-case it.
	a, _ := setup(t, testModels())

	for _, path := range []string{"/api/v1/nodes", "/api/v1/replicas"} {
		mux := http.NewServeMux()
		a.Register(mux)
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, httptest.NewRequest("GET", path, nil))

		if got := strings.TrimSpace(rec.Body.String()); got != "[]" {
			t.Errorf("%s returned %q, expected []", path, got)
		}
	}
}

type fakeModels struct {
	added   []catalog.Model
	actors  []string
	removed []string
	err     error
}

func (f *fakeModels) Add(_ context.Context, m catalog.Model, by string) error {
	if f.err != nil {
		return f.err
	}
	f.added = append(f.added, m)
	f.actors = append(f.actors, by)
	return nil
}

func (f *fakeModels) Remove(_ context.Context, name string) error {
	if f.err != nil {
		return f.err
	}
	f.removed = append(f.removed, name)
	return nil
}

func TestAddingAModelOverTheAPI(t *testing.T) {
	a, _ := setup(t, testModels())
	fm := &fakeModels{}
	a.SetModels(fm)
	mux := http.NewServeMux()
	a.Register(mux)

	body := `{"name":"new","source":"/models/new","profiles":[{"name":"p","engine":"llamacpp"}]}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(body))
	req.Header.Set("X-Llmgw-Actor", "ilker")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if len(fm.added) != 1 || fm.added[0].Name != "new" {
		t.Fatalf("added = %+v", fm.added)
	}
	// Who made the change, kept so a catalog that grew over months can say
	// where each entry came from.
	if len(fm.actors) != 1 || fm.actors[0] != "ilker" {
		t.Errorf("actor = %v", fm.actors)
	}
}

func TestABadModelDefinitionIsTheCallersFault(t *testing.T) {
	// Every failure here is the definition being wrong — a collision, a missing
	// fallback, a profile that cannot be placed. None are server faults, and
	// returning 500 for them would send someone looking at the wrong logs.
	a, _ := setup(t, testModels())
	a.SetModels(&fakeModels{err: errors.New("model \"x\" collides")})
	mux := http.NewServeMux()
	a.Register(mux)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(`{"name":"x"}`))
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusBadRequest {
		t.Errorf("status %d, expected 400", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "collides") {
		t.Errorf("the reason was lost: %s", rec.Body.String())
	}
}

func TestWithoutAStoreTheAPIRefusesRatherThanForgets(t *testing.T) {
	// Accepting a model that vanishes on the next restart is worse than saying
	// it cannot be kept.
	a, _ := setup(t, testModels()) // SetModels never called
	mux := http.NewServeMux()
	a.Register(mux)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/models", strings.NewReader(`{"name":"x"}`))
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusNotImplemented {
		t.Errorf("status %d, expected 501", rec.Code)
	}
	if !strings.Contains(rec.Body.String(), "restart") {
		t.Errorf("the reason does not explain the consequence: %s", rec.Body.String())
	}
}

func TestRemovingAModelOverTheAPI(t *testing.T) {
	a, _ := setup(t, testModels())
	fm := &fakeModels{}
	a.SetModels(fm)
	mux := http.NewServeMux()
	a.Register(mux)

	req := httptest.NewRequest(http.MethodDelete, "/api/v1/models/gone", nil)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
	}
	if len(fm.removed) != 1 || fm.removed[0] != "gone" {
		t.Errorf("removed = %v", fm.removed)
	}
}