llmgw/internal/agent/trace_test.go

package agent

import (
	"testing"
	"time"

	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

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

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

	// The control plane's side, as the reconciler sends it: the command
	// dispatched under a span of its own, over the real stream.
	ctx, sent := obs.Tracer("test").Start(t.Context(), "send start",
		trace.WithSpanKind(trace.SpanKindProducer))
	if err := h.svc.StartShard(ctx, shardCommand(t, "shard-a")); err != nil {
		t.Fatal(err)
	}
	sent.End()
	h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Ready
	})

	// The node's part ends just after it reports the shard ready.
	tid := sent.SpanContext().TraceID()
	node := rec.Await(t, tid, "start shard", 5*time.Second)
	if node.Parent().SpanID() != sent.SpanContext().SpanID() {
		t.Fatal("the node's start is not beneath the control plane's command")
	}
	if node.SpanKind() != trace.SpanKindConsumer {
		t.Errorf("the node's start is a %v span", node.SpanKind())
	}
	if a := obstest.Attrs(node); a["llmgw.node"] != "n1" || a["llmgw.replica"] != "rep-1" ||
		a["llmgw.model"] != "test-model" || a["llmgw.engine"] != "vllm" {
		t.Errorf("the node's start does not say what it started: %v", a)
	}
	if node.Status().Code == codes.Error {
		t.Errorf("a start that served is %v", node.Status())
	}
	for _, name := range []string{"fetch weights", "launch engine", "wait for engine"} {
		if !obstest.ChildOf(rec.One(t, tid, name), node) {
			t.Errorf("%s is not beneath the node's start", name)
		}
	}
	if a := obstest.Attrs(rec.One(t, tid, "fetch weights")); a["llmgw.weights.outcome"] != "local" {
		t.Errorf("the weights do not say how they were found: %v", a)
	}
	if a := obstest.Attrs(rec.One(t, tid, "wait for engine")); a["llmgw.polls"] == "" || a["llmgw.polls"] == "0" {
		t.Errorf("the readiness wait does not say how often it asked: %v", a)
	}

	// A stop joins the trace it was sent in, too.
	sctx, stop := obs.Tracer("test").Start(t.Context(), "stop replica")
	if err := h.svc.StopShard(sctx, "n1", "shard-a"); err != nil {
		t.Fatal(err)
	}
	stop.End()
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-a" && s.state == supervisor.Stopped
	})
	stid := stop.SpanContext().TraceID()
	nodeStop := rec.Await(t, stid, "stop shard", 5*time.Second)
	if !obstest.ChildOf(nodeStop, rec.One(t, stid, "stop replica")) {
		t.Error("the node's stop is not beneath the control plane's")
	}
	if !obstest.ChildOf(rec.One(t, stid, "stop engine"), nodeStop) {
		t.Error("stopping the engine is not beneath the node's stop")
	}
}

func TestAnAdoptionIsTracedAndItsStopsAreInIt(t *testing.T) {
	rec := obstest.Record(t)
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
	})
	if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-orphaned")); err != nil {
		t.Fatal(err)
	}
	h.rec.await(t, 15*time.Second, func(s status) bool {
		return s.shardID == "shard-orphaned" && s.state == supervisor.Starting
	})

	// The control plane, restarted, claims nothing the node is running.
	h.rec.markUnclaimed("shard-orphaned")
	h.ln.drop()
	h.rec.await(t, 20*time.Second, func(s status) bool {
		return s.shardID == "shard-orphaned" && s.state == supervisor.Stopped
	})

	adopt := rec.Named("adopt")
	if len(adopt) != 1 {
		t.Fatalf("%d adoptions traced", len(adopt))
	}
	if a := obstest.Attrs(adopt[0]); a["llmgw.node"] != "n1" || a["llmgw.shards"] != "1" || a["llmgw.unclaimed"] != "1" {
		t.Errorf("the adoption does not say what it found: %v", a)
	}
	stop := rec.Await(t, adopt[0].SpanContext().TraceID(), "stop shard", 5*time.Second)
	if !obstest.ChildOf(stop, adopt[0]) {
		t.Error("stopping the unclaimed shard is not beneath the adoption")
	}
}

func TestAStartThatFailsOnTheNodeSaysWhyInTheTrace(t *testing.T) {
	rec := obstest.Record(t)
	h := start(t, func(c *Config) {
		c.Engines.VLLM.Bin = fakeEngineScript(t, `echo "CUDA out of memory" >&2; exit 1`)
	})

	ctx, sent := obs.Tracer("test").Start(t.Context(), "send start")
	if err := h.svc.StartShard(ctx, shardCommand(t, "shard-b")); err != nil {
		t.Fatal(err)
	}
	sent.End()
	h.rec.await(t, 30*time.Second, func(s status) bool {
		return s.shardID == "shard-b" && s.state == supervisor.Failed
	})

	tid := sent.SpanContext().TraceID()
	node := rec.Await(t, tid, "start shard", 5*time.Second)
	if node.Status().Code != codes.Error {
		t.Errorf("a start that failed is %v", node.Status())
	}
	if wait := rec.One(t, tid, "wait for engine"); wait.Status().Code != codes.Error {
		t.Errorf("the readiness wait of an engine that exited is %v", wait.Status())
	}
}