llmgw/internal/obs/trace.go

package obs

import (
	"context"
	"fmt"
	"net/url"
	"os"
	"time"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
)

// Tracing is where traces go, over OTLP/HTTP.
//
// In server.yaml: one span per inference request, with admission, scheduling
// and each attempt at an engine beneath it, and the trace carried on to the
// engine in a traceparent header — the request's X-Request-Id is on the span,
// so a trace and `llmgw ctl request` name the same request. A control plane
// also traces every reconcile pass that does something, each replica from
// its placement until it serves, and each job `ctl model serve` and `ctl
// model update` start.
//
// In agent.yaml: the node's part of starting a replica — fetching the
// weights, launching the engine, waiting for it to answer — inside the
// control plane's trace of that start. Point both at the same collector.
type Tracing struct {
	// Endpoint is the collector, for example http://otel.internal:4318.
	// Empty falls back to OTEL_EXPORTER_OTLP_ENDPOINT, and to no tracing.
	Endpoint string `koanf:"endpoint"`
	// SampleRatio is the share of traces kept, 0 to 1, of those nobody
	// upstream already decided. Zero takes 1. A request with a sampled
	// traceparent of its own is kept, and an agent keeps every start the
	// control plane kept, whatever their own ratio says.
	SampleRatio float64 `koanf:"sample_ratio"`
}

// TracingOptions says where traces go.
type TracingOptions struct {
	Endpoint    string  // http(s)://host:4318; empty reads OTEL_EXPORTER_OTLP_ENDPOINT
	SampleRatio float64 // of traces the caller did not already decide; zero takes 1
	Service     string
	Version     string
	// Instance tells apart the processes of one service: the node, for an
	// agent. Empty leaves it out.
	Instance string
}

// SetupTracing installs the global tracer provider and the W3C propagator,
// and returns what flushes it on the way out. With no endpoint it installs
// nothing, and every span is the no-op one: tracing that is off costs a
// function call per span.
func SetupTracing(ctx context.Context, o TracingOptions) (func(context.Context) error, error) {
	noop := func(context.Context) error { return nil }
	endpoint := o.Endpoint
	if endpoint == "" {
		endpoint = os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
	}
	if endpoint == "" {
		return noop, nil
	}
	u, err := url.Parse(endpoint)
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
		return noop, fmt.Errorf("tracing endpoint %q is not an http(s) address", endpoint)
	}
	opts := []otlptracehttp.Option{otlptracehttp.WithEndpoint(u.Host)}
	if u.Scheme == "http" {
		opts = append(opts, otlptracehttp.WithInsecure())
	}
	if u.Path != "" && u.Path != "/" {
		opts = append(opts, otlptracehttp.WithURLPath(u.Path))
	}
	exp, err := otlptracehttp.New(ctx, opts...)
	if err != nil {
		return noop, err
	}
	ratio := o.SampleRatio
	if ratio <= 0 || ratio > 1 {
		ratio = 1
	}
	attrs := []attribute.KeyValue{
		attribute.String("service.name", o.Service),
		attribute.String("service.version", o.Version),
	}
	if o.Instance != "" {
		attrs = append(attrs, attribute.String("service.instance.id", o.Instance))
	}
	tp := sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exp, sdktrace.WithBatchTimeout(5*time.Second)),
		sdktrace.WithResource(resource.NewSchemaless(attrs...)),
		// A caller that sent a sampled traceparent has already decided; the
		// ratio is for requests that arrive without one.
		sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(ratio))),
	)
	otel.SetTracerProvider(tp)
	otel.SetTextMapPropagator(propagation.TraceContext{})
	return tp.Shutdown, nil
}

// Tracer is the named tracer of whatever provider is installed when it is
// called. Looked up per use rather than kept, so a provider installed after a
// component was built — as tests do — is the one its spans go to. With none
// installed it is the no-op tracer, and a span costs a function call.
func Tracer(name string) trace.Tracer { return otel.Tracer(name) }

// End closes a span, marking it failed when err is set.
func End(span trace.Span, err error) {
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
	}
	span.End()
}

// Inject is the trace context of ctx as a message carries it: traceparent
// and, when there is one, tracestate. Nil when ctx holds no span, which is
// every context while tracing is off, so a message then carries nothing.
func Inject(ctx context.Context) map[string]string {
	if !trace.SpanContextFromContext(ctx).IsValid() {
		return nil
	}
	m := propagation.MapCarrier{}
	otel.GetTextMapPropagator().Inject(ctx, m)
	if len(m) == 0 {
		return nil
	}
	return m
}

// Extract is ctx continuing the trace a message carried, as Inject wrote it.
// A message that carried none leaves ctx as it was.
func Extract(ctx context.Context, carried map[string]string) context.Context {
	if len(carried) == 0 {
		return ctx
	}
	return otel.GetTextMapPropagator().Extract(ctx, propagation.MapCarrier(carried))
}