llmgw/test/e2e/ondemand_test.go

package e2e

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"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/router"
)

// onDemandHarness runs a control plane whose only model loads when asked for.
//
// End to end rather than in the reconciler's own tests, because on-demand
// loading is four components agreeing: the router notices nothing is serving,
// the control plane decides the model may be started, the reconciler places it,
// and the endpoint list carries it back. Every one of those is tested alone.
// The wiring between them is where this session's real bugs have been — a
// collector registered after the registry was built, a readiness check that
// never ran — and none of them failed a unit test.
func onDemandHarness(t *testing.T, idle time.Duration) *controlHarness {
	t.Helper()

	cat, err := catalog.New([]catalog.Model{{
		Name: testModel, Source: "/models/test",
		Scaling:  catalog.Scaling{Min: 0, Max: 1, IdleTimeout: idle},
		Profiles: []catalog.Profile{{Name: "single", GPUClass: "h100-80g", TP: 1, Weight: 1}},
	}})
	if err != nil {
		t.Fatal(err)
	}

	cfg := config.Default()
	cfg.Router.ScrapeInterval = 200 * time.Millisecond
	cfg.Router.StaleAfter = 3 * time.Second
	cfg.Router.Admission.Enabled = false
	cfg.Router.ActivationWait = 10 * time.Second

	rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	disp := newSimDispatcher(t)

	ccfg := control.DefaultConfig()
	ccfg.Reconciler.Interval = 100 * time.Millisecond
	ccfg.Reconciler.BaseBackoff = 200 * time.Millisecond
	ccfg.Reconciler.MaxBackoff = 500 * time.Millisecond
	ccfg.Reconciler.UnreachableAfter = 300 * time.Millisecond
	ccfg.Reconciler.LostAfter = 600 * time.Millisecond
	ccfg.Reconciler.SettleFor = 400 * time.Millisecond
	ccfg.Reconciler.IdleTimeout = idle
	ccfg.PublishInterval = 100 * time.Millisecond

	ctl := control.New(ccfg, catalog.NewLive(cat), disp, rt, slog.New(slog.DiscardHandler))
	disp.onStatus = ctl.Reconciler().OnShardStatus
	rt.SetActivator(ctl.Reconciler())

	ctl.Inventory().Upsert(inventory.Node{
		ID: "n0", MgmtAddr: "127.0.0.1", State: inventory.Ready,
		GPUs: []discovery.GPU{{
			Index: 0, Name: "NVIDIA H100 80GB HBM3",
			MemoryTotalMB: 81559, ComputeCapability: "9.0",
		}},
	})

	ctx, cancel := context.WithCancel(context.Background())
	rt.Start(ctx)
	ctl.Start(ctx)

	h := &controlHarness{
		router: httptest.NewServer(rt.Handler()),
		rt:     rt, ctl: ctl, disp: disp, inv: ctl.Inventory(),
	}
	// The node has to keep reporting. UnreachableAfter is 300ms in this
	// harness and every wait below is longer than that, so a single heartbeat
	// would have the node go quiet and nothing would be placed at all — which
	// looks exactly like on-demand loading failing.
	beat, stopBeat := context.WithCancel(ctx)
	go func() {
		t := time.NewTicker(50 * time.Millisecond)
		defer t.Stop()
		for {
			select {
			case <-beat.Done():
				return
			case <-t.C:
				h2 := h
				if h2 != nil {
					h2.heartbeat("n0")
				}
			}
		}
	}()

	t.Cleanup(func() {
		stopBeat()
		h.router.Close()
		ctl.Close()
		cancel()
		rt.Stop()
	})
	return h
}

func TestARequestLoadsAnOnDemandModelAndIsServed(t *testing.T) {
	h := onDemandHarness(t, time.Hour)
	// Settled, and still nothing running: that is the feature, not a fault.
	time.Sleep(700 * time.Millisecond)
	if n := h.disp.liveShards(); n != 0 {
		t.Fatalf("%d shards running before anyone asked for the model", n)
	}

	// One ordinary request, which should start the model and then be answered
	// by it rather than refused.
	resp, body := chat(t, h.router.URL, testModel)
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status %d: %s", resp.StatusCode, body)
	}
	if h.disp.liveShards() == 0 {
		t.Error("the request was served but nothing is running")
	}
}

func TestAnIdleOnDemandModelIsStoppedEndToEnd(t *testing.T) {
	// Short enough to observe, long enough that the model is not evicted while
	// the first request is still in flight.
	h := onDemandHarness(t, 2*time.Second)
	time.Sleep(700 * time.Millisecond)

	if resp, body := chat(t, h.router.URL, testModel); resp.StatusCode != http.StatusOK {
		t.Fatalf("status %d: %s", resp.StatusCode, body)
	}
	if h.disp.liveShards() == 0 {
		t.Fatal("the model never started")
	}

	// Left alone past its idle timeout, with the node still heartbeating so
	// this is eviction rather than node loss.
	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if h.disp.liveShards() == 0 {
			return
		}
		time.Sleep(200 * time.Millisecond)
	}
	t.Error("an idle on-demand model was never stopped")
}

func TestAnUnknownModelIsStillAnErrorWithOnDemandEnabled(t *testing.T) {
	// On-demand must not turn every typo into "loading", or nothing is ever
	// wrong and the useful error disappears.
	h := onDemandHarness(t, time.Hour)
	time.Sleep(700 * time.Millisecond)

	resp, body := chat(t, h.router.URL, "a-model-nobody-configured")
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("status %d: %s", resp.StatusCode, body)
	}
	if resp.Header.Get("Retry-After") != "" {
		t.Error("an unknown model was presented as retryable")
	}
	// With an on-demand-only cluster and nothing loaded, naming what is served
	// means naming nothing — and saying so plainly is the honest answer.
	if !strings.Contains(body, "no model is currently being served") {
		t.Errorf("the useful error was lost: %s", body)
	}
}

// chat sends one completion and returns the response and its body.
func chat(t *testing.T, base, model string) (*http.Response, string) {
	t.Helper()
	body := fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":8}`, model)
	req, err := http.NewRequest(http.MethodPost, base+"/v1/chat/completions", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	buf := make([]byte, 4096)
	n, _ := resp.Body.Read(buf)
	return resp, string(buf[:n])
}