llmgw/internal/agent/agent_test.go

package agent

import (
	"context"
	"crypto/tls"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
)

// statusRecorder captures what the control plane hears about shards.
type statusRecorder struct{ ch chan status }

type status struct {
	replicaID, shardID string
	state              supervisor.State
	err                string
}

func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} }

func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
	select {
	case r.ch <- status{replicaID, shardID, st, errMsg}:
	default:
	}
}

func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
	t.Helper()
	deadline := time.After(limit)
	var seen []status
	for {
		select {
		case s := <-r.ch:
			seen = append(seen, s)
			if pred(s) {
				return s
			}
		case <-deadline:
			t.Fatalf("no matching status within %v; saw %+v", limit, seen)
			return status{}
		}
	}
}

// fakeEngineScript writes a script that stands in for vLLM.
//
// It parses the real --port flag the vLLM adapter generates rather than being
// told out of band, so the adapter's argument construction is exercised too.
func fakeEngineScript(t *testing.T, body string) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "fake-vllm")
	script := `#!/bin/sh
port=""
while [ $# -gt 0 ]; do
  case "$1" in
    --port) port="$2"; shift 2 ;;
    *) shift ;;
  esac
done
` + body + "\n"
	if err := os.WriteFile(path, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}
	return path
}

// servingEngine answers /health and /v1/models, which is what readiness needs.
func servingEngine(t *testing.T) string {
	t.Helper()
	py, err := exec.LookPath("python3")
	if err != nil {
		t.Skip("python3 is required for the fake engine")
	}
	return fakeEngineScript(t, `exec `+py+` -c '
import http.server, json, sys
port = int(sys.argv[1])
class H(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/health":
            self.send_response(200); self.end_headers(); return
        if self.path == "/v1/models":
            b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode()
            self.send_response(200)
            self.send_header("content-type","application/json")
            self.send_header("content-length", str(len(b)))
            self.end_headers(); self.wfile.write(b); return
        self.send_response(404); self.end_headers()
    def log_message(self, *a): pass
http.server.HTTPServer(("127.0.0.1", port), H).serve_forever()
' "$port"`)
}

type harness struct {
	svc *agentsvc.Service
	inv *inventory.Store
	rec *statusRecorder
	srv *httptest.Server
	ag  *Agent
}

// start wires a real agent to a real service over h2c, so this exercises the
// protocol rather than two structs calling each other.
func start(t *testing.T, cfgFns ...func(*Config)) *harness {
	t.Helper()

	inv := inventory.NewStore()
	rec := newRecorder()
	svc := agentsvc.New(inv, rec, slog.New(slog.DiscardHandler))

	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)
	// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
	srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))

	cfg := DefaultConfig()
	cfg.NodeID = "n1"
	cfg.Server = srv.URL
	cfg.MgmtAddr = "127.0.0.1"
	cfg.HeartbeatInterval = 50 * time.Millisecond
	cfg.ReconnectBase = 20 * time.Millisecond
	cfg.ReconnectMax = 100 * time.Millisecond
	cfg.Discovery.SkipGPUs = true
	cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
	cfg.Superv = supervisor.Config{
		WorkDir: t.TempDir(), LogLines: 64,
		StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
	}
	for _, f := range cfgFns {
		f(&cfg)
	}

	client := &http.Client{Transport: &http2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, addr)
		},
	}}

	ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() { _ = ag.Run(ctx) }()

	h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag}
	t.Cleanup(func() {
		cancel()
		ag.Supervisor().StopAll(context.Background())
		srv.Close()
	})
	h.waitRegistered(t)
	return h
}

func (h *harness) waitRegistered(t *testing.T) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("the agent never registered")
}

// shardCommand builds a single-node vLLM shard on a free port.
func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand {
	t.Helper()
	return reconciler.ShardCommand{
		ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1",
		Spec: engines.ShardSpec{
			ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1,
			GPUIndices: []int{0}, HTTPPort: freePort(t),
			Profile: engines.Profile{
				Name: "p", Engine: "vllm", TP: 1,
				ModelPath: "/models/x", ServedModelName: "test-model",
			},
		},
		StartupTimeout: 15 * time.Second,
	}
}

