Skip to content

Commit 5dfaa87

Browse files
rarguelloFrarguelloF
committed
propagate experiment info
1 parent 45e659c commit 5dfaa87

3 files changed

Lines changed: 127 additions & 10 deletions

File tree

‎internal/llmobs/llmobs.go‎

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@ var (
4747
)
4848

4949
const (
50-
baggageKeyExperimentID = "_ml_obs.experiment_id"
50+
baggageKeyExperimentID = "_ml_obs.experiment_id"
51+
baggageKeyExperimentRunID = "_ml_obs.experiment_run_id"
52+
baggageKeyExperimentRunIteration = "_ml_obs.experiment_run_iteration"
5153
)
5254

5355
const (
@@ -717,23 +719,52 @@ func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg
717719
}
718720
}
719721

720-
if experimentID := apmSpan.BaggageItem(baggageKeyExperimentID); experimentID != "" {
721-
span.scope = "experiments"
722+
experimentID := apmSpan.BaggageItem(baggageKeyExperimentID)
723+
experimentRunID := apmSpan.BaggageItem(baggageKeyExperimentRunID)
724+
experimentRunIteration := apmSpan.BaggageItem(baggageKeyExperimentRunIteration)
725+
if experimentID != "" || experimentRunID != "" || experimentRunIteration != "" {
726+
if span.llmCtx.tags == nil {
727+
span.llmCtx.tags = make(map[string]string)
728+
}
729+
if experimentID != "" {
730+
span.scope = "experiments"
731+
span.llmCtx.tags["experiment_id"] = experimentID
732+
}
733+
if experimentRunID != "" {
734+
span.llmCtx.tags["run_id"] = experimentRunID
735+
}
736+
if experimentRunIteration != "" {
737+
span.llmCtx.tags["run_iteration"] = experimentRunIteration
738+
}
722739
}
723740

724741
log.Debug("llmobs: starting LLMObs span: %s, span_kind: %s, ml_app: %s", spanName, kind, span.mlApp)
725742
return span, contextWithActiveLLMSpan(ctx, span)
726743
}
727744

