Skip to content

Commit b54a554

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
feat: Add Resource to tracer provider in Agent Engine templates
PiperOrigin-RevId: 763962855
1 parent 0a26a20 commit b54a554

File tree

6 files changed

+122
-23
lines changed

6 files changed

+122
-23
lines changed

‎vertexai/agent_engines/_utils.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -685,6 +685,20 @@ def _import_opentelemetry_or_warn() -> Optional[types.ModuleType]:
685685
return None
686686

687687

688+
def _import_opentelemetry_sdk_resources_or_warn() -> Optional[types.ModuleType]:
689+
"""Tries to import the opentelemetry.sdk.trace module."""
690+
try:
691+
import opentelemetry.sdk.resources # noqa:F401
692+
693+
return opentelemetry.sdk.resources
694+
except ImportError:
695+
LOGGER.warning(
696+
"Failed to import opentelemetry.sdk.resources. Please call "
697+
"'pip install google-cloud-aiplatform[agent_engines]'."
698+
)
699+
return None
700+
701+
688702
def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]:
689703
"""Tries to import the opentelemetry.sdk.trace module."""
690704
try:

‎vertexai/agent_engines/templates/ag2.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,16 +96,23 @@ def _default_instrumentor_builder(project_id: str):
9696
openinference_autogen = _utils._import_openinference_autogen_or_warn()
9797
opentelemetry = _utils._import_opentelemetry_or_warn()
9898
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
99+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
99100
if all(
100101
(
101102
cloud_trace_exporter,
102103
cloud_trace_v2,
103104
openinference_autogen,
104105
opentelemetry,
105106
opentelemetry_sdk_trace,
107+
opentelemetry_sdk_resources,
106108
)
107109
):
108110
import google.auth
111+
import os
112+
113+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
114+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
115+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
109116

110117
credentials, _ = google.auth.default()
111118
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -119,6 +126,13 @@ def _default_instrumentor_builder(project_id: str):
119126
span_exporter=span_exporter,
120127
)
121128
)
129+
resource = opentelemetry_sdk_trace.Resource.create(
130+
attributes={
131+
"gcp.project_id": project_id,
132+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
133+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
134+
}
135+
)
122136
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
123137
# Get the appropriate tracer provider:
124138
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -127,7 +141,7 @@ def _default_instrumentor_builder(project_id: str):
127141
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
128142
# If none of the above is set, we log a warning, and
129143
# create a tracer provider.
130-
if not tracer_provider:
144+
if AGENT_ENGINE_ID or not tracer_provider:
131145
from google.cloud.aiplatform import base
132146

133147
_LOGGER = base.Logger(__name__)
@@ -137,13 +151,17 @@ def _default_instrumentor_builder(project_id: str):
137151
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
138152
"or _PROXY_TRACER_PROVIDER."
139153
)
140-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
154+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
155+
resource=resource,
156+
)
141157
opentelemetry.trace.set_tracer_provider(tracer_provider)
142158
# Avoids AttributeError:
143159
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
144160
# attribute 'add_span_processor'.
145161
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
146-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
162+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
163+
resource=resource,
164+
)
147165
opentelemetry.trace.set_tracer_provider(tracer_provider)
148166
# Avoids OpenTelemetry client already exists error.
149167
_override_active_span_processor(

‎vertexai/agent_engines/templates/langchain.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,16 +175,23 @@ def _default_instrumentor_builder(project_id: str):
175175
openinference_langchain = _utils._import_openinference_langchain_or_warn()
176176
opentelemetry = _utils._import_opentelemetry_or_warn()
177177
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
178+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
178179
if all(
179180
(
180181
cloud_trace_exporter,
181182
cloud_trace_v2,
182183
openinference_langchain,
183184
opentelemetry,
184185
opentelemetry_sdk_trace,
186+
opentelemetry_sdk_resources,
185187
)
186188
):
187189
import google.auth
190+
import os
191+
192+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
193+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
194+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
188195

189196
credentials, _ = google.auth.default()
190197
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -198,6 +205,13 @@ def _default_instrumentor_builder(project_id: str):
198205
span_exporter=span_exporter,
199206
)
200207
)
208+
resource = opentelemetry_sdk_trace.Resource.create(
209+
attributes={
210+
"gcp.project_id": project_id,
211+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
212+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
213+
}
214+
)
201215
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
202216
# Get the appropriate tracer provider:
203217
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -206,7 +220,7 @@ def _default_instrumentor_builder(project_id: str):
206220
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
207221
# If none of the above is set, we log a warning, and
208222
# create a tracer provider.
209-
if not tracer_provider:
223+
if AGENT_ENGINE_ID or not tracer_provider:
210224
from google.cloud.aiplatform import base
211225

