llmgw/test/e2e/harness_test.go

// Package e2e runs the router against real sim-engines over real HTTP.
//
// These tests exist to check the phase 1 exit criterion from the plan, so they
// deliberately use the whole stack: no fakes, no in-process shortcuts beyond
// httptest listeners.
package e2e

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router"
	"gitea.mynodes.xyz/ilker/llmgw/internal/sim"
)

const testModel = "test-model"

type engine struct {
	eng *sim.Engine
	srv *httptest.Server
}

type harness struct {
	engines []engine
	router  *httptest.Server
	rt      *router.Router
}

// simConfig is fast enough for tests while still exercising queueing.
func simConfig(name string, seed int64) sim.Config {
	c := sim.Default()
	c.Model = testModel
	c.Seed = seed
	c.Tick = time.Millisecond
	c.MaxNumSeqs = 512
	c.NumGPUBlocks = 200000
	c.PrefillTokensPerSec = 5e6
	c.DecodeTokensPerSec = 2000
	c.BatchContention = 0
	_ = name
	return c
}

// engineOpt tunes the sim-engines a harness starts.
type engineOpt func(*sim.Config)

// engineCapacity is the batch capacity used by withLoadedEngines, chosen to sit
// in the range a real vLLM deployment runs at. Making it much smaller would
// shift the bottleneck to slot starvation inside the engine, which admission
// cannot do anything about — the test would then be measuring the wrong thing.
const engineCapacity = 64

// withLoadedEngines sizes the engines so that a few hundred concurrent requests
// genuinely overflow them, which is what the priority isolation test needs.
func withLoadedEngines() func(*config.Server) {
	return nil // marker; see start()
}

// tightPrefixCache shrinks the engines' prefix cache so the working set does
// not fit in any one of them.
//
// That condition is the whole reason tier 1 exists. With a cache large enough
// for every prompt, scattering is free — every replica ends up holding
// everything — and prefix routing can only add constraint. It is worth knowing
// the feature does nothing there rather than assuming it helps everywhere.
var tightPrefixCache = func(c *sim.Config) { c.PrefixCacheBlocks = 96 }

// start brings up n sim-engines behind a router.
//
// A nil tune function is the "small engines" marker, which keeps the call sites
// readable without a second variadic parameter.
func start(t *testing.T, n int, tune ...func(*config.Server)) *harness {
	t.Helper()
	return startTuned(t, n, nil, tune...)
}

// startTuned is start with the sim-engines themselves adjusted.
func startTuned(t *testing.T, n int, engineOpts []engineOpt, tune ...func(*config.Server)) *harness {
	t.Helper()

	small := false
	kept := tune[:0]
	for _, f := range tune {
		if f == nil {
			small = true
			continue
		}
		kept = append(kept, f)
	}
	tune = kept

	h := &harness{}
	cfg := config.Default()
	cfg.Router.ScrapeInterval = 200 * time.Millisecond
	cfg.Router.ScrapeTimeout = 150 * time.Millisecond
	cfg.Router.StaleAfter = 2 * time.Second

	for i := range n {
		sc := simConfig(fmt.Sprintf("sim-%d", i), int64(i+1))
		if small {
			sc.MaxNumSeqs = engineCapacity
			sc.DecodeTokensPerSec = 800
			sc.BatchContention = 0.005
		}
		for _, opt := range engineOpts {
			opt(&sc)
		}
		e := sim.New(sc)
		srv := httptest.NewServer(sim.NewServer(e, slog.New(slog.DiscardHandler)))
		h.engines = append(h.engines, engine{eng: e, srv: srv})
		cfg.Endpoints = append(cfg.Endpoints, config.Endpoint{
			ID:     fmt.Sprintf("e%d", i),
			URL:    srv.URL,
			Models: []string{testModel},
			Pool:   "test",
			Weight: 1,
		})
	}
	for _, f := range tune {
		f(&cfg)
	}

	rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	rt.Start(ctx)

	h.rt = rt
	h.router = httptest.NewServer(rt.Handler())

	t.Cleanup(func() {
		h.router.Close()
		cancel()
		rt.Stop()
		for _, e := range h.engines {
			e.srv.Close()
			e.eng.Close()
		}
	})

	// Let the first scrape land so scheduling starts from real data.
	h.waitForFreshMetrics(t)
	return h
}

func (h *harness) waitForFreshMetrics(t *testing.T) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		fresh := 0
		for _, s := range h.rt.Registry().Snapshots(testModel) {
			if !s.Stale {
				fresh++
			}
		}
		if fresh == len(h.engines) {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("endpoint metrics never became fresh")
}

// running counts requests actively generating across all engines.
func (h *harness) running() int {
	n := 0
	for _, e := range h.engines {
		n += e.eng.Stats().Running
	}
	return n
}

func (h *harness) totalRunning() int {
	n := 0
	for _, e := range h.engines {
		s := e.eng.Stats()
		n += s.Running + s.Waiting
	}
	return n
}

// chatBody builds a streaming chat request.
func chatBody(prompt string, maxTokens int, stream bool) []byte {
	b, _ := json.Marshal(map[string]any{
		"model":      testModel,
		"messages":   []map[string]string{{"role": "user", "content": prompt}},
		"max_tokens": maxTokens,
		"stream":     stream,
	})
	return b
}

// streamRequest sends a streaming request and counts the content chunks.
func streamRequest(ctx context.Context, url string, body []byte, hdr map[string]string) (chunks int, ttfb time.Duration, err error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return 0, 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	for k, v := range hdr {
		req.Header.Set(k, v)
	}

	start := time.Now()
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return 0, 0, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<12))
		return 0, 0, fmt.Errorf("status %d: %s", resp.StatusCode, b)
	}

	br := bufio.NewReader(resp.Body)
	for {
		line, rerr := br.ReadString('\n')
		if strings.HasPrefix(line, "data: ") {
			if chunks == 0 {
				ttfb = time.Since(start)
			}
			if !strings.Contains(line, "[DONE]") {
				chunks++
			}
		}
		if rerr != nil {
			if rerr == io.EOF {
				return chunks, ttfb, nil
			}
			return chunks, ttfb, rerr
		}
	}
}

// envInt reads an integer from the environment, falling back to def.
func envInt(key string, def int) int {
	v := os.Getenv(key)
	if v == "" {
		return def
	}
	var n int
	if _, err := fmt.Sscanf(v, "%d", &n); err != nil || n <= 0 {
		return def
	}
	return n
}