func freePort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func TestRegistrationCarriesInventory(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Hostname = "gpu-a-01"
		c.Labels = map[string]string{"rack": "r12"}
		c.Fabrics = []inventory.FabricAttachment{
			{Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200},
		}
	})

	n, ok := h.inv.Get("n1")
	if !ok {
		t.Fatal("node missing from inventory")
	}
	if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" {
		t.Errorf("identity wrong: %+v", n)
	}
	if n.Labels["rack"] != "r12" {
		t.Errorf("labels lost: %v", n.Labels)
	}
	// The fabric label cannot be inferred from the wire, so losing it in
	// transit would silently disable multi-node placement.
	if !n.HasFabric("ib-a") {
		t.Fatalf("fabric attachment lost: %+v", n.Fabrics)
	}
	if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) {
		t.Errorf("HCAs=%v", got)
	}
}

func TestHeartbeatKeepsNodeAlive(t *testing.T) {
	h := start(t)

	n, _ := h.inv.Get("n1")
	first := n.LastSeen

	time.Sleep(200 * time.Millisecond)
	n, _ = h.inv.Get("n1")
	if !n.LastSeen.After(first) {
		t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost")
	}
}

func TestStartAndStopShardEndToEnd(t *testing.T) {
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-a")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}

	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Ready
	})
	if got.replicaID != "rep-1" {
		t.Errorf("replica id lost in transit: %q", got.replicaID)
	}

	if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Stopped
	})
}

func TestFailedShardCarriesDiagnosis(t *testing.T) {
	// When a model fails to load the engine's output is the only diagnosis, so
	// it has to survive the trip to the control plane.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`)
	})

	if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-bad" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("the failure carried no error message")
	}
}

func TestEngineThatNeverBecomesReady(t *testing.T) {
	// A process that runs but never serves must time out rather than hang the
	// replica forever.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})

	cmd := shardCommand(t, "shard-stuck")
	cmd.StartupTimeout = 500 * time.Millisecond
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 20*time.Second, func(s status) bool {
		return s.shardID == "shard-stuck" && s.state == supervisor.Failed
	})
}

func TestValidationRejectsBeforeLaunch(t *testing.T) {
	// An engine that refuses a profile does so after minutes of loading;
	// catching it at dispatch costs nothing.
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = servingEngine(t)
	})

	cmd := shardCommand(t, "shard-invalid")
	cmd.Spec.Profile.ModelPath = "" // validation requires one

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	got := h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-invalid" && s.state == supervisor.Failed
	})
	if got.err == "" {
		t.Error("validation failure carried no reason")
	}
}

func TestDispatchToUnconnectedNode(t *testing.T) {
	h := start(t)

	err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"})
	if err == nil {
		t.Fatal("dispatching to an unconnected node succeeded")
	}
	if !errors.Is(err, agentsvc.ErrNotConnected) {
		t.Errorf("err=%v, expected ErrNotConnected", err)
	}
}

func TestAgentReconnects(t *testing.T) {
	// A control plane restart must not require touching every GPU node.
	h := start(t)
	h.srv.CloseClientConnections()

	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if len(h.svc.ConnectedNodes()) == 1 {
			if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil {
				t.Fatalf("the reconnected stream does not accept commands: %v", err)
			}
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the agent never reconnected")
}

func TestIdempotentStart(t *testing.T) {
	// The control plane retries commands; a repeat must not start a second
	// process (plan §5.6).
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`)
	})

	cmd := shardCommand(t, "shard-idem")
	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-idem" && s.state == supervisor.Starting
	})

	sup := h.ag.Supervisor()
	sh, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	pid := sh.Status().PID

	if err := h.svc.StartShard(t.Context(), cmd); err != nil {
		t.Fatal(err)
	}
	time.Sleep(400 * time.Millisecond)

	sh2, err := sup.Get("shard-idem")
	if err != nil {
		t.Fatal(err)
	}
	if sh2.Status().PID != pid {
		t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
	}
}