212226
_LOGGER = base.Logger(__name__)
@@ -216,13 +230,17 @@ def _default_instrumentor_builder(project_id: str):
216230
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
217231
"or _PROXY_TRACER_PROVIDER."
218232
)
219-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
233+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
234+
resource=resource,
235+
)
220236
opentelemetry.trace.set_tracer_provider(tracer_provider)
221237
# Avoids AttributeError:
222238
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
223239
# attribute 'add_span_processor'.
224240
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
225-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
241+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
242+
resource=resource,
243+
)
226244
opentelemetry.trace.set_tracer_provider(tracer_provider)
227245
# Avoids OpenTelemetry client already exists error.
228246
_override_active_span_processor(

‎vertexai/agent_engines/templates/langgraph.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -166,16 +166,23 @@ def _default_instrumentor_builder(project_id: str):
166166
openinference_langchain = _utils._import_openinference_langchain_or_warn()
167167
opentelemetry = _utils._import_opentelemetry_or_warn()
168168
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
169+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
169170
if all(
170171
(
171172
cloud_trace_exporter,
172173
cloud_trace_v2,
173174
openinference_langchain,
174175
opentelemetry,
175176
opentelemetry_sdk_trace,
177+
opentelemetry_sdk_resources,
176178
)
177179
):
178180
import google.auth
181+
import os
182+
183+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
184+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
185+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
179186

180187
credentials, _ = google.auth.default()
181188
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -189,6 +196,13 @@ def _default_instrumentor_builder(project_id: str):
189196
span_exporter=span_exporter,
190197
)
191198
)
199+
resource = opentelemetry_sdk_trace.Resource.create(
200+
attributes={
201+
"gcp.project_id": project_id,
202+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
203+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
204+
}
205+
)
192206
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
193207
# Get the appropriate tracer provider:
194208
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -197,7 +211,7 @@ def _default_instrumentor_builder(project_id: str):
197211
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
198212
# If none of the above is set, we log a warning, and
199213
# create a tracer provider.
200-
if not tracer_provider:
214+
if AGENT_ENGINE_ID or not tracer_provider:
201215
from google.cloud.aiplatform import base
202216

