llmgw/internal/control/reconciler/stuck_test.go

package reconciler

import (
	"context"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

// TestAReplicaThatNeverComesUpIsEventuallyFailed pins down the hole found on a
// real cluster: two replicas sat in Starting for eight minutes past a two-minute
// startup timeout, holding GPUs and counting towards the model's replica count
// while serving nothing.
//
// Every other way out of Starting needs the agent to report something. When it
// never does, this is the only thing that moves.
func TestAReplicaThatNeverComesUpIsEventuallyFailed(t *testing.T) {
	m := testModel(1, catalog.Profile{
		Name: "tp4", GPUClass: "h100-80g", TP: 4, StartupTimeout: 2 * time.Minute,
	})
	h := newHarness(t, m, gpuNode("1", 8))
	ctx := context.Background()

	h.r.Reconcile(ctx)
	reps := h.r.Replicas()
	if len(reps) != 1 || reps[0].State != Starting {
		t.Fatalf("expected one starting replica, got %+v", reps)
	}
	// Captured rather than assumed: replica ids carry a random suffix so a
	// restarted control plane cannot collide with engines from its last life.
	stuckID := reps[0].ID

	// The agent says nothing at all, which is exactly the case.
	h.advance(90 * time.Second)
	h.r.Reconcile(ctx)
	if got := h.r.Replicas()[0].State; got != Starting {
		t.Fatalf("failed early, at %s: %s", 90*time.Second, got)
	}

	h.advance(60 * time.Second)
	h.r.Reconcile(ctx)

	// Failing it is not the outcome that matters; being unstuck is. The same
	// reconcile reaps the failure and places a replacement, which is the whole
	// point of failing it.
	reps = h.r.Replicas()
	for _, rep := range reps {
		if rep.ID == stuckID && rep.State == Starting {
			t.Fatal("the replica is still stuck in starting past its deadline")
		}
	}
	// The replacement waits out the model's backoff, which is the point of the
	// backoff: a model that cannot start should not be restarted in a loop.
	h.advance(10 * time.Minute)
	h.r.Reconcile(ctx)

	reps = h.r.Replicas()
	var replaced bool
	for _, rep := range reps {
		if rep.ID != stuckID {
			replaced = true
		}
	}
	if !replaced {
		t.Errorf("no replacement was ever placed; replicas are %+v", reps)
	}

	// And the shards of the stuck one were told to stop, or its GPUs stay
	// allocated to a replica nobody is waiting for any more.
	h.disp.mu.Lock()
	stopped := len(h.disp.stopped)
	h.disp.mu.Unlock()
	if stopped == 0 {
		t.Error("the stuck replica's shards were never stopped")
	}
}

func TestTheStuckReplicaSaysWhyItFailed(t *testing.T) {
	// Reaping removes the evidence, so the reason has to be recorded when the
	// deadline passes rather than found afterwards.
	m := testModel(1, catalog.Profile{
		Name: "tp4", GPUClass: "h100-80g", TP: 4, StartupTimeout: time.Minute,
	})
	h := newHarness(t, m, gpuNode("1", 8))
	h.r.Reconcile(context.Background())

	h.advance(2 * time.Minute)
	h.r.expireStuckReplicas()

	reps := h.r.Replicas()
	if len(reps) != 1 || reps[0].State != Failed {
		t.Fatalf("replicas: %+v", reps)
	}
	if !strings.Contains(reps[0].Err, "did not come up") {
		t.Errorf("the error does not say what happened: %q", reps[0].Err)
	}
}

func TestAReplicaThatIsLoadingIsNotFailedEarly(t *testing.T) {
	// Loading a large model legitimately takes a long time. The deadline is
	// what the profile says, and nothing shorter.
	m := testModel(1, catalog.Profile{
		Name: "tp4", GPUClass: "h100-80g", TP: 4, StartupTimeout: time.Hour,
	})
	h := newHarness(t, m, gpuNode("1", 8))
	ctx := context.Background()
	h.r.Reconcile(ctx)

	h.advance(30 * time.Minute)
	h.r.Reconcile(ctx)
	if got := h.r.Replicas()[0].State; got == Failed {
		t.Error("a replica still inside its startup timeout was failed")
	}
}

// TestTheFleetWinsOverAReplicaWeOnlyPlaced is the other half of the same
// incident. Replica ids come from the model name and a counter, so a control
// plane that restarts and places before its agents reconnect builds a different
// replica under the same id, on a freshly allocated port. Dropping the agent's
// report then leaves the cluster pointing at a port nothing listens on while the
// engine that is actually serving is treated as a stray.
func TestTheFleetWinsOverAReplicaWeOnlyPlaced(t *testing.T) {
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	ctx := context.Background()

	h.r.Reconcile(ctx)
	placed := h.r.Replicas()[0]
	if placed.State != Starting {
		t.Fatalf("state = %s", placed.State)
	}

	// An agent reports the same replica id, running on the port it was actually
	// started with.
	unclaimed := h.r.Adopt("1", []AdoptedShard{{
		NodeID: "1", ReplicaID: placed.ID, ShardID: placed.ID + "-r0",
		Rank: 0, GPUIndices: []int{0, 1, 2, 3}, HTTPPort: 20000,
		State: supervisor.Ready, Model: "llama", Profile: "tp4",
		Endpoint: "http://node1:20000",
	}})

	if len(unclaimed) != 0 {
		t.Fatalf("the running engine was disowned: %+v", unclaimed)
	}
	reps := h.r.Replicas()
	if len(reps) != 1 {
		t.Fatalf("replicas: %+v", reps)
	}
	if reps[0].Endpoint != "http://node1:20000" {
		t.Errorf("endpoint = %q; the cluster is still pointing at what we placed",
			reps[0].Endpoint)
	}
	if reps[0].State != Ready {
		t.Errorf("state = %s, want ready", reps[0].State)
	}
}

func TestAnAdoptedShardDoesNotDisplaceOneThatIsActuallyRunning(t *testing.T) {
	// The other direction. If our copy has a shard that has come up, two things
	// really are running and the one we did not ask for is the stray — the
	// agent is told to stop it rather than the live one being replaced.
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	ctx := context.Background()
	h.r.Reconcile(ctx)

	placed := h.r.Replicas()[0]
	h.r.OnShardStatus(placed.ID, placed.Shards[0].ID, supervisor.Ready, "")

	unclaimed := h.r.Adopt("1", []AdoptedShard{{
		NodeID: "1", ReplicaID: placed.ID, ShardID: "some-other-shard",
		Rank: 0, HTTPPort: 29999, State: supervisor.Ready,
		Model: "llama", Profile: "tp4", Endpoint: "http://node1:29999",
	}})

	if len(unclaimed) != 1 {
		t.Fatalf("the stray was not disowned: %+v", unclaimed)
	}
	if got := h.r.Replicas()[0].Endpoint; got == "http://node1:29999" {
		t.Error("a stray replaced the replica that was actually serving")
	}
}

func TestReplicaIDsDoNotRepeatAcrossRestarts(t *testing.T) {
	// The collision that caused the bind-conflict loop on a real cluster. A
	// control plane that restarts begins counting again, so the replica it
	// places next was called llama-1 — which is also what an engine started
	// before the restart, still running on another node, called itself. Two
	// nodes claimed the same shard id and every placement landed on a port an
	// orphan was already holding.
	seen := map[string]bool{}
	for i := range 1000 {
		// The same counter value every time, which is exactly what a restart
		// reproduces.
		id := replicaID("llama", 1)
		if seen[id] {
			t.Fatalf("id %q repeated after %d placements", id, i)
		}
		seen[id] = true
	}
}

func TestAPortHeldByAStrayIsNotHandedOut(t *testing.T) {
	// Stopping is not instant, and an engine keeps its socket until it exits.
	// Giving its port to a new shard makes that shard fail to bind, which fails
	// its replica, which places another one on the same port.
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	ctx := context.Background()
	h.r.Reconcile(ctx)

	placed := h.r.Replicas()[0]
	h.r.OnShardStatus(placed.ID, placed.Shards[0].ID, supervisor.Ready, "")

	// An agent reports a stray on the same replica id, holding a different
	// port. It is disowned — and its port must not become available.
	strayPort := placed.Shards[0].HTTPPort + 2
	unclaimed := h.r.Adopt("1", []AdoptedShard{{
		NodeID: "1", ReplicaID: placed.ID, ShardID: "stray-r0",
		Rank: 0, HTTPPort: strayPort, State: supervisor.Ready,
		Model: "llama", Profile: "tp4", Endpoint: "http://node1:29999",
	}})
	if len(unclaimed) != 1 {
		t.Fatalf("the stray was not disowned: %+v", unclaimed)
	}

	for range 5 {
		if got := h.r.allocatePorts(nil); got == strayPort {
			t.Fatalf("port %d was handed out while a stray still holds it", got)
		}
	}
}

// TestAnEngineThatDiesWhileItsAgentSurvivesIsNoticed is the mirror image of
// adoption, and the half that was missing.
//
// Observed on a real cluster: the engine was killed, the agent restarted and
// registered with nothing running, and the replica stayed Ready pointing at a
// dead port for as long as anyone watched. Traffic survived — the breaker took
// the endpoint out — but the cluster ran a third short with no sign of it.
func TestAnEngineThatDiesWhileItsAgentSurvivesIsNoticed(t *testing.T) {
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	ctx := context.Background()
	h.r.Reconcile(ctx)
	h.markReady()

	rep := h.r.Replicas()[0]
	if rep.State != Ready {
		t.Fatalf("state = %s", rep.State)
	}
	shardID := rep.Shards[0].ID

	// The agent reconnects and says it is running nothing at all.
	h.r.Adopt("1", nil)

	reps := h.r.Replicas()
	if len(reps) != 1 {
		t.Fatalf("replicas: %+v", reps)
	}
	if reps[0].State == Ready {
		t.Fatal("a replica whose shard the agent does not have is still Ready")
	}
	if !strings.Contains(reps[0].Err, "not running") {
		t.Errorf("the error does not say what happened: %q", reps[0].Err)
	}

	// The GPUs come back, or the replacement has nowhere to go.
	if free := len(h.inv.FreeGPUs("1")); free != 8 {
		t.Errorf("%d of 8 GPUs are free; the dead shard still holds some", free)
	}
	_ = shardID

	// And a replacement is placed once the backoff elapses. Several ticks,
	// because failing, reaping and placing are separate steps by design.
	for range 6 {
		h.advance(10 * time.Minute)
		h.r.Reconcile(ctx)
		h.markReady()
	}
	h.r.Reconcile(ctx)

	var ready int
	for _, r := range h.r.Replicas() {
		if r.State == Ready {
			ready++
		}
	}
	if ready != 1 {
		t.Errorf("%d ready replicas after the rebuild; want 1", ready)
	}
}

func TestAdoptionAndForgettingDoNotFightEachOther(t *testing.T) {
	// The shard the agent does report must survive the same registration that
	// forgets the ones it does not.
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	ctx := context.Background()
	h.r.Reconcile(ctx)
	h.markReady()

	rep := h.r.Replicas()[0]
	unclaimed := h.r.Adopt("1", []AdoptedShard{{
		NodeID: "1", ReplicaID: rep.ID, ShardID: rep.Shards[0].ID,
		Rank: 0, GPUIndices: rep.Shards[0].GPUIndices, HTTPPort: rep.Shards[0].HTTPPort,
		State: supervisor.Ready, Model: "llama", Profile: "tp4",
		Endpoint: rep.Endpoint,
	}})
	if len(unclaimed) != 0 {
		t.Fatalf("a shard the agent reported was disowned: %+v", unclaimed)
	}
	if got := h.r.Replicas()[0].State; got != Ready {
		t.Errorf("state = %s; the reported shard was forgotten anyway", got)
	}
}

func TestAnotherNodesShardsAreNotForgotten(t *testing.T) {
	// One agent's registration says nothing about any other node.
	h := newHarness(t, testModel(2), gpuNode("1", 8), gpuNode("2", 8))
	ctx := context.Background()
	for range 4 {
		h.r.Reconcile(ctx)
		h.markReady()
		h.advance(time.Second)
	}
	if n := len(h.r.Replicas()); n != 2 {
		t.Fatalf("%d replicas", n)
	}

	h.r.Adopt("1", nil)

	var onTwo int
	for _, rep := range h.r.Replicas() {
		for _, s := range rep.Shards {
			if s.NodeID == "2" && !s.State.Terminal() {
				onTwo++
			}
		}
	}
	if onTwo == 0 {
		t.Error("node 1's registration forgot node 2's shards")
	}
}

// TestAReplicaMissingAShardIsNotReady is what a control plane restart exposed:
// half of a two-node replica was adopted, every shard it still had was healthy,
// and it was reported Ready with no endpoint — so the router was handed a
// replica nothing could route to.
func TestAReplicaMissingAShardIsNotReady(t *testing.T) {
	rep := &Replica{
		ID: "r1", WantShards: 2,
		Shards: []Shard{{ID: "r1-r1", Rank: 1, State: supervisor.Ready}},
	}
	rep.deriveState(time.Now())

	if rep.State == Ready {
		t.Fatal("a replica with one of two shards is Ready")
	}
	if !strings.Contains(rep.Err, "1 of 2") {
		t.Errorf("the error does not say what is missing: %q", rep.Err)
	}
}

func TestACompleteReplicaIsReady(t *testing.T) {
	rep := &Replica{
		ID: "r1", WantShards: 2,
		Shards: []Shard{
			{ID: "r1-r0", Rank: 0, State: supervisor.Ready},
			{ID: "r1-r1", Rank: 1, State: supervisor.Ready},
		},
	}
	rep.deriveState(time.Now())
	if rep.State != Ready {
		t.Errorf("state = %s, want ready", rep.State)
	}
}

func TestAReplicaStillComingUpIsNotFailed(t *testing.T) {
	// Every shard present, not all ready yet: that is Starting, not a shortage.
	rep := &Replica{
		ID: "r1", WantShards: 2,
		Shards: []Shard{
			{ID: "r1-r0", Rank: 0, State: supervisor.Ready},
			{ID: "r1-r1", Rank: 1, State: supervisor.Loading},
		},
	}
	rep.deriveState(time.Now())
	if rep.State != Starting {
		t.Errorf("state = %s, want starting", rep.State)
	}
}

func TestAMultiNodeReplicaKnowsHowManyShardsItWants(t *testing.T) {
	// End to end through placement, because the count has to reach the replica
	// from the profile rather than from whatever turned up.
	m := testModel(1, catalog.Profile{
		Name: "tp8", GPUClass: "h100-80g", TP: 8, Nodes: 2, GPUsPerNode: 4,
	})
	h := newHarness(t, m, gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
	h.r.Reconcile(context.Background())

	reps := h.r.Replicas()
	if len(reps) != 1 {
		t.Fatalf("replicas: %+v", reps)
	}
	if reps[0].WantShards != 2 {
		t.Errorf("WantShards = %d, want 2", reps[0].WantShards)
	}
}