chore(llmobs): dac strip io from vertex #13693

jsimpher · 2025-06-17T16:56:34Z

Remove potentially sensitive i/o data from apm spans. This way, prompt and completion data will only appear on the llm obs spans, which are/will be subject to data access controls.

Mostly, this just removes io tag sets. A few things (mostly metrics) have llmobs tags dependent on span tags, so there is a bit more refactoring there.

Let me know if I removed anything that should really stay, or if I missed something that should be restricted.

Checklist

PR author has checked that all the criteria below are met
The PR description includes an overview of the change
The PR description articulates the motivation for the change
The change includes tests OR the PR description describes a testing strategy
The PR description notes risks associated with the change, if any
Newly-added code is easy to change
The change follows the library release note guidelines
The change includes or references documentation updates if necessary
Backport labels are set (if applicable)

Reviewer Checklist

Reviewer has checked that all the criteria below are met
Title is accurate
All changes are related to the pull request's stated goal
Avoids breaking API changes
Testing strategy adequately addresses listed risks
Newly-added code is easy to change
Release note makes sense to a user of the library
If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment
Backport labels are set in a manner that is consistent with the release branch maintenance policy

github-actions · 2025-06-17T16:57:36Z

CODEOWNERS have been resolved as:

.github/CODEOWNERS                                                      @DataDog/python-guild @DataDog/apm-core-python
ddtrace/contrib/internal/vertexai/_utils.py                             @DataDog/ml-observability
ddtrace/contrib/internal/vertexai/patch.py                              @DataDog/ml-observability
ddtrace/llmobs/_integrations/utils.py                                   @DataDog/ml-observability
ddtrace/llmobs/_integrations/vertexai.py                                @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_error.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_multiple_messages.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_stream.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_stream_error.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_stream_tool.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_system_prompt.json  @DataDog/apm-python @DataDog/ml-observability
tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_tool.json  @DataDog/apm-python @DataDog/ml-observability

github-actions · 2025-06-17T17:19:09Z

Bootstrap import analysis

Comparison of import times between this PR and base.

Summary

The average import time from this PR is: 281 ± 5 ms.

The average import time from base is: 281 ± 5 ms.

The import time difference between this PR and base is: 0.1 ± 0.2 ms.

The difference is not statistically significant (z = 0.27).

Import time breakdown

The following import paths have shrunk:

ddtrace.auto 1.826 ms (0.65%)

ddtrace.bootstrap.sitecustomize 1.154 ms (0.41%)

ddtrace.bootstrap.preload 1.154 ms (0.41%)

ddtrace.internal.remoteconfig.client 0.615 ms (0.22%)

ddtrace 0.672 ms (0.24%)

ddtrace.internal._unpatched 0.028 ms (0.01%)

json 0.028 ms (0.01%)

json.decoder 0.028 ms (0.01%)

re 0.028 ms (0.01%)

enum 0.028 ms (0.01%)

types 0.028 ms (0.01%)

pr-commenter · 2025-06-17T17:43:11Z

Benchmarks

Benchmark execution time: 2025-06-25 15:29:34

Comparing candidate commit 1e5ac47 in PR branch jsimpher/dac-strip-io-from-vertex with baseline commit 40f2c37 in branch main.

Found 0 performance improvements and 2 performance regressions! Performance is the same for 559 metrics, 3 unstable metrics.

scenario:iastaspects-replace_aspect

🟥 execution_time [+529.696ns; +610.039ns] or [+11.261%; +12.969%]

scenario:iastaspectsospath-ospathsplitdrive_aspect

🟥 execution_time [+392.961ns; +447.254ns] or [+10.806%; +12.299%]

jsimpher · 2025-06-24T13:51:13Z

ddtrace/contrib/internal/vertexai/patch.py

@@ -3,6 +3,7 @@
 from typing import Dict

 import vertexai
+from vertexai.generative_models import GenerativeModel  # noqa:F401


Removing this import (which used to exist on ./_utils.py but is no longer used there) breaks tests. As far as I can tell, it may have to do with vertex lazy loading, meaning that .generative_models may not exist yet when patching.
If this is a known thing that we have a standard way to deal with, let me know.

Hmm interesting, I would probably put this import in the patch function directly (as long as that works) so that we only import it when we actually do the patching.

Still fails. Looking closer, it will actually fail on the assertion that we are unpatched before calling .patch, since the thing we check isn't imported. There might be a way to move this to the test itself (a little tricky since its the base tests, not the vertex ones). Do you think that would be worth it?

tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_stream_tool.json

ncybul · 2025-06-24T17:40:25Z

ddtrace/contrib/internal/vertexai/_utils.py

-from vertexai.generative_models import GenerativeModel
-from vertexai.generative_models import Part
-
-from ddtrace.internal.utils import get_argument_value
-from ddtrace.llmobs._integrations.utils import get_generation_config_google
-from ddtrace.llmobs._integrations.utils import get_system_instructions_from_google_model
-from ddtrace.llmobs._integrations.utils import tag_request_content_part_google
-from ddtrace.llmobs._integrations.utils import tag_response_part_google