203217
base.Logger(__name__).warning(
@@ -206,13 +220,17 @@ def _default_instrumentor_builder(project_id: str):
206220
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
207221
"or _PROXY_TRACER_PROVIDER."
208222
)
209-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
223+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
224+
resource=resource,
225+
)
210226
opentelemetry.trace.set_tracer_provider(tracer_provider)
211227
# Avoids AttributeError:
212228
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
213229
# attribute 'add_span_processor'.
214230
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
215-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
231+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
232+
resource=resource,
233+
)
216234
opentelemetry.trace.set_tracer_provider(tracer_provider)
217235
# Avoids OpenTelemetry client already exists error.
218236
_override_active_span_processor(

‎vertexai/preview/reasoning_engines/templates/adk.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -176,15 +176,22 @@ def _default_instrumentor_builder(project_id: str):
176176
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
177177
opentelemetry = _utils._import_opentelemetry_or_warn()
178178
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
179+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
179180
if all(
180181
(
181182
cloud_trace_exporter,
182183
cloud_trace_v2,
183184
opentelemetry,
184185
opentelemetry_sdk_trace,
186+
opentelemetry_sdk_resources,
185187
)
186188
):
187189
import google.auth
190+
import os
191+
192+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
193+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
194+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
188195

189196
credentials, _ = google.auth.default()
190197
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -196,6 +203,13 @@ def _default_instrumentor_builder(project_id: str):
196203
span_processor = opentelemetry_sdk_trace.export.BatchSpanProcessor(
197204
span_exporter=span_exporter,
198205
)
206+
resource = opentelemetry_sdk_trace.Resource.create(
207+
attributes={
208+
"gcp.project_id": project_id,
209+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
210+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
211+
}
212+
)
199213
tracer_provider = opentelemetry.trace.get_tracer_provider()
200214
# Get the appropriate tracer provider:
201215
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -204,7 +218,7 @@ def _default_instrumentor_builder(project_id: str):
204218
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
205219
# If none of the above is set, we log a warning, and
206220
# create a tracer provider.
207-
if not tracer_provider:
221+
if AGENT_ENGINE_ID or not tracer_provider:
208222
from google.cloud.aiplatform import base
209223

210224
_LOGGER = base.Logger(__name__)
@@ -214,13 +228,17 @@ def _default_instrumentor_builder(project_id: str):
214228
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
215229
"or _PROXY_TRACER_PROVIDER."
216230
)
217-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
231+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
232+
resource=resource,
233+
)
218234
opentelemetry.trace.set_tracer_provider(tracer_provider)
219235
# Avoids AttributeError:
220236
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
221237
# attribute 'add_span_processor'.
222238
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
223-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
239+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
240+
resource=resource,
241+
)
224242
opentelemetry.trace.set_tracer_provider(tracer_provider)
225243
# Avoids OpenTelemetry client already exists error.
226244
_override_active_span_processor(

‎vertexai/preview/reasoning_engines/templates/llama_index.py

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,7 @@ def set_up(self):
384384
that can not be serialized.
385385
"""
386386
if self._enable_tracing:
387-
from vertexai.reasoning_engines import _utils
387+
from vertexai.agent_engines import _utils
388388

389389
cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
390390
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
@@ -393,16 +393,25 @@ def set_up(self):
393393
)
394394
opentelemetry = _utils._import_opentelemetry_or_warn()
395395
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
396+
opentelemetry_sdk_resources = (
397+
_utils._import_opentelemetry_sdk_resources_or_warn()
398+
)
396399
if all(
397400
(
398401
cloud_trace_exporter,
399402
cloud_trace_v2,
400403
openinference_llama_index,
401404
opentelemetry,
402405
opentelemetry_sdk_trace,
406+
opentelemetry_sdk_resources,
403407
)
404408
):
405409
import google.auth
410+
import os
411+
412+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
413+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
414+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
406415

407416
credentials, _ = google.auth.default()
408417
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
@@ -416,17 +425,17 @@ def set_up(self):
416425
span_exporter=span_exporter,
417426
)
418427
)
428+
resource = opentelemetry_sdk_trace.Resource.create(
429+
attributes={
430+
"gcp.project_id": self._project,
431+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
432+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
433+
}
434+
)
419435
tracer_provider: TracerProvider = (
420436
opentelemetry.trace.get_tracer_provider()
421437
)
422-
# Get the appropriate tracer provider:
423-
# 1. If _TRACER_PROVIDER is already set, use that.
424-
# 2. Otherwise, if the OTEL_PYTHON_TRACER_PROVIDER environment
425-
# variable is set, use that.
426-
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
427-
# If none of the above is set, we log a warning, and
428-
# create a tracer provider.
429-
if not tracer_provider:
438+
if AGENT_ENGINE_ID or not tracer_provider:
430439
from google.cloud.aiplatform import base
431440

432441
_LOGGER = base.Logger(__name__)
@@ -436,13 +445,17 @@ def set_up(self):
436445
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
437446
"or _PROXY_TRACER_PROVIDER."
438447
)
439-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
448+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
449+
resource=resource,
450+
)
440451
opentelemetry.trace.set_tracer_provider(tracer_provider)
441452
# Avoids AttributeError:
442453
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
443454
# attribute 'add_span_processor'.
444455
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
445-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
456+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
457+
resource=resource,
458+
)
446459
opentelemetry.trace.set_tracer_provider(tracer_provider)
447460
# Avoids OpenTelemetry client already exists error.
448461
_override_active_span_processor(

0 commit comments

Comments
 (0)