llmgw/internal/control/adminapi/api_test.go

package adminapi

import (
	"context"
	"encoding/json"
	"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)
	}
	ctl := control.New(control.DefaultConfig(), 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)
		}
	}
}