Looks like we are still using these helpers in the google generative ai integration. Once we remove the APM tagging there, we can get rid of these helpers altogether :)

ddtrace/contrib/internal/vertexai/_utils.py

ncybul

Left a few comments which should hopefully be quick. Going to approve for now, but if there's anything you are unsure about, feel free to reach out for another review!

brettlangdon · 2025-06-25T13:28:47Z

This should probably get a release note?

emmettbutler · 2025-06-25T15:10:39Z

.github/CODEOWNERS

@@ -177,6 +175,8 @@ tests/contrib/crewai                                          @DataDog/ml-observ
 tests/contrib/openai_agents                                   @DataDog/ml-observability
 tests/contrib/litellm                                         @DataDog/ml-observability
 .gitlab/tests/llmobs.yml                                      @DataDog/ml-observability
+# MLObs snapshot tests
+tests/snapshots/tests.contrib.vertexai.*                     @DataDog/apm-python @DataDog/ml-observability


I don't think apm-python needs to be included on this

Ended up pulling this into its own pr to reassign all the llmobs integration test snapshots; with your feedback integrated

jsimpher added 2 commits June 17, 2025 12:08

remove io from vertex ai integration

95c14e8

llmobs reads metrics from kwargs instead of putting them on span

89fad2c

jsimpher changed the title ~~Jsimpher/dac strip io from vertex~~ dac strip io from vertex Jun 17, 2025

jsimpher added 3 commits June 17, 2025 16:13

moved metrics logic to llmobs

fb280b6

fixed no response case

65398cc

update snapshots

29c1527

jsimpher changed the title ~~dac strip io from vertex~~ chore(llmobs): dac strip io from vertex Jun 22, 2025

jsimpher added 5 commits June 23, 2025 10:26

hoisted import

d97af97

ruff

f6a897f

move it back

f48dbfc

hopefully fix import

a537804

blakc

2f17324

jsimpher commented Jun 24, 2025

View reviewed changes

jsimpher and others added 2 commits June 24, 2025 09:52

added comment for unused import

3003ef6

Merge branch 'main' into jsimpher/dac-strip-io-from-vertex

42b0470

jsimpher marked this pull request as ready for review June 24, 2025 16:01

jsimpher requested review from a team as code owners June 24, 2025 16:01

jsimpher requested review from juanjux and mabdinur June 24, 2025 16:01

ncybul reviewed Jun 24, 2025

View reviewed changes

tests/snapshots/tests.contrib.vertexai.test_vertexai.test_vertexai_completion_stream_tool.json Outdated Show resolved Hide resolved

ncybul reviewed Jun 24, 2025

View reviewed changes

ddtrace/contrib/internal/vertexai/_utils.py Outdated Show resolved Hide resolved

ncybul approved these changes Jun 24, 2025

View reviewed changes

jsimpher added 2 commits June 24, 2025 14:05

removed stream boolean tag

a9ecd89

update snapshots

cee9268

update codeowners

1e5ac47

jsimpher requested review from a team as code owners June 25, 2025 14:39

emmettbutler approved these changes Jun 25, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

chore(llmobs): dac strip io from vertex #13693

chore(llmobs): dac strip io from vertex #13693

jsimpher commented Jun 17, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jun 17, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jun 17, 2025 •

edited

Loading

Uh oh!

pr-commenter bot commented Jun 17, 2025 •

edited

Loading

Uh oh!

jsimpher Jun 24, 2025

Uh oh!

ncybul Jun 24, 2025

Uh oh!

jsimpher Jun 24, 2025

Uh oh!

Uh oh!

ncybul Jun 24, 2025

Uh oh!

Uh oh!

ncybul left a comment •

edited

Loading

Uh oh!

brettlangdon commented Jun 25, 2025

Uh oh!

emmettbutler Jun 25, 2025

Uh oh!

jsimpher Jun 25, 2025

Uh oh!

Uh oh!

chore(llmobs): dac strip io from vertex #13693

Are you sure you want to change the base?

chore(llmobs): dac strip io from vertex #13693

Conversation

jsimpher commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklist

Reviewer Checklist

Uh oh!

github-actions bot commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Bootstrap import analysis

Summary

Import time breakdown

Uh oh!

pr-commenter bot commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmarks

scenario:iastaspects-replace_aspect

scenario:iastaspectsospath-ospathsplitdrive_aspect

Uh oh!

jsimpher Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

ncybul Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

jsimpher Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ncybul Jun 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ncybul left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brettlangdon commented Jun 25, 2025

Uh oh!

emmettbutler Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

jsimpher Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jsimpher commented Jun 17, 2025 •

edited

Loading

github-actions bot commented Jun 17, 2025 •

edited

Loading

github-actions bot commented Jun 17, 2025 •

edited

Loading

pr-commenter bot commented Jun 17, 2025 •

edited

Loading

ncybul left a comment •

edited

Loading