Skip to content

Commit 8a891c9

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 c43de0a commit 8a891c9

File tree

6 files changed

+117
-23
lines changed

6 files changed

+117
-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: 20 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,12 @@ def _default_instrumentor_builder(project_id: str):
119126
span_exporter=span_exporter,
120127
)
121128
)
129+
resource = opentelemetry_sdk_trace.Resource.create(
130+
attributes={
131+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
132+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
133+
}
134+
)
122135
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
123136
# Get the appropriate tracer provider:
124137
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -127,7 +140,7 @@ def _default_instrumentor_builder(project_id: str):
127140
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
128141
# If none of the above is set, we log a warning, and
129142
# create a tracer provider.
130-
if not tracer_provider:
143+
if AGENT_ENGINE_ID or not tracer_provider:
131144
from google.cloud.aiplatform import base
132145

133146
_LOGGER = base.Logger(__name__)
@@ -137,13 +150,17 @@ def _default_instrumentor_builder(project_id: str):
137150
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
138151
"or _PROXY_TRACER_PROVIDER."
139152
)
140-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
153+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
154+
resource=resource,
155+
)
141156
opentelemetry.trace.set_tracer_provider(tracer_provider)
142157
# Avoids AttributeError:
143158
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
144159
# attribute 'add_span_processor'.
145160
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
146-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
161+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
162+
resource=resource,
163+
)
147164
opentelemetry.trace.set_tracer_provider(tracer_provider)
148165
# Avoids OpenTelemetry client already exists error.
149166
_override_active_span_processor(

‎vertexai/agent_engines/templates/langchain.py

Lines changed: 20 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,12 @@ def _default_instrumentor_builder(project_id: str):
198205
span_exporter=span_exporter,
199206
)
200207
)
208+
resource = opentelemetry_sdk_trace.Resource.create(
209+
attributes={
210+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
211+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
212+
}
213+
)
201214
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
202215
# Get the appropriate tracer provider:
203216
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -206,7 +219,7 @@ def _default_instrumentor_builder(project_id: str):
206219
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
207220
# If none of the above is set, we log a warning, and
208221
# create a tracer provider.
209-
if not tracer_provider:
222+
if AGENT_ENGINE_ID or not tracer_provider:
210223
from google.cloud.aiplatform import base
211224

212225
_LOGGER = base.Logger(__name__)
@@ -216,13 +229,17 @@ def _default_instrumentor_builder(project_id: str):
216229
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
217230
"or _PROXY_TRACER_PROVIDER."
218231
)
219-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
232+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
233+
resource=resource,
234+
)
220235
opentelemetry.trace.set_tracer_provider(tracer_provider)
221236
# Avoids AttributeError:
222237
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
223238
# attribute 'add_span_processor'.
224239
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
225-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
240+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
241+
resource=resource,
242+
)
226243
opentelemetry.trace.set_tracer_provider(tracer_provider)
227244
# Avoids OpenTelemetry client already exists error.
228245
_override_active_span_processor(

‎vertexai/agent_engines/templates/langgraph.py

Lines changed: 20 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,12 @@ def _default_instrumentor_builder(project_id: str):
189196
span_exporter=span_exporter,
190197
)
191198
)
199+
resource = opentelemetry_sdk_trace.Resource.create(
200+
attributes={
201+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
202+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
203+
}
204+
)
192205
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
193206
# Get the appropriate tracer provider:
194207
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -197,7 +210,7 @@ def _default_instrumentor_builder(project_id: str):
197210
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
198211
# If none of the above is set, we log a warning, and
199212
# create a tracer provider.
200-
if not tracer_provider:
213+
if AGENT_ENGINE_ID or not tracer_provider:
201214
from google.cloud.aiplatform import base
202215

203216
base.Logger(__name__).warning(
@@ -206,13 +219,17 @@ def _default_instrumentor_builder(project_id: str):
206219
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
207220
"or _PROXY_TRACER_PROVIDER."
208221
)
209-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
222+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
223+
resource=resource,
224+
)
210225
opentelemetry.trace.set_tracer_provider(tracer_provider)
211226
# Avoids AttributeError:
212227
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
213228
# attribute 'add_span_processor'.
214229
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
215-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
230+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
231+
resource=resource,
232+
)
216233
opentelemetry.trace.set_tracer_provider(tracer_provider)
217234
# Avoids OpenTelemetry client already exists error.
218235
_override_active_span_processor(

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

Lines changed: 20 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,12 @@ 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+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
209+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
210+
}
211+
)
199212
tracer_provider = opentelemetry.trace.get_tracer_provider()
200213
# Get the appropriate tracer provider:
201214
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -204,7 +217,7 @@ def _default_instrumentor_builder(project_id: str):
204217
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
205218
# If none of the above is set, we log a warning, and
206219
# create a tracer provider.
207-
if not tracer_provider:
220+
if AGENT_ENGINE_ID or not tracer_provider:
208221
from google.cloud.aiplatform import base
209222

210223
_LOGGER = base.Logger(__name__)
@@ -214,13 +227,17 @@ def _default_instrumentor_builder(project_id: str):
214227
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
215228
"or _PROXY_TRACER_PROVIDER."
216229
)
217-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
230+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
231+
resource=resource,
232+
)
218233
opentelemetry.trace.set_tracer_provider(tracer_provider)
219234
# Avoids AttributeError:
220235
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
221236
# attribute 'add_span_processor'.
222237
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
223-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
238+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
239+
resource=resource,
240+
)
224241
opentelemetry.trace.set_tracer_provider(tracer_provider)
225242
# Avoids OpenTelemetry client already exists error.
226243
_override_active_span_processor(

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

Lines changed: 23 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,16 @@ def set_up(self):
416425
span_exporter=span_exporter,
417426
)
418427
)
428+
resource = opentelemetry_sdk_trace.Resource.create(
429+
attributes={
430+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
431+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
432+
}
433+
)
419434
tracer_provider: TracerProvider = (
420435
opentelemetry.trace.get_tracer_provider()
421436
)
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:
437+
if AGENT_ENGINE_ID or not tracer_provider:
430438
from google.cloud.aiplatform import base
431439

432440
_LOGGER = base.Logger(__name__)
@@ -436,13 +444,17 @@ def set_up(self):
436444
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
437445
"or _PROXY_TRACER_PROVIDER."
438446
)
439-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
447+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
448+
resource=resource,
449+
)
440450
opentelemetry.trace.set_tracer_provider(tracer_provider)
441451
# Avoids AttributeError:
442452
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
443453
# attribute 'add_span_processor'.
444454
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
445-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
455+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
456+
resource=resource,
457+
)
446458
opentelemetry.trace.set_tracer_provider(tracer_provider)
447459
# Avoids OpenTelemetry client already exists error.
448460
_override_active_span_processor(

0 commit comments

Comments
 (0)