Add head-based sampling for LLM Observability traces [MLOB-7815]

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -48,6 +48,8 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private static final String CONTEXT_VARIABLE_KEYS = "_dd_context_variable_keys";
   private static final String QUERY_VARIABLE_KEYS = "_dd_query_variable_keys";
   private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
+  private static final String SAMPLE_RATE_TAG_INTERNAL = "sample_rate";
+  private static final String SAMPLING_DECISION_TAG_INTERNAL = "sampling_decision";
 
   private static final String SERVICE = LLMOBS_TAG_PREFIX + "service";
   private static final String VERSION = LLMOBS_TAG_PREFIX + "version";
@@ -58,6 +60,8 @@ public class DDLLMObsSpan implements LLMObsSpan {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class);
 
+  private static final LLMObsSampler CONFIGURED_SAMPLER = LLMObsSampler.fromConfig();
+
   private final AgentSpan span;
   private final String spanKind;
   private final String mlApp;
@@ -73,6 +77,17 @@ public DDLLMObsSpan(
       String sessionId,
       @Nonnull String serviceName,
       WellKnownTags wellKnownTags) {
+    this(kind, spanName, mlApp, sessionId, serviceName, wellKnownTags, CONFIGURED_SAMPLER);
+  }
+
+  DDLLMObsSpan(
+      @Nonnull String kind,
+      String spanName,
+      @Nonnull String mlApp,
+      String sessionId,
+      @Nonnull String serviceName,
+      WellKnownTags wellKnownTags,
+      @Nonnull LLMObsSampler sampler) {
 
     if (null == spanName || spanName.isEmpty()) {
       spanName = kind;
@@ -101,11 +116,13 @@ public DDLLMObsSpan(
     spanKind = kind;
     this.mlApp = mlApp;
     span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp);
-    // Resolve effective parent_id and session_id from the LLMObs context, both gated on
-    // trace-id consistency. A stale context from a different trace (e.g. async boundary
-    // leakage) must not contribute either tag.
+    // Resolve effective parent_id, session_id and sampling decision from the LLMObs context, all
+    // gated on trace-id consistency. A stale context from a different trace (e.g. async boundary
+    // leakage) must not contribute any of them.
     AgentSpanContext parent = LLMObsContext.current();
     String parentSpanID = LLMObsContext.ROOT_SPAN_ID;
+    String sampleRate = null;
+    String samplingDecision = null;
     if (null != parent) {
       if (parent.getTraceId() != span.getTraceId()) {
         LOGGER.error(
@@ -125,16 +142,33 @@ public DDLLMObsSpan(
             sessionId = inherited;
           }
         }
+        // Inherit the sampling decision from the context if present.
+        sampleRate = LLMObsContext.currentSampleRate();
+        samplingDecision = LLMObsContext.currentSamplingDecision();
       }
     }
 
+    // Root of an LLMObs trace: decide once, keyed on the APM trace ID so that this decision is
+    // reproducible in any other service that observes the same trace at the same rate. Stamped at
+    // every rate, including the default of 1.0, matching dd-trace-py.
+    if (samplingDecision == null) {
+      sampleRate = sampler.formattedRate();
+      samplingDecision =
+          sampler.sample(span.getTraceId().toLong())
+              ? LLMObsContext.SAMPLING_DECISION_SAMPLED
+              : LLMObsContext.SAMPLING_DECISION_DROPPED;
+    }
+    span.setTag(LLMOBS_TAG_PREFIX + SAMPLE_RATE_TAG_INTERNAL, sampleRate);
+    span.setTag(LLMOBS_TAG_PREFIX + SAMPLING_DECISION_TAG_INTERNAL, samplingDecision);
+
     this.hasSessionId = sessionId != null && !sessionId.isEmpty();
     if (this.hasSessionId) {
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
     span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
-    // Propagate the effective sessionId to descendant LLMObs spans via the context.
-    scope = LLMObsContext.attach(span.spanContext(), sessionId);
+    // Propagate the effective sessionId and sampling decision to descendant LLMObs spans via the
+    // context.
+    scope = LLMObsContext.attach(span.spanContext(), sessionId, sampleRate, samplingDecision);
   }
 
   @Override

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsSampler.java

@@ -0,0 +1,93 @@
+package datadog.trace.llmobs.domain;
+
+import datadog.trace.api.Config;
+
+/**
+ * Head-based retention sampler for LLM Observability traces.
+ *
+ * <p>Duplicates the arithmetic in {@code DeterministicSampler} because {@code agent-llmobs} does
+ * not depend on {@code dd-trace-core}; {@code ApiSecurityDownstreamSamplerImpl} does the same.
+ */
+public final class LLMObsSampler {
+
+  private static final long KNUTH_FACTOR = 1111111111111111111L;
+  private static final double MAX = Math.pow(2, 64) - 1;
+
+  private final double rate;
+  private final long threshold;
+  private final String formattedRate;
+
+  public static LLMObsSampler fromConfig() {
+    return new LLMObsSampler(Config.get().getLlmObsSampleRate());
+  }
+
+  public LLMObsSampler(final double rate) {
+    // NaN is clamped to 1.0 rather than left alone: every comparison against it is false, which
+    // would otherwise leave the cutoff and the reported rate disagreeing about what happened.
+    final double bounded = Double.isNaN(rate) ? 1.0 : rate;
+    this.rate = bounded < 0.0 ? 0.0 : (bounded > 1.0 ? 1.0 : bounded);
+    this.threshold = cutoff(this.rate);
+    this.formattedRate = formatRate(this.rate);
+  }
+
+  /**
+   * The configured rate, formatted for the wire. Derived from the configured {@code double} rather
+   * than from a narrowed {@code float}, so the reported rate is the rate actually applied.
+   */
+  public String formattedRate() {
+    return formattedRate;
+  }
+
+  /**
+   * @param samplingId the low-order 64 bits of the APM trace ID.
+   * @return whether the trace is retained.
+   */
+  public boolean sample(final long samplingId) {
+    // unsigned 64 bit comparison with cutoff/threshold
+    return samplingId * KNUTH_FACTOR + Long.MIN_VALUE <= threshold;
+  }
+
+  private static long cutoff(final double rate) {
+    if (rate < 0.5) {
+      return (long) (rate * MAX) + Long.MIN_VALUE;
+    }
+    if (rate < 1.0) {
+      return (long) ((rate * MAX) + Long.MIN_VALUE);
+    }
+    return Long.MAX_VALUE;
+  }
+
+  /**
+   * Formats a sampling rate with up to 6 decimal digits of precision and no trailing zeros. Mirrors
+   * {@code format_rate} in dd-trace-py, which stamps {@code sample_rate} through the same helper it
+   * uses for {@code _dd.p.ksr}, so the two languages report the same rate as the same string.
+   */
+  static String formatRate(final double rate) {
+    if (rate <= 0.0) {
+      return "0";
+    }
+    if (rate >= 1.0) {
+      return "1";
+    }
+    long rounded = Math.round(rate * 1_000_000L);
+    if (rounded <= 0) {
+      return "0";
+    }
+    if (rounded >= 1_000_000L) {
+      return "1";
+    }
+    // Build "0.DDDDDD", then trim trailing zeros.
+    char[] chars = new char[8];
+    chars[0] = '0';
+    chars[1] = '.';
+    for (int i = 7; i >= 2; i--) {
+      chars[i] = (char) ('0' + (rounded % 10));
+      rounded /= 10;
+    }
+    int end = 8;
+    while (chars[end - 1] == '0') {
+      end--;
+    }
+    return new String(chars, 0, end);
+  }
+}

dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanSamplingTest.java

@@ -0,0 +1,133 @@
+package datadog.trace.llmobs.domain;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.api.WellKnownTags;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.core.CoreTracer;
+import java.lang.reflect.Field;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Covers where the head-based sampling decision is made and where it is inherited: the decision is
+ * computed once, at the root of an LLMObs trace, and every descendant reports it verbatim.
+ */
+class DDLLMObsSpanSamplingTest {
+  private static final String SAMPLE_RATE_TAG = "_ml_obs_tag.sample_rate";
+  private static final String SAMPLING_DECISION_TAG = "_ml_obs_tag.sampling_decision";
+
+  private static final Field SPAN_FIELD;
+
+  private static CoreTracer tracer;
+
+  static {
+    try {
+      SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span");
+      SPAN_FIELD.setAccessible(true);
+    } catch (ReflectiveOperationException error) {
+      throw new ExceptionInInitializerError(error);
+    }
+  }
+
+  @BeforeAll
+  static void installTracer() {
+    tracer = CoreTracer.builder().build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+  }
+
+  @AfterAll
+  static void closeTracer() {
+    TracerInstaller.forceInstallGlobalTracer(null);
+    tracer.close();
+  }
+
+  @Test
+  void stampsRetainedDecisionAtTheDefaultRate() throws IllegalAccessException {
+    // The fields are stamped at every rate, including 1.0, matching dd-trace-py.
+    DDLLMObsSpan llmObsSpan = newSpan(new LLMObsSampler(1.0));
+    try {
+      AgentSpan span = spanOf(llmObsSpan);
+      assertEquals("1", span.getTag(SAMPLING_DECISION_TAG));
+      assertEquals("1", span.getTag(SAMPLE_RATE_TAG));
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  @Test
+  void stampsDroppedDecisionOnRoot() throws IllegalAccessException {
+    DDLLMObsSpan llmObsSpan = newSpan(new LLMObsSampler(0.0));
+    try {
+      AgentSpan span = spanOf(llmObsSpan);
+      assertEquals("0", span.getTag(SAMPLING_DECISION_TAG));
+      assertEquals("0", span.getTag(SAMPLE_RATE_TAG));
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  @Test
+  void childInheritsTheRootDecisionInsteadOfRecomputingIt() throws IllegalAccessException {
+    // The child's sampler drops everything. If the decision were recomputed per span, the child
+    // would report "0" and the trace would be torn in half at the intake.
+    DDLLMObsSpan root = newSpan(new LLMObsSampler(1.0));
+    try {
+      AgentSpan rootSpan = spanOf(root);
+      // Inheritance is gated on the two spans sharing an APM trace, so the root's APM span has to
+      // be active for the child to be started under it.
+      try (AgentScope ignored = AgentTracer.activateSpan(rootSpan)) {
+        DDLLMObsSpan child = newSpan(new LLMObsSampler(0.0));
+        try {
+          AgentSpan childSpan = spanOf(child);
+          assertEquals("1", childSpan.getTag(SAMPLING_DECISION_TAG));
+          assertEquals(
+              rootSpan.getTag(SAMPLE_RATE_TAG),
+              childSpan.getTag(SAMPLE_RATE_TAG),
+              "every span in a trace must report the rate the decision was made at");
+        } finally {
+          child.finish();
+        }
+      }
+    } finally {
+      root.finish();
+    }
+  }
+
+  @Test
+  void childOfADroppedRootStaysDropped() throws IllegalAccessException {
+    // Symmetric case. Because a decision is stamped at every rate, "the context carries a decision"
+    // is an unambiguous signal, so a child never mistakes an inherited drop for being a root.
+    DDLLMObsSpan root = newSpan(new LLMObsSampler(0.0));
+    try {
+      try (AgentScope ignored = AgentTracer.activateSpan(spanOf(root))) {
+        DDLLMObsSpan child = newSpan(new LLMObsSampler(1.0));
+        try {
+          AgentSpan childSpan = spanOf(child);
+          assertEquals("0", childSpan.getTag(SAMPLING_DECISION_TAG));
+          assertEquals("0", childSpan.getTag(SAMPLE_RATE_TAG));
+        } finally {
+          child.finish();
+        }
+      }
+    } finally {
+      root.finish();
+    }
+  }
+
+  private static DDLLMObsSpan newSpan(LLMObsSampler sampler) {
+    WellKnownTags tags =
+        new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java");
+    return new DDLLMObsSpan(
+        Tags.LLMOBS_LLM_SPAN_KIND, "span", "ml-app", null, "service", tags, sampler);
+  }
+
+  private static AgentSpan spanOf(DDLLMObsSpan llmObsSpan) throws IllegalAccessException {
+    return (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+  }
+}