Skip to content

Commit fe11a42

Browse files
genesorgenesor
authored
feat: add deterministic sampling with knuth based on TraceId (#5712)
* add threshold computation * add traceId BigInt func * add span context to sampler * add jsdoc and backport old Sampler as RandomSampler * accept both Span and SpanContext in isSampled func * update calls to isSampled with span * move langchain plugin to deterministic sampler
1 parent 1639cf7 commit fe11a42

File tree

21 files changed

+315
-88
lines changed

21 files changed

+315
-88
lines changed

‎benchmark/core.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ suite
9393
sampler = new Sampler(0.5)
9494
},
9595
fn () {
96-
sampler.isSampled()
96+
sampler.isSampled(span)
9797
}
9898
})
9999
.add('format', {

‎packages/datadog-plugin-google-cloud-vertexai/src/tracing.js

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,18 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
2525
bindStart (ctx) {
2626
const { instance, request, resource, stream } = ctx
2727

28-
const tags = this.tagRequest(request, instance, stream)
29-
3028
const span = this.startSpan('vertexai.request', {
3129
service: this.config.service,
3230
resource,
3331
kind: 'client',
3432
meta: {
35-
[MEASURED]: 1,
36-
...tags
33+
[MEASURED]: 1
3734
}
3835
}, false)
3936

37+
const tags = this.tagRequest(request, instance, stream, span)
38+
span.addTags(tags)
39+
4040
const store = storage('legacy').getStore() || {}
4141
ctx.currentStore = { ...store, span }
4242

@@ -51,14 +51,23 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
5151

5252
const response = result?.response
5353
if (response) {
54-
const tags = this.tagResponse(response)
54+
const tags = this.tagResponse(response, span)
5555
span.addTags(tags)
5656
}
5757

5858
span.finish()
5959
}
6060

61-
tagRequest (request, instance, stream) {
61+
/**
62+
* Generate the request tags.
63+
*
64+
* @param {Object} request
65+
* @param {Object} instance
66+
* @param {boolean} stream
67+
* @param {Span} span
68+
* @returns {Object}
69+
*/
70+
tagRequest (request, instance, stream, span) {
6271
const model = extractModel(instance)
6372
const tags = {
6473
'vertexai.request.model': model
@@ -81,7 +90,7 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
8190
tags['vertexai.request.stream'] = true
8291
}
8392

84-
if (!this.isPromptCompletionSampled()) return tags
93+
if (!this.isPromptCompletionSampled(span)) return tags
8594

8695
const systemInstructions = extractSystemInstructions(instance)
8796

@@ -141,8 +150,16 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
141150
}
142151
}
143152

144-
tagResponse (response) {
153+
/**
154+
* Generate the response tags.
155+
*
156+
* @param {Object} response
157+
* @param {Span} span
158+
* @returns {Object}
159+
*/
160+
tagResponse (response, span) {
145161
const tags = {}
162+
const isSampled = this.isPromptCompletionSampled(span)
146163

147164
const candidates = response.candidates
148165
for (const [candidateIdx, candidate] of candidates.entries()) {
@@ -154,7 +171,7 @@ class GoogleCloudVertexAITracingPlugin extends TracingPlugin {
154171
const role = candidateContent.role
155172
tags[`vertexai.response.candidates.${candidateIdx}.content.role`] = role
156173

157-
if (!this.isPromptCompletionSampled()) continue
174+
if (!isSampled) continue
158175

159176
const parts = candidateContent.parts
160177
for (const [partIdx, part] of parts.entries()) {

‎packages/datadog-plugin-langchain/src/handlers/chain.js

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,10 @@
33
const LangChainHandler = require('./default')
44

55
class LangChainChainHandler extends LangChainHandler {
6-
getSpanStartTags (ctx) {
6+
getSpanStartTags (ctx, provider, span) {
77
const tags = {}
88

9-
if (!this.isPromptCompletionSampled()) return tags
9+
if (!this.isPromptCompletionSampled(span)) return tags
1010

1111
let inputs = ctx.args?.[0]
1212
inputs = Array.isArray(inputs) ? inputs : [inputs]
@@ -28,10 +28,10 @@ class LangChainChainHandler extends LangChainHandler {
2828
return tags
2929
}
3030

31-
getSpanEndTags (ctx) {
31+
getSpanEndTags (ctx, span) {
3232
const tags = {}
3333

34-
if (!this.isPromptCompletionSampled()) return tags
34+
if (!this.isPromptCompletionSampled(span)) return tags
3535

3636
let outputs = ctx.result
3737
outputs = Array.isArray(outputs) ? outputs : [outputs]

‎packages/datadog-plugin-langchain/src/handlers/embedding.js

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
const LangChainHandler = require('./default')
44

55
class LangChainEmbeddingHandler extends LangChainHandler {
6-
getSpanStartTags (ctx) {
6+
getSpanStartTags (ctx, provider, span) {
77
const tags = {}
88

99
const inputTexts = ctx.args?.[0]
1010

11-
const sampled = this.isPromptCompletionSampled()
11+
const sampled = this.isPromptCompletionSampled(span)
1212
if (typeof inputTexts === 'string') {
1313
// embed query
1414
if (sampled) {

‎packages/datadog-plugin-langchain/src/handlers/language_models/chat_model.js

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ const LangChainLanguageModelHandler = require('.')
55
const COMPLETIONS = 'langchain.response.completions'
66

77
class LangChainChatModelHandler extends LangChainLanguageModelHandler {
8-
getSpanStartTags (ctx, provider) {
8+
getSpanStartTags (ctx, provider, span) {
99
const tags = {}
1010

1111
const inputs = ctx.args?.[0]
@@ -15,7 +15,7 @@ class LangChainChatModelHandler extends LangChainLanguageModelHandler {
1515

1616
for (const messageIndex in messageSet) {
1717
const message = messageSet[messageIndex]
18-
if (this.isPromptCompletionSampled()) {
18+
if (this.isPromptCompletionSampled(span)) {
1919
tags[`langchain.request.messages.${messageSetIndex}.${messageIndex}.content`] =
2020
this.normalize(message.content) || ''
2121
}
@@ -39,11 +39,13 @@ class LangChainChatModelHandler extends LangChainLanguageModelHandler {
3939
return tags
4040
}
4141

42-
getSpanEndTags (ctx) {
42+
getSpanEndTags (ctx, span) {
4343
const { result } = ctx
4444

4545
const tags = {}
4646

47+
const sampled = this.isPromptCompletionSampled(span)
48+
4749
this.extractTokenMetrics(ctx.currentStore?.span, result)
4850

4951
for (const messageSetIdx in result?.generations) {
@@ -56,7 +58,7 @@ class LangChainChatModelHandler extends LangChainLanguageModelHandler {
5658
const message = chatCompletion.message
5759
let toolCalls = message.tool_calls
5860

59-
if (text && this.isPromptCompletionSampled()) {
61+
if (text && sampled) {
6062
tags[
6163
`${COMPLETIONS}.${messageSetIdx}.${chatCompletionIdx}.content`
6264
] = this.normalize(text)

‎packages/datadog-plugin-langchain/src/handlers/language_models/llm.js

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
const LangChainLanguageModelHandler = require('.')
44

55
class LangChainLLMHandler extends LangChainLanguageModelHandler {
6-
getSpanStartTags (ctx, provider) {
6+
getSpanStartTags (ctx, provider, span) {
77
const tags = {}
88

99
const prompts = ctx.args?.[0]
1010
for (const promptIdx in prompts) {
11-
if (!this.isPromptCompletionSampled()) continue
11+
if (!this.isPromptCompletionSampled(span)) continue
1212

1313
const prompt = prompts[promptIdx]
1414
tags[`langchain.request.prompts.${promptIdx}.content`] = this.normalize(prompt) || ''
@@ -30,16 +30,17 @@ class LangChainLLMHandler extends LangChainLanguageModelHandler {
3030
return tags
3131
}
3232

33-
getSpanEndTags (ctx) {
33+
getSpanEndTags (ctx, span) {
3434
const { result } = ctx
3535

3636
const tags = {}
37+
const sampled = this.isPromptCompletionSampled(span)
3738

3839
this.extractTokenMetrics(ctx.currentStore?.span, result)
3940

4041
for (const completionIdx in result?.generations) {
4142
const completion = result.generations[completionIdx]
42-
if (this.isPromptCompletionSampled()) {
43+
if (sampled) {
4344
tags[`langchain.response.completions.${completionIdx}.text`] = this.normalize(completion[0].text) || ''
4445
}
4546

‎packages/datadog-plugin-langchain/src/tracing.js

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,23 +49,24 @@ class BaseLangChainTracingPlugin extends TracingPlugin {
4949
const provider = handler.extractProvider(instance)
5050
const model = handler.extractModel(instance)
5151

52-
const tags = handler.getSpanStartTags(ctx, provider) || []
53-
54-
if (apiKey) tags[API_KEY] = apiKey
55-
if (provider) tags[PROVIDER] = provider
56-
if (model) tags[MODEL] = model
57-
if (type) tags[TYPE] = type
58-
5952
const span = this.startSpan('langchain.request', {
6053
service: this.config.service,
6154
resource,
6255
kind: 'client',
6356
meta: {
64-
[MEASURED]: 1,
65-
...tags
57+
[MEASURED]: 1
6658
}
6759
}, false)
6860

61+
const tags = handler.getSpanStartTags(ctx, provider, span) || []
62+
63+
if (apiKey) tags[API_KEY] = apiKey
64+
if (provider) tags[PROVIDER] = provider
65+
if (model) tags[MODEL] = model
66+
if (type) tags[TYPE] = type
67+
68+
span.addTags(tags)
69+
6970
const store = storage('legacy').getStore() || {}
7071
ctx.currentStore = { ...store, span }
7172

@@ -78,7 +79,7 @@ class BaseLangChainTracingPlugin extends TracingPlugin {
7879
const { type } = ctx
7980

8081
const handler = this.handlers[type]
81-
const tags = handler.getSpanEndTags(ctx) || {}
82+
const tags = handler.getSpanEndTags(ctx, span) || {}
8283

8384
span.addTags(tags)
8485

‎packages/datadog-plugin-openai/src/tracing.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,7 @@ class OpenAiTracingPlugin extends TracingPlugin {
326326
sendLog (methodName, span, tags, openaiStore, error) {
327327
if (!openaiStore) return
328328
if (!Object.keys(openaiStore).length) return
329-
if (!this.sampler.isSampled()) return
329+
if (!this.sampler.isSampled(span)) return
330330

331331
const log = {
332332
status: error ? 'error' : 'info',

‎packages/dd-trace/src/id.js

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ class Identifier {
2626
: toNumberString(this._buffer, radix)
2727
}
2828

29+
toBigInt () {
30+
return Buffer.from(this._buffer).readBigUInt64BE(0)
31+
}
32+
2933
toBuffer () {
3034
return this._buffer
3135
}

‎packages/dd-trace/src/opentracing/tracer.js

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,12 @@ class DatadogTracer {
107107
}
108108
}
109109

110+
/**
111+
* Get the span context from a span or a span context.
112+
*
113+
* @param {Span|SpanContext} spanContext
114+
* @returns {SpanContext}
115+
*/
110116
function getContext (spanContext) {
111117
if (spanContext instanceof Span) {
112118
spanContext = spanContext.context()

0 commit comments

Comments
 (0)