Skip to content

Commit 1e2f4d4

Browse files
rarguelloFrarguelloF
authored
feat(llmobs): experiments multi-run (#4660)
<!-- * New contributors are highly encouraged to read our [CONTRIBUTING](/CONTRIBUTING.md) documentation. * Commit and PR titles should be prefixed with the general area of the pull request's change. --> ### What does this PR do? <!-- * A brief description of the change being made with this pull request. * If the description here cannot be expressed in a succinct form, consider opening multiple pull requests instead of a single one. --> - Adds `llmobs/experiment.WithRuns(n)` option to enable multiple runs (https://docs.datadoghq.com/llm_observability/experiments/advanced_runs/#multiple-runs). - Implements experiment lifecycle status, so now the backend is notified of the experiment's state throughout execution: running when it starts, completed or failed when it finishes (based on whether any errors occurred), and interrupted if the context is cancelled before completion. - Propagate `experiment_id`, `run_id`, and `run_iteration` across distributed boundaries via APM baggage keys. Any LLMObs span started in a downstream service automatically inherits these fields as tags and gets scope=experiments. ### Motivation <!-- * What inspired you to submit this pull request? * Link any related GitHub issues or PRs here. * If this resolves a GitHub issue, include "Fixes #XXXX" to link the issue and auto-close it on merge. --> ### Reviewer's Checklist <!-- * Authors can use this list as a reference to ensure that there are no problems during the review but the signing off is to be done by the reviewer(s). --> - [x] Changed code has unit tests for its functionality at or near 100% coverage. - [ ] [System-Tests](https://github.com/DataDog/system-tests/) covering this feature have been added and enabled with the va.b.c-dev version tag. - [ ] There is a benchmark for any new code, or changes to existing code. - [ ] If this interacts with the agent in a new way, a system test has been added. - [ ] New code is free of linting errors. You can check this by running `make lint` locally. - [ ] New code doesn't break existing tests. You can check this by running `make test` locally. - [ ] Add an appropriate team label so this PR gets put in the right place for the release notes. - [ ] All generated files are up to date. You can check this by running `make generate` locally. - [ ] Non-trivial go.mod changes, e.g. adding new modules, are reviewed by @DataDog/dd-trace-go-guild. Make sure all nested modules are up to date by running `make fix-modules` locally. Unsure? Have a question? Request a review! Co-authored-by: rodrigo.arguello <[email protected]>
1 parent 7c0ac15 commit 1e2f4d4

6 files changed

Lines changed: 865 additions & 64 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 (
@@ -729,23 +731,52 @@ func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg
729731
}
730732
}
731733

732-
if experimentID := apmSpan.BaggageItem(baggageKeyExperimentID); experimentID != "" {
733-
span.scope = "experiments"
734+
experimentID := apmSpan.BaggageItem(baggageKeyExperimentID)
735+
experimentRunID := apmSpan.BaggageItem(baggageKeyExperimentRunID)
736+
experimentRunIteration := apmSpan.BaggageItem(baggageKeyExperimentRunIteration)
737+
if experimentID != "" || experimentRunID != "" || experimentRunIteration != "" {
738+
if span.llmCtx.tags == nil {
739+
span.llmCtx.tags = make(map[string]string)
740+
}
741+
if experimentID != "" {
742+
span.scope = "experiments"
743+
span.llmCtx.tags["experiment_id"] = experimentID
744+
}
745+
if experimentRunID != "" {
746+
span.llmCtx.tags["run_id"] = experimentRunID
747+
}
748+
if experimentRunIteration != "" {
749+
span.llmCtx.tags["run_iteration"] = experimentRunIteration
750+
}
734751
}
735752

736753
log.Debug("llmobs: starting LLMObs span: %s, span_kind: %s, ml_app: %s", spanName, kind, span.mlApp)
737754
return span, contextWithActiveLLMSpan(ctx, span)
738755
}
739756

