177 lines
6.3 KiB
Go
177 lines
6.3 KiB
Go
package proxy_test
|
|
|
|
// v2.3 task 01: control-plane requests. A client that manages its own slots (Boxmaker) polls
|
|
// /slots, reads /props, tokenizes and steers a running completion through
|
|
// /v1/chat/completions/control. Those calls follow the route's lease like any other request but
|
|
// must never wait for, or take, a slot: /control is sent while the client's own stream holds one.
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// controlClient gives every control call a short deadline: a call that queues behind a full host
|
|
// is the bug, and it must fail the test rather than hang it.
|
|
var controlClient = &http.Client{Timeout: 2 * time.Second}
|
|
|
|
func (r *rig) do(method, path, body string) *http.Response {
|
|
r.t.Helper()
|
|
var rd *strings.Reader
|
|
if body != "" {
|
|
rd = strings.NewReader(body)
|
|
}
|
|
var req *http.Request
|
|
var err error
|
|
if rd != nil {
|
|
req, err = http.NewRequest(method, r.front.URL+path, rd)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
} else {
|
|
req, err = http.NewRequest(method, r.front.URL+path, nil)
|
|
}
|
|
if err != nil {
|
|
r.t.Fatal(err)
|
|
}
|
|
resp, err := controlClient.Do(req)
|
|
if err != nil {
|
|
r.t.Fatalf("%s %s: %v", method, path, err)
|
|
}
|
|
return resp
|
|
}
|
|
|
|
func (r *rig) rows(route string) int64 {
|
|
r.t.Helper()
|
|
counts, err := r.store.StatusCounts(time.Time{})
|
|
if err != nil {
|
|
r.t.Fatal(err)
|
|
}
|
|
var n int64
|
|
for _, c := range counts {
|
|
if c.Route == route {
|
|
n += c.Count
|
|
}
|
|
}
|
|
return n
|
|
}
|
|
|
|
// A GET names its model in the query string: /slots?model=alpha-only must reach the host that
|
|
// has alpha-only loaded, not whichever host the route's default model would pick.
|
|
func TestGetModelComesFromTheQuery(t *testing.T) {
|
|
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
|
|
r := newRig(t, twoHosts, alpha, beta)
|
|
|
|
resp := r.do(http.MethodGet, "/r/slots?model=alpha-only", "")
|
|
drain(resp)
|
|
if resp.StatusCode != 200 || resp.Header.Get("X-Crossbar-Host") != "alpha" {
|
|
t.Fatalf("GET /r/slots?model=alpha-only: %d on %q, want 200 on alpha", resp.StatusCode, resp.Header.Get("X-Crossbar-Host"))
|
|
}
|
|
if got := alpha.lastReq(); got.method != "GET" || got.path != "/slots?model=alpha-only" {
|
|
t.Errorf("alpha saw %s %s, want GET /slots?model=alpha-only", got.method, got.path)
|
|
}
|
|
// The same for beta-only, so a lucky default cannot pass the test.
|
|
resp = r.do(http.MethodGet, "/r/slots?model=beta-only", "")
|
|
drain(resp)
|
|
if resp.Header.Get("X-Crossbar-Host") != "beta" {
|
|
t.Errorf("GET /r/slots?model=beta-only went to %q, want beta", resp.Header.Get("X-Crossbar-Host"))
|
|
}
|
|
}
|
|
|
|
// /slots and /tokenize are proxied; the per-slot actions under /slots/ (save, restore, erase)
|
|
// are not.
|
|
func TestControlPathsAllowed(t *testing.T) {
|
|
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
|
|
r := newRig(t, twoHosts, alpha, beta)
|
|
for _, tc := range []struct {
|
|
method, path, body string
|
|
want int
|
|
}{
|
|
{http.MethodGet, "/r/slots", "", 200},
|
|
{http.MethodGet, "/r/slots?model=shared", "", 200},
|
|
{http.MethodPost, "/r/tokenize", `{"model":"shared","content":"hello"}`, 200},
|
|
{http.MethodPost, "/r/v1/chat/completions/control", `{"id":"chatcmpl-1","action":"reasoning_end","model":"shared"}`, 200},
|
|
{http.MethodGet, "/r/slots/0", "", 404},
|
|
{http.MethodPost, "/r/slots/0?action=erase", "", 404},
|
|
{http.MethodPost, "/r/slots/0?action=save", `{"filename":"x"}`, 404},
|
|
} {
|
|
resp := r.do(tc.method, tc.path, tc.body)
|
|
body := drain(resp)
|
|
if resp.StatusCode != tc.want {
|
|
t.Errorf("%s %s: %d %s, want %d", tc.method, tc.path, resp.StatusCode, body, tc.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// With every slot on both hosts taken and the queue full, control-plane calls still go straight
|
|
// through: no 503, no wait, no slot taken, no accounting row.
|
|
func TestControlRequestsNeverTakeASlot(t *testing.T) {
|
|
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
|
|
r := newRig(t, twoHosts, alpha, beta)
|
|
|
|
// Take every "shared" slot (parallel 2 on each host) and the one queue place per host.
|
|
var releases []func()
|
|
for _, host := range []string{"alpha", "beta"} {
|
|
for i := 0; i < 2; i++ {
|
|
rel, _, err := r.lim.Acquire(context.Background(), host, "shared")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
releases = append(releases, rel)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
go func() { _, _, _ = r.lim.Acquire(ctx, host, "shared") }()
|
|
waitUntil(t, func() bool { return r.lim.Queued(host, "shared") == 1 })
|
|
}
|
|
defer func() {
|
|
for _, rel := range releases {
|
|
rel()
|
|
}
|
|
}()
|
|
|
|
for _, tc := range []struct{ method, path, body string }{
|
|
{http.MethodGet, "/r/slots?model=shared", ""},
|
|
{http.MethodGet, "/r/props?model=shared", ""},
|
|
{http.MethodHead, "/r/props?model=shared", ""},
|
|
{http.MethodGet, "/r/v1/models", ""},
|
|
{http.MethodPost, "/r/tokenize", `{"model":"shared","content":"hello"}`},
|
|
{http.MethodPost, "/r/v1/chat/completions/control", `{"id":"chatcmpl-1","action":"reasoning_end","model":"shared"}`},
|
|
} {
|
|
resp := r.do(tc.method, tc.path, tc.body)
|
|
body := drain(resp)
|
|
if resp.StatusCode != 200 {
|
|
t.Errorf("%s %s with the host full: %d %s, want 200", tc.method, tc.path, resp.StatusCode, body)
|
|
}
|
|
}
|
|
for _, host := range []string{"alpha", "beta"} {
|
|
if n := r.lim.InFlight(host, "shared"); n != 2 {
|
|
t.Errorf("%s in flight = %d after control calls, want 2 (control takes no slot)", host, n)
|
|
}
|
|
}
|
|
time.Sleep(100 * time.Millisecond) // a row is written after the answer; give a stray one time to land
|
|
if n := r.rows("r"); n != 0 {
|
|
t.Errorf("control calls wrote %d accounting rows, want 0", n)
|
|
}
|
|
|
|
// A chat completion on the same full route still queues or is refused as before: the bypass
|
|
// is for control calls only.
|
|
resp := r.do(http.MethodPost, "/r/v1/chat/completions", conversation(1, 1))
|
|
drain(resp)
|
|
if resp.StatusCode != http.StatusServiceUnavailable {
|
|
t.Errorf("chat on a full route: %d, want 503 (queue full)", resp.StatusCode)
|
|
}
|
|
}
|
|
|
|
// A chat completion is not a control call just because its path starts the same way.
|
|
func TestChatIsNotControl(t *testing.T) {
|
|
alpha, beta := newUpstream(t, "alpha"), newUpstream(t, "beta")
|
|
r := newRig(t, twoHosts, alpha, beta)
|
|
resp := r.do(http.MethodPost, "/r/v1/chat/completions", conversation(1, 1))
|
|
drain(resp)
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("chat: %d", resp.StatusCode)
|
|
}
|
|
waitUntil(t, func() bool { return r.rows("r") == 1 }) // the row lands just after the answer
|
|
}
|