|
| 1 | +/* |
| 2 | +Copyright 2025 The Kubernetes Authors. |
| 3 | +
|
| 4 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +you may not use this file except in compliance with the License. |
| 6 | +You may obtain a copy of the License at |
| 7 | +
|
| 8 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | +Unless required by applicable law or agreed to in writing, software |
| 11 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +See the License for the specific language governing permissions and |
| 14 | +limitations under the License. |
| 15 | +*/ |
| 16 | + |
| 17 | +package scheduling |
| 18 | + |
| 19 | +import ( |
| 20 | + "context" |
| 21 | + "testing" |
| 22 | + |
| 23 | + "go.opentelemetry.io/otel" |
| 24 | + "go.opentelemetry.io/otel/attribute" |
| 25 | + sdktrace "go.opentelemetry.io/otel/sdk/trace" |
| 26 | + "go.opentelemetry.io/otel/sdk/trace/tracetest" |
| 27 | + "go.opentelemetry.io/otel/trace" |
| 28 | + k8stypes "k8s.io/apimachinery/pkg/types" |
| 29 | + |
| 30 | + fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer" |
| 31 | + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" |
| 32 | + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" |
| 33 | +) |
| 34 | + |
| 35 | +// setupSpanRecorder installs an in-memory span recorder as the global tracer |
| 36 | +// provider and returns it, restoring the previous provider on cleanup. |
| 37 | +func setupSpanRecorder(t *testing.T) *tracetest.SpanRecorder { |
| 38 | + t.Helper() |
| 39 | + recorder := tracetest.NewSpanRecorder() |
| 40 | + tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) |
| 41 | + origTP := otel.GetTracerProvider() |
| 42 | + otel.SetTracerProvider(tp) |
| 43 | + t.Cleanup(func() { otel.SetTracerProvider(origTP) }) |
| 44 | + return recorder |
| 45 | +} |
| 46 | + |
| 47 | +func findSpans(spans []sdktrace.ReadOnlySpan, name string) []sdktrace.ReadOnlySpan { |
| 48 | + var out []sdktrace.ReadOnlySpan |
| 49 | + for _, s := range spans { |
| 50 | + if s.Name() == name { |
| 51 | + out = append(out, s) |
| 52 | + } |
| 53 | + } |
| 54 | + return out |
| 55 | +} |
| 56 | + |
| 57 | +func spanAttributes(span sdktrace.ReadOnlySpan) map[attribute.Key]attribute.Value { |
| 58 | + attrs := make(map[attribute.Key]attribute.Value) |
| 59 | + for _, kv := range span.Attributes() { |
| 60 | + attrs[kv.Key] = kv.Value |
| 61 | + } |
| 62 | + return attrs |
| 63 | +} |
| 64 | + |
| 65 | +func newTestEndpoints(names ...string) []fwksched.Endpoint { |
| 66 | + endpoints := make([]fwksched.Endpoint, len(names)) |
| 67 | + for i, name := range names { |
| 68 | + endpoints[i] = fwksched.NewEndpoint( |
| 69 | + &fwkdl.EndpointMetadata{NamespacedName: k8stypes.NamespacedName{Name: name}}, nil, nil) |
| 70 | + } |
| 71 | + return endpoints |
| 72 | +} |
| 73 | + |
| 74 | +func TestRunFilterPluginsSingleSpan(t *testing.T) { |
| 75 | + recorder := setupSpanRecorder(t) |
| 76 | + |
| 77 | + filter := &testPlugin{ |
| 78 | + typedName: fwkplugin.TypedName{Type: "test-filter", Name: "instance-a"}, |
| 79 | + FilterRes: []k8stypes.NamespacedName{{Name: "pod1"}, {Name: "pod2"}}, |
| 80 | + } |
| 81 | + profile := NewSchedulerProfile().WithFilters(filter) |
| 82 | + endpoints := newTestEndpoints("pod1", "pod2", "pod3") |
| 83 | + |
| 84 | + ctx, root := otel.Tracer("test").Start(context.Background(), "root") |
| 85 | + result := profile.runFilterPlugins(ctx, &fwksched.InferenceRequest{TargetModel: "m1", RequestID: "r1"}, endpoints) |
| 86 | + root.End() |
| 87 | + |
| 88 | + if len(result) != 2 { |
| 89 | + t.Fatalf("runFilterPlugins returned %d endpoints, want 2", len(result)) |
| 90 | + } |
| 91 | + |
| 92 | + spans := findSpans(recorder.Ended(), "filter_endpoints") |
| 93 | + if len(spans) != 1 { |
| 94 | + t.Fatalf("got %d filter_endpoints spans, want 1", len(spans)) |
| 95 | + } |
| 96 | + span := spans[0] |
| 97 | + if span.SpanKind() != trace.SpanKindInternal { |
| 98 | + t.Errorf("span kind = %v, want Internal", span.SpanKind()) |
| 99 | + } |
| 100 | + if span.Parent().SpanID() != root.SpanContext().SpanID() { |
| 101 | + t.Errorf("parent span ID = %v, want root %v", span.Parent().SpanID(), root.SpanContext().SpanID()) |
| 102 | + } |
| 103 | + |
| 104 | + attrs := spanAttributes(span) |
| 105 | + if got := attrs["llm_d.epp.filter.candidate_endpoints"].AsInt64(); got != 3 { |
| 106 | + t.Errorf("candidate_endpoints = %d, want 3", got) |
| 107 | + } |
| 108 | + if got := attrs["llm_d.epp.filter.filtered_endpoints"].AsInt64(); got != 2 { |
| 109 | + t.Errorf("filtered_endpoints = %d, want 2", got) |
| 110 | + } |
| 111 | + if got := attrs["gen_ai.request.model"].AsString(); got != "m1" { |
| 112 | + t.Errorf("gen_ai.request.model = %q, want %q", got, "m1") |
| 113 | + } |
| 114 | + if got := attrs["gen_ai.request.id"].AsString(); got != "r1" { |
| 115 | + t.Errorf("gen_ai.request.id = %q, want %q", got, "r1") |
| 116 | + } |
| 117 | +} |
| 118 | + |
| 119 | +// A filter chain emits one span over the whole stage: candidate is the input to |
| 120 | +// the first filter, filtered is the output of the last filter that ran. |
| 121 | +func TestRunFilterPluginsChainEmitsOneSpan(t *testing.T) { |
| 122 | + recorder := setupSpanRecorder(t) |
| 123 | + |
| 124 | + filterA := &testPlugin{ |
| 125 | + typedName: fwkplugin.TypedName{Type: "filter-a", Name: "a"}, |
| 126 | + FilterRes: []k8stypes.NamespacedName{{Name: "pod1"}, {Name: "pod2"}}, |
| 127 | + } |
| 128 | + filterB := &testPlugin{ |
| 129 | + typedName: fwkplugin.TypedName{Type: "filter-b", Name: "b"}, |
| 130 | + FilterRes: []k8stypes.NamespacedName{{Name: "pod1"}}, |
| 131 | + } |
| 132 | + profile := NewSchedulerProfile().WithFilters(filterA, filterB) |
| 133 | + endpoints := newTestEndpoints("pod1", "pod2", "pod3") |
| 134 | + |
| 135 | + ctx, root := otel.Tracer("test").Start(context.Background(), "root") |
| 136 | + result := profile.runFilterPlugins(ctx, &fwksched.InferenceRequest{}, endpoints) |
| 137 | + root.End() |
| 138 | + |
| 139 | + if len(result) != 1 { |
| 140 | + t.Fatalf("runFilterPlugins returned %d endpoints, want 1", len(result)) |
| 141 | + } |
| 142 | + spans := findSpans(recorder.Ended(), "filter_endpoints") |
| 143 | + if len(spans) != 1 { |
| 144 | + t.Fatalf("got %d filter_endpoints spans, want 1 for the whole chain", len(spans)) |
| 145 | + } |
| 146 | + attrs := spanAttributes(spans[0]) |
| 147 | + if got := attrs["llm_d.epp.filter.candidate_endpoints"].AsInt64(); got != 3 { |
| 148 | + t.Errorf("candidate_endpoints = %d, want 3", got) |
| 149 | + } |
| 150 | + if got := attrs["llm_d.epp.filter.filtered_endpoints"].AsInt64(); got != 1 { |
| 151 | + t.Errorf("filtered_endpoints = %d, want 1", got) |
| 152 | + } |
| 153 | +} |
| 154 | + |
| 155 | +func TestRunFilterPluginsDrainBreakStillEndsSpan(t *testing.T) { |
| 156 | + recorder := setupSpanRecorder(t) |
| 157 | + |
| 158 | + drain := &testPlugin{ |
| 159 | + typedName: fwkplugin.TypedName{Type: "drain-filter", Name: "drain"}, |
| 160 | + FilterRes: []k8stypes.NamespacedName{}, |
| 161 | + } |
| 162 | + never := &testPlugin{ |
| 163 | + typedName: fwkplugin.TypedName{Type: "never-filter", Name: "never"}, |
| 164 | + FilterRes: []k8stypes.NamespacedName{{Name: "pod1"}}, |
| 165 | + } |
| 166 | + profile := NewSchedulerProfile().WithFilters(drain, never) |
| 167 | + endpoints := newTestEndpoints("pod1", "pod2") |
| 168 | + |
| 169 | + ctx, root := otel.Tracer("test").Start(context.Background(), "root") |
| 170 | + result := profile.runFilterPlugins(ctx, &fwksched.InferenceRequest{}, endpoints) |
| 171 | + root.End() |
| 172 | + |
| 173 | + if len(result) != 0 { |
| 174 | + t.Fatalf("runFilterPlugins returned %d endpoints, want 0", len(result)) |
| 175 | + } |
| 176 | + spans := findSpans(recorder.Ended(), "filter_endpoints") |
| 177 | + if len(spans) != 1 { |
| 178 | + t.Fatalf("got %d filter_endpoints spans, want 1 (span must end on drain)", len(spans)) |
| 179 | + } |
| 180 | + if got := spanAttributes(spans[0])["llm_d.epp.filter.filtered_endpoints"].AsInt64(); got != 0 { |
| 181 | + t.Errorf("filtered_endpoints = %d, want 0", got) |
| 182 | + } |
| 183 | + if never.FilterCallCount != 0 { |
| 184 | + t.Errorf("second filter ran %d times after drain, want 0", never.FilterCallCount) |
| 185 | + } |
| 186 | +} |
| 187 | + |
| 188 | +// childSpanFilter starts a child span from the context it is given, so a test |
| 189 | +// can assert the delegate runs inside the filter span. |
| 190 | +type childSpanFilter struct{ typedName fwkplugin.TypedName } |
| 191 | + |
| 192 | +func (f *childSpanFilter) TypedName() fwkplugin.TypedName { return f.typedName } |
| 193 | + |
| 194 | +func (f *childSpanFilter) Filter(ctx context.Context, _ *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) []fwksched.Endpoint { |
| 195 | + _, span := otel.Tracer("test").Start(ctx, "inner_filter_span") |
| 196 | + span.End() |
| 197 | + return endpoints |
| 198 | +} |
| 199 | + |
| 200 | +func TestRunFilterPluginsNestsDelegateSpan(t *testing.T) { |
| 201 | + recorder := setupSpanRecorder(t) |
| 202 | + |
| 203 | + filter := &childSpanFilter{typedName: fwkplugin.TypedName{Type: "child", Name: "c"}} |
| 204 | + profile := NewSchedulerProfile().WithFilters(filter) |
| 205 | + endpoints := newTestEndpoints("pod1") |
| 206 | + |
| 207 | + ctx, root := otel.Tracer("test").Start(context.Background(), "root") |
| 208 | + profile.runFilterPlugins(ctx, &fwksched.InferenceRequest{}, endpoints) |
| 209 | + root.End() |
| 210 | + |
| 211 | + outer := findSpans(recorder.Ended(), "filter_endpoints") |
| 212 | + inner := findSpans(recorder.Ended(), "inner_filter_span") |
| 213 | + if len(outer) != 1 || len(inner) != 1 { |
| 214 | + t.Fatalf("got %d filter_endpoints and %d inner spans, want 1 each", len(outer), len(inner)) |
| 215 | + } |
| 216 | + if inner[0].Parent().SpanID() != outer[0].SpanContext().SpanID() { |
| 217 | + t.Errorf("inner span parent = %v, want filter_endpoints span %v", |
| 218 | + inner[0].Parent().SpanID(), outer[0].SpanContext().SpanID()) |
| 219 | + } |
| 220 | +} |
| 221 | + |
| 222 | +func TestRunFilterPluginsOmitsEmptyGenAI(t *testing.T) { |
| 223 | + recorder := setupSpanRecorder(t) |
| 224 | + |
| 225 | + filter := &testPlugin{ |
| 226 | + typedName: fwkplugin.TypedName{Type: "f", Name: "n"}, |
| 227 | + FilterRes: []k8stypes.NamespacedName{{Name: "pod1"}}, |
| 228 | + } |
| 229 | + profile := NewSchedulerProfile().WithFilters(filter) |
| 230 | + endpoints := newTestEndpoints("pod1") |
| 231 | + |
| 232 | + ctx, root := otel.Tracer("test").Start(context.Background(), "root") |
| 233 | + profile.runFilterPlugins(ctx, &fwksched.InferenceRequest{}, endpoints) |
| 234 | + root.End() |
| 235 | + |
| 236 | + spans := findSpans(recorder.Ended(), "filter_endpoints") |
| 237 | + if len(spans) != 1 { |
| 238 | + t.Fatalf("got %d filter_endpoints spans, want 1", len(spans)) |
| 239 | + } |
| 240 | + attrs := spanAttributes(spans[0]) |
| 241 | + if _, ok := attrs["gen_ai.request.model"]; ok { |
| 242 | + t.Error("gen_ai.request.model set for empty TargetModel") |
| 243 | + } |
| 244 | + if _, ok := attrs["gen_ai.request.id"]; ok { |
| 245 | + t.Error("gen_ai.request.id set for empty RequestID") |
| 246 | + } |
| 247 | +} |
0 commit comments