728-
// StartExperimentSpan starts a new experiment span with the given name, experiment ID, and configuration.
745+
// ExperimentInfo holds the experiment identifiers propagated via baggage to distributed child spans.
746+
type ExperimentInfo struct {
747+
ID string
748+
RunID string
749+
RunIteration int
750+
}
751+
752+
// StartExperimentSpan starts a new experiment span with the given name and configuration.
753+
// ExperimentInfo fields are propagated via baggage so distributed child spans inherit them.
729754
// Returns the created span and a context containing the span.
730-
func (l *LLMObs) StartExperimentSpan(ctx context.Context, name string, experimentID string, cfg StartSpanConfig) (*Span, context.Context) {
755+
func (l *LLMObs) StartExperimentSpan(ctx context.Context, name string, params ExperimentInfo, cfg StartSpanConfig) (*Span, context.Context) {
731756
span, ctx := l.StartSpan(ctx, SpanKindExperiment, name, cfg)
732757

733-
if experimentID != "" {
734-
span.apm.SetBaggageItem(baggageKeyExperimentID, experimentID)
758+
if params.ID != "" {
759+
span.apm.SetBaggageItem(baggageKeyExperimentID, params.ID)
735760
span.scope = "experiments"
736761
}
762+
if params.RunID != "" {
763+
span.apm.SetBaggageItem(baggageKeyExperimentRunID, params.RunID)
764+
}
765+
if params.RunIteration > 0 {
766+
span.apm.SetBaggageItem(baggageKeyExperimentRunIteration, fmt.Sprintf("%d", params.RunIteration))
767+
}
737768
return span, ctx
738769
}
739770

‎internal/llmobs/llmobs_test.go‎

Lines changed: 84 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,58 @@ func TestStartSpan(t *testing.T) {
258258
assert.Equal(t, llmWorkflow1.ParentID, llmLLM1.SpanID, "workflow-1 parent ID should be the llm-1 span ID")
259259
assert.Equal(t, llmAgent1.ParentID, llmWorkflow1.SpanID, "agent-1 parent ID should be the workflow-1 span ID")
260260
})
261+
t.Run("distributed-context-propagation-experiment-baggage", func(t *testing.T) {
262+
tt, ll := testTracer(t)
263+
264+
experimentID := "exp-dist-123"
265+
experimentRunID := "run-uuid-xyz"
266+
experimentRunIteration := 3
267+
268+
h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
269+
ctx := req.Context()
270+
serverSpan, _ := ll.StartSpan(ctx, llmobs.SpanKindLLM, "server-llm", llmobs.StartSpanConfig{})
271+
defer serverSpan.Finish(llmobs.FinishSpanConfig{})
272+
w.Write([]byte("ok"))
273+
})
274+
srv, cl := testClientServer(t, h)
275+
276+
genSpans := func() {
277+
experimentSpan, ctx := ll.StartExperimentSpan(context.Background(), "client-experiment", llmobs.ExperimentInfo{
278+
ID: experimentID,
279+
RunID: experimentRunID,
280+
RunIteration: experimentRunIteration,
281+
}, llmobs.StartSpanConfig{})
282+
defer experimentSpan.Finish(llmobs.FinishSpanConfig{})
283+
284+
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+"/", nil)
285+
require.NoError(t, err)
286+
resp, err := cl.Do(req)
287+
require.NoError(t, err)
288+
require.Equal(t, http.StatusOK, resp.StatusCode)
289+
_ = resp.Body.Close()
290+
}
291+
genSpans()
292+
293+
_ = tt.WaitForSpans(t, 4) // experiment + server-llm (LLMObs) + HTTP client + HTTP server (APM)
294+
llmSpans := tt.WaitForLLMObsSpans(t, 2)
295+
296+
var experimentLLM, serverLLM *llmobstransport.LLMObsSpanEvent
297+
for i := range llmSpans {
298+
switch llmSpans[i].Name {
299+
case "client-experiment":
300+
experimentLLM = &llmSpans[i]
301+
case "server-llm":
302+
serverLLM = &llmSpans[i]
303+
}
304+
}
305+
require.NotNil(t, experimentLLM, "client experiment span should exist")
306+
require.NotNil(t, serverLLM, "server LLM span should exist")
307+
308+
assert.Equal(t, "experiments", serverLLM.DDAttributes.Scope, "server span should inherit experiments scope via baggage")
309+
assert.Equal(t, experimentID, findTag(serverLLM.Tags, "experiment_id"), "server span should inherit experiment_id via baggage")
310+
assert.Equal(t, experimentRunID, findTag(serverLLM.Tags, "run_id"), "server span should inherit run_id via baggage")
311+
assert.Equal(t, fmt.Sprintf("%d", experimentRunIteration), findTag(serverLLM.Tags, "run_iteration"), "server span should inherit run_iteration via baggage")
312+
})
261313
t.Run("custom-start-and-finish-times", func(t *testing.T) {
262314
tt, ll := testTracer(t)
263315

@@ -2080,7 +2132,7 @@ func TestDDAttributes(t *testing.T) {
20802132
ctx := context.Background()
20812133

20822134
experimentID := "test-experiment-123"
2083-
span, _ := ll.StartExperimentSpan(ctx, "test-experiment", experimentID, llmobs.StartSpanConfig{})
2135+
span, _ := ll.StartExperimentSpan(ctx, "test-experiment", llmobs.ExperimentInfo{ID: experimentID}, llmobs.StartSpanConfig{})
20842136
span.Finish(llmobs.FinishSpanConfig{})
20852137

20862138
apmSpans := tt.WaitForSpans(t, 1)
@@ -2107,7 +2159,7 @@ func TestDDAttributes(t *testing.T) {
21072159
ctx := context.Background()
21082160

21092161
experimentID := "test-experiment-456"
2110-
parentSpan, ctx := ll.StartExperimentSpan(ctx, "parent-experiment", experimentID, llmobs.StartSpanConfig{})
2162+
parentSpan, ctx := ll.StartExperimentSpan(ctx, "parent-experiment", llmobs.ExperimentInfo{ID: experimentID}, llmobs.StartSpanConfig{})
21112163
childSpan, _ := ll.StartSpan(ctx, llmobs.SpanKindLLM, "child-llm", llmobs.StartSpanConfig{})
21122164

21132165
childSpan.Finish(llmobs.FinishSpanConfig{})
@@ -2129,6 +2181,36 @@ func TestDDAttributes(t *testing.T) {
21292181

21302182
assert.Equal(t, "experiments", parentLLM.DDAttributes.Scope, "Parent scope should be 'experiments'")
21312183
assert.Equal(t, "experiments", childLLM.DDAttributes.Scope, "Child scope should be 'experiments' via baggage propagation")
2184+
assert.Contains(t, childLLM.Tags, "experiment_id:"+experimentID, "Child span should inherit experiment_id tag from baggage")
2185+
})
2186+
t.Run("child-span-inherits-run-id-and-run-iteration-from-baggage", func(t *testing.T) {
2187+
tt, ll := testTracer(t)
2188+
ctx := context.Background()
2189+
2190+
experimentID := "test-experiment-789"
2191+
experimentRunID := "run-uuid-abc"
2192+
experimentRunIteration := 2
2193+
parentSpan, ctx := ll.StartExperimentSpan(ctx, "parent-experiment", llmobs.ExperimentInfo{
2194+
ID: experimentID,
2195+
RunID: experimentRunID,
2196+
RunIteration: experimentRunIteration,
2197+
}, llmobs.StartSpanConfig{})
2198+
childSpan, _ := ll.StartSpan(ctx, llmobs.SpanKindLLM, "child-llm", llmobs.StartSpanConfig{})
2199+
2200+
childSpan.Finish(llmobs.FinishSpanConfig{})
2201+
parentSpan.Finish(llmobs.FinishSpanConfig{})
2202+
2203+
llmSpans := tt.WaitForLLMObsSpans(t, 2)
2204+
2205+
var childLLM *llmobstransport.LLMObsSpanEvent
2206+
for i := range llmSpans {
2207+
if llmSpans[i].Name == "child-llm" {
2208+
childLLM = &llmSpans[i]
2209+
}
2210+
}
2211+
require.NotNil(t, childLLM, "Child LLM span should exist")
2212+
assert.Contains(t, childLLM.Tags, "run_id:"+experimentRunID, "Child span should inherit run_id from baggage")
2213+
assert.Contains(t, childLLM.Tags, "run_iteration:2", "Child span should inherit run_iteration from baggage")
21322214
})
21332215
t.Run("child-span-trace-ids", func(t *testing.T) {
21342216
tt, ll := testTracer(t)

‎llmobs/experiment/experiment.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,11 @@ func (e *Experiment) runTaskForRecord(ctx context.Context, llmobs *illmobs.LLMOb
450450
err error
451451
)
452452

453-
span, ctx := llmobs.StartExperimentSpan(ctx, e.task.Name(), e.id, illmobs.StartSpanConfig{})
453+
span, ctx := llmobs.StartExperimentSpan(ctx, e.task.Name(), illmobs.ExperimentInfo{
454+
ID: e.id,
455+
RunID: run.ID,
456+
RunIteration: run.Iteration,
457+
}, illmobs.StartSpanConfig{})
454458
defer func() { span.Finish(illmobs.FinishSpanConfig{Error: err}) }()
455459

456460
tags := make(map[string]string)

0 commit comments

Comments
 (0)