740-
// StartExperimentSpan starts a new experiment span with the given name, experiment ID, and configuration.
757+
// ExperimentInfo holds the experiment identifiers propagated via baggage to distributed child spans.
758+
type ExperimentInfo struct {
759+
ID string
760+
RunID string
761+
RunIteration int
762+
}
763+
764+
// StartExperimentSpan starts a new experiment span with the given name and configuration.
765+
// ExperimentInfo fields are propagated via baggage so distributed child spans inherit them.
741766
// Returns the created span and a context containing the span.
742-
func (l *LLMObs) StartExperimentSpan(ctx context.Context, name string, experimentID string, cfg StartSpanConfig) (*Span, context.Context) {
767+
func (l *LLMObs) StartExperimentSpan(ctx context.Context, name string, params ExperimentInfo, cfg StartSpanConfig) (*Span, context.Context) {
743768
span, ctx := l.StartSpan(ctx, SpanKindExperiment, name, cfg)
744769

745-
if experimentID != "" {
746-
span.apm.SetBaggageItem(baggageKeyExperimentID, experimentID)
770+
if params.ID != "" {
771+
span.apm.SetBaggageItem(baggageKeyExperimentID, params.ID)
747772
span.scope = "experiments"
748773
}
774+
if params.RunID != "" {
775+
span.apm.SetBaggageItem(baggageKeyExperimentRunID, params.RunID)
776+
}
777+
if params.RunIteration > 0 {
778+
span.apm.SetBaggageItem(baggageKeyExperimentRunIteration, fmt.Sprintf("%d", params.RunIteration))
779+
}
749780
return span, ctx
750781
}
751782

‎internal/llmobs/llmobs_test.go‎

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

@@ -2083,7 +2135,7 @@ func TestDDAttributes(t *testing.T) {
20832135
ctx := context.Background()
20842136

20852137
experimentID := "test-experiment-123"
2086-
span, _ := ll.StartExperimentSpan(ctx, "test-experiment", experimentID, llmobs.StartSpanConfig{})
2138+
span, _ := ll.StartExperimentSpan(ctx, "test-experiment", llmobs.ExperimentInfo{ID: experimentID}, llmobs.StartSpanConfig{})
20872139
span.Finish(llmobs.FinishSpanConfig{})
20882140

20892141
apmSpans := tt.WaitForSpans(t, 1)
@@ -2110,7 +2162,7 @@ func TestDDAttributes(t *testing.T) {
21102162
ctx := context.Background()
21112163

21122164
experimentID := "test-experiment-456"
2113-
parentSpan, ctx := ll.StartExperimentSpan(ctx, "parent-experiment", experimentID, llmobs.StartSpanConfig{})
2165+
parentSpan, ctx := ll.StartExperimentSpan(ctx, "parent-experiment", llmobs.ExperimentInfo{ID: experimentID}, llmobs.StartSpanConfig{})
21142166
childSpan, _ := ll.StartSpan(ctx, llmobs.SpanKindLLM, "child-llm", llmobs.StartSpanConfig{})
21152167

21162168
childSpan.Finish(llmobs.FinishSpanConfig{})
@@ -2132,6 +2184,36 @@ func TestDDAttributes(t *testing.T) {
21322184

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

‎internal/llmobs/transport/dne.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,12 @@ type RequestAttributesExperimentCreate struct {
124124
Config map[string]any `json:"config,omitempty"`
125125
DatasetVersion int `json:"dataset_version,omitempty"`
126126
EnsureUnique bool `json:"ensure_unique,omitempty"`
127+
RunCount int `json:"run_count,omitempty"`
128+
}
129+
130+
type RequestAttributesExperimentUpdate struct {
131+
Status string `json:"status,omitempty"`
132+
Error string `json:"error,omitempty"`
127133
}
128134

129135
type RequestAttributesExperimentPushEvents struct {
@@ -156,6 +162,7 @@ type (
156162
CreateProjectRequest = Request[RequestAttributesProjectCreate]
157163

158164
CreateExperimentRequest = Request[RequestAttributesExperimentCreate]
165+
UpdateExperimentRequest = Request[RequestAttributesExperimentUpdate]
159166
PushExperimentEventsRequest = Request[RequestAttributesExperimentPushEvents]
160167
)
161168

@@ -441,6 +448,7 @@ func (c *Transport) CreateExperiment(
441448
expConfig map[string]any,
442449
tags []string,
443450
description string,
451+
runs int,
444452
) (*ExperimentView, error) {
445453
path := endpointPrefixDNE + "/experiments"
446454
method := http.MethodPost
@@ -461,6 +469,7 @@ func (c *Transport) CreateExperiment(
461469
Config: expConfig,
462470
DatasetVersion: datasetVersion,
463471
EnsureUnique: true,
472+
RunCount: runs,
464473
},
465474
},
466475
}
@@ -483,6 +492,29 @@ func (c *Transport) CreateExperiment(
483492
return &exp, nil
484493
}
485494

495+
func (c *Transport) UpdateExperimentStatus(ctx context.Context, id, status, errSummary string) error {
496+
path := fmt.Sprintf("%s/experiments/%s", endpointPrefixDNE, url.PathEscape(id))
497+
498+
body := UpdateExperimentRequest{
499+
Data: RequestData[RequestAttributesExperimentUpdate]{
500+
Type: resourceTypeExperiments,
501+
Attributes: RequestAttributesExperimentUpdate{
502+
Status: status,
503+
Error: errSummary,
504+
},
505+
},
506+
}
507+
508+
result, err := c.jsonRequest(ctx, http.MethodPatch, path, subdomainDNE, body, defaultTimeout)
509+
if err != nil {
510+
return err
511+
}
512+
if result.statusCode != http.StatusOK {
513+
return fmt.Errorf("unexpected status %d: %s", result.statusCode, string(result.body))
514+
}
515+
return nil
516+
}
517+
486518
func (c *Transport) PushExperimentEvents(
487519
ctx context.Context,
488520
experimentID string,

0 commit comments

Comments
 (0)