llmgw/internal/control/reconciler/trace_test.go

package reconciler

import (
	"errors"
	"strings"
	"testing"

	"go.opentelemetry.io/otel/codes"

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

func TestAPlacementIsTracedFromThePassToTheCommandAndTheReplicaServing(t *testing.T) {
	rec := obstest.Record(t)
	h := newHarness(t, testModel(1), gpuNode("1", 8))

	h.r.Reconcile(t.Context())

	pass := rec.One(t, [16]byte{}, "reconcile")
	tid := pass.SpanContext().TraceID()
	up := rec.One(t, tid, "scale up")
	place := rec.One(t, tid, "place")
	if !obstest.ChildOf(up, pass) || !obstest.ChildOf(place, up) {
		t.Fatal("the placement is not beneath the scale-up, beneath the pass")
	}
	if a := obstest.Attrs(place); a["llmgw.model"] != "llama" || a["llmgw.profile"] != "tp4" ||
		a["llmgw.gpu_class"] != "h100-80g" || a["llmgw.nodes"] != `["1"]` {
		t.Errorf("the placement does not say where it went: %v", a)
	}
	if a := obstest.Attrs(up); a["llmgw.outcome"] != "placed" || a["llmgw.desired"] != "1" {
		t.Errorf("the scale-up does not say what it did: %v", a)
	}

	// The command went under a span of its own, beneath the replica's start,
	// which is still open: the replica is loading.
	send := rec.One(t, tid, "send start")
	if got := h.disp.traces; len(got) != 1 || got[0].SpanID() != send.SpanContext().SpanID() {
		t.Fatalf("the command was not sent under its span: %v", got)
	}
	if len(rec.Named("replica start")) != 0 {
		t.Fatal("the replica's start ended before the replica served")
	}

	h.markReady()
	start := rec.One(t, tid, "replica start")
	if !obstest.ChildOf(start, up) || !obstest.ChildOf(send, start) {
		t.Error("the replica's start is not beneath the scale-up, or the command beneath it")
	}
	a := obstest.Attrs(start)
	if a["llmgw.state"] != "ready" || a["llmgw.profile"] != "tp4" || a["llmgw.ready_after_seconds"] == "" {
		t.Errorf("the replica's start does not say how it ended: %v", a)
	}
	if start.Status().Code != codes.Ok {
		t.Errorf("a replica that serves is %v", start.Status())
	}
	events := start.Events()
	if len(events) == 0 || events[len(events)-1].Name != "shard ready" {
		t.Errorf("the shard's progress is not on the replica's start: %v", events)
	}
	if sc := h.r.Replicas()[0].StartTrace(); sc.SpanID() != start.SpanContext().SpanID() {
		t.Error("the replica does not name its start's span for a job to link to")
	}
}

func TestAPassThatDoesNothingLeavesNoTrace(t *testing.T) {
	rec := obstest.Record(t)
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	h.r.Reconcile(t.Context())
	h.markReady()
	before := len(rec.Ended())

	// Every pass from here finds the model as it should be.
	for range 3 {
		h.r.Reconcile(t.Context())
	}
	if got := rec.Ended()[before:]; len(got) != 0 {
		var names []string
		for _, s := range got {
			names = append(names, s.Name())
		}
		t.Errorf("passes with nothing to do were traced: %v", names)
	}
}

func TestAPlacementThatFitsNowhereSaysWhy(t *testing.T) {
	rec := obstest.Record(t)
	m := testModel(1, catalog.Profile{Name: "tp16", GPUClass: "h100-80g", TP: 16})
	h := newHarness(t, m, gpuNode("1", 8))

	h.r.Reconcile(t.Context())

	place := rec.One(t, [16]byte{}, "place")
	if place.Status().Code != codes.Error {
		t.Errorf("a placement that found nowhere is %v", place.Status())
	}
	if a := obstest.Attrs(place); a["llmgw.rejections"] == "0" || !strings.Contains(a["llmgw.explain"], "tp16") {
		t.Errorf("the placement does not say what it turned down: %v", a)
	}
	if a := obstest.Attrs(rec.One(t, place.SpanContext().TraceID(), "scale up")); a["llmgw.outcome"] != "no room" {
		t.Errorf("the scale-up says %v", a)
	}
}

func TestAFailedStartEndsItsTraceWithTheEnginesWords(t *testing.T) {
	rec := obstest.Record(t)
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	h.r.Reconcile(t.Context())
	cmd := h.disp.startedShards()[0]

	h.r.OnShardStatus(cmd.ReplicaID, cmd.ShardID, supervisor.Failed, "CUDA out of memory")
	start := rec.One(t, [16]byte{}, "replica start")
	if start.Status().Code != codes.Error || !strings.Contains(start.Status().Description, "CUDA out of memory") {
		t.Errorf("the failed start says %+v", start.Status())
	}

	// The next pass reaps it, in a trace of its own, saying when it is tried
	// again.
	h.r.Reconcile(t.Context())
	reap := rec.One(t, [16]byte{}, "reap replica")
	a := obstest.Attrs(reap)
	if a["llmgw.replica"] != cmd.ReplicaID || a["llmgw.failures_in_a_row"] != "1" || a["llmgw.retry_in_seconds"] == "" {
		t.Errorf("the reaping does not say what happens next: %v", a)
	}
}

func TestADispatchThatFailsIsTracedAsFailed(t *testing.T) {
	rec := obstest.Record(t)
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	h.disp.failNext = errors.New("no agent connected")

	h.r.Reconcile(t.Context())

	send := rec.One(t, [16]byte{}, "send start")
	if send.Status().Code != codes.Error {
		t.Errorf("a command that was not sent is %v", send.Status())
	}
	if start := rec.One(t, [16]byte{}, "replica start"); start.Status().Code != codes.Error {
		t.Errorf("a replica none of whose shards was sent is %v", start.Status())
	}
}

func TestAPreemptionIsTracedWithItsVictims(t *testing.T) {
	h := fillThenWant(t, gpuModel("vital", catalog.ClassCritical, 1), gpuModel("chat", "", 1), gpuNode("1", 8))
	rec := obstest.Record(t)

	h.r.Reconcile(t.Context())

	pre := rec.One(t, [16]byte{}, "preempt")
	a := obstest.Attrs(pre)
	if a["llmgw.for"] != "vital" || !strings.Contains(a["llmgw.victims"], "chat") {
		t.Errorf("the preemption does not say who gave way to whom: %v", a)
	}
	drain := rec.One(t, pre.SpanContext().TraceID(), "drain replica")
	if !obstest.ChildOf(drain, pre) || obstest.Attrs(drain)["llmgw.reason"] != "preempted for vital" {
		t.Errorf("the victim's drain is not beneath the preemption: %v", obstest.Attrs(drain))
	}
}

func TestAStopSaysWhy(t *testing.T) {
	h := newHarness(t, testModel(1), gpuNode("1", 8))
	h.r.Reconcile(t.Context())
	h.markReady()
	rec := obstest.Record(t)

	h.setModels(t) // the model leaves the catalog
	h.r.Reconcile(t.Context())

	stop := rec.One(t, [16]byte{}, "stop replica")
	if a := obstest.Attrs(stop); a["llmgw.model"] != "llama" || a["llmgw.reason"] != "its model left the catalog" {
		t.Errorf("the stop does not say why: %v", a)
	}
}