Skip to content

Commit 62a5746

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 62a5746

File tree

7 files changed

+163
-72
lines changed

7 files changed

+163
-72
lines changed

‎setup.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -146,20 +146,18 @@
146146

147147
reasoning_engine_extra_require = [
148148
"cloudpickle >= 3.0, < 4.0",
149-
"google-cloud-trace < 2",
150149
"opentelemetry-sdk < 2",
151-
"opentelemetry-exporter-gcp-trace < 2",
150+
"opentelemetry-exporter-otlp < 2",
152151
"pydantic >= 2.11.1, < 3",
153152
"typing_extensions",
154153
]
155154

156155
agent_engines_extra_require = [
157156
"packaging >= 24.0",
158157
"cloudpickle >= 3.0, < 4.0",
159-
"google-cloud-trace < 2",
160158
"google-cloud-logging < 4",
161159
"opentelemetry-sdk < 2",
162-
"opentelemetry-exporter-gcp-trace < 2",
160+
"opentelemetry-exporter-otlp < 2",
163161
"pydantic >= 2.11.1, < 3",
164162
"typing_extensions",
165163
]

‎vertexai/agent_engines/_utils.py

Lines changed: 28 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:
@@ -713,6 +727,20 @@ def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]:
713727
return None
714728

715729

730+
def _import_otlp_trace_exporter_or_warn() -> Optional[types.ModuleType]:
731+
"""Tries to import the google.cloud.trace_v2 module."""
732+
try:
733+
import opentelemetry.exporter.otlp.proto.http.trace_exporter
734+
735+
return opentelemetry.exporter.otlp.proto.http.trace_exporter
736+
except ImportError:
737+
LOGGER.warning(
738+
"opentelemetry-exporter-otlp is not installed. Please "
739+
"call 'pip install google-cloud-aiplatform[agent_engines]'."
740+
)
741+
return None
742+
743+
716744
def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]:
717745
"""Tries to import the opentelemetry.exporter.cloud_trace module."""
718746
try:

‎vertexai/agent_engines/templates/ag2.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -91,34 +91,44 @@ def _default_runnable_builder(
9191
def _default_instrumentor_builder(project_id: str):
9292
from vertexai.agent_engines import _utils
9393

94-
cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
95-
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
94+
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
9695
openinference_autogen = _utils._import_openinference_autogen_or_warn()
9796
opentelemetry = _utils._import_opentelemetry_or_warn()
9897
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
98+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
9999
if all(
100100
(
101-
cloud_trace_exporter,
102-
cloud_trace_v2,
101+
otlp_trace_exporter,
103102
openinference_autogen,
104103
opentelemetry,
105104
opentelemetry_sdk_trace,
105+
opentelemetry_sdk_resources,
106106
)
107107
):
108108
import google.auth
109+
import os
110+
111+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
112+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
113+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
109114

110115
credentials, _ = google.auth.default()
111-
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
112-
project_id=project_id,
113-
client=cloud_trace_v2.TraceServiceClient(
114-
credentials=credentials.with_quota_project(project_id),
115-
),
116+
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
117+
session=google.auth.transport.requests.AuthorizedSession(credentials),
118+
endpoint="https://telemetry.googleapis.com/v1/traces",
116119
)
117120
span_processor: SpanProcessor = (
118121
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
119122
span_exporter=span_exporter,
120123
)
121124
)
125+
resource = opentelemetry_sdk_trace.Resource.create(
126+
attributes={
127+
"gcp.project_id": project_id,
128+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
129+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
130+
}
131+
)
122132
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
123133
# Get the appropriate tracer provider:
124134
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -127,7 +137,7 @@ def _default_instrumentor_builder(project_id: str):
127137
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
128138
# If none of the above is set, we log a warning, and
129139
# create a tracer provider.
130-
if not tracer_provider:
140+
if AGENT_ENGINE_ID or not tracer_provider:
131141
from google.cloud.aiplatform import base
132142

133143
_LOGGER = base.Logger(__name__)
@@ -137,13 +147,17 @@ def _default_instrumentor_builder(project_id: str):
137147
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
138148
"or _PROXY_TRACER_PROVIDER."
139149
)
140-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
150+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
151+
resource=resource,
152+
)
141153
opentelemetry.trace.set_tracer_provider(tracer_provider)
142154
# Avoids AttributeError:
143155
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
144156
# attribute 'add_span_processor'.
145157
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
146-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
158+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
159+
resource=resource,
160+
)
147161
opentelemetry.trace.set_tracer_provider(tracer_provider)
148162
# Avoids OpenTelemetry client already exists error.
149163
_override_active_span_processor(

‎vertexai/agent_engines/templates/langchain.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -170,34 +170,44 @@ def _default_runnable_builder(
170170
def _default_instrumentor_builder(project_id: str):
171171
from vertexai.agent_engines import _utils
172172

173-
cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
174-
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
173+
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
175174
openinference_langchain = _utils._import_openinference_langchain_or_warn()
176175
opentelemetry = _utils._import_opentelemetry_or_warn()
177176
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
177+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
178178
if all(
179179
(
180-
cloud_trace_exporter,
181-
cloud_trace_v2,
180+
otlp_trace_exporter,
182181
openinference_langchain,
183182
opentelemetry,
184183
opentelemetry_sdk_trace,
184+
opentelemetry_sdk_resources,
185185
)
186186
):
187187
import google.auth
188+
import os
189+
190+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
191+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
192+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
188193

189194
credentials, _ = google.auth.default()
190-
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
191-
project_id=project_id,
192-
client=cloud_trace_v2.TraceServiceClient(
193-
credentials=credentials.with_quota_project(project_id),
194-
),
195+
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
196+
session=google.auth.transport.requests.AuthorizedSession(credentials),
197+
endpoint="https://telemetry.googleapis.com/v1/traces",
195198
)
196199
span_processor: SpanProcessor = (
197200
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
198201
span_exporter=span_exporter,
199202
)
200203
)
204+
resource = opentelemetry_sdk_trace.Resource.create(
205+
attributes={
206+
"gcp.project_id": project_id,
207+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
208+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
209+
}
210+
)
201211
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
202212
# Get the appropriate tracer provider:
203213
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -206,7 +216,7 @@ def _default_instrumentor_builder(project_id: str):
206216
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
207217
# If none of the above is set, we log a warning, and
208218
# create a tracer provider.
209-
if not tracer_provider:
219+
if AGENT_ENGINE_ID or not tracer_provider:
210220
from google.cloud.aiplatform import base
211221

212222
_LOGGER = base.Logger(__name__)
@@ -216,13 +226,17 @@ def _default_instrumentor_builder(project_id: str):
216226
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
217227
"or _PROXY_TRACER_PROVIDER."
218228
)
219-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
229+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
230+
resource=resource,
231+
)
220232
opentelemetry.trace.set_tracer_provider(tracer_provider)
221233
# Avoids AttributeError:
222234
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
223235
# attribute 'add_span_processor'.
224236
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
225-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
237+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
238+
resource=resource,
239+
)
226240
opentelemetry.trace.set_tracer_provider(tracer_provider)
227241
# Avoids OpenTelemetry client already exists error.
228242
_override_active_span_processor(

‎vertexai/agent_engines/templates/langgraph.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -161,34 +161,44 @@ def _default_runnable_builder(
161161
def _default_instrumentor_builder(project_id: str):
162162
from vertexai.agent_engines import _utils
163163

164-
cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
165-
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
164+
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
166165
openinference_langchain = _utils._import_openinference_langchain_or_warn()
167166
opentelemetry = _utils._import_opentelemetry_or_warn()
168167
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
168+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
169169
if all(
170170
(
171-
cloud_trace_exporter,
172-
cloud_trace_v2,
171+
otlp_trace_exporter,
173172
openinference_langchain,
174173
opentelemetry,
175174
opentelemetry_sdk_trace,
175+
opentelemetry_sdk_resources,
176176
)
177177
):
178178
import google.auth
179+
import os
180+
181+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
182+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
183+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
179184

180185
credentials, _ = google.auth.default()
181-
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
182-
project_id=project_id,
183-
client=cloud_trace_v2.TraceServiceClient(
184-
credentials=credentials.with_quota_project(project_id),
185-
),
186+
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
187+
session=google.auth.transport.requests.AuthorizedSession(credentials),
188+
endpoint="https://telemetry.googleapis.com/v1/traces",
186189
)
187190
span_processor: SpanProcessor = (
188191
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
189192
span_exporter=span_exporter,
190193
)
191194
)
195+
resource = opentelemetry_sdk_trace.Resource.create(
196+
attributes={
197+
"gcp.project_id": project_id,
198+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
199+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
200+
}
201+
)
192202
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
193203
# Get the appropriate tracer provider:
194204
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -197,7 +207,7 @@ def _default_instrumentor_builder(project_id: str):
197207
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
198208
# If none of the above is set, we log a warning, and
199209
# create a tracer provider.
200-
if not tracer_provider:
210+
if AGENT_ENGINE_ID or not tracer_provider:
201211
from google.cloud.aiplatform import base
202212

203213
base.Logger(__name__).warning(
@@ -206,13 +216,17 @@ def _default_instrumentor_builder(project_id: str):
206216
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
207217
"or _PROXY_TRACER_PROVIDER."
208218
)
209-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
219+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
220+
resource=resource,
221+
)
210222
opentelemetry.trace.set_tracer_provider(tracer_provider)
211223
# Avoids AttributeError:
212224
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
213225
# attribute 'add_span_processor'.
214226
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
215-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
227+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
228+
resource=resource,
229+
)
216230
opentelemetry.trace.set_tracer_provider(tracer_provider)
217231
# Avoids OpenTelemetry client already exists error.
218232
_override_active_span_processor(

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

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -172,30 +172,40 @@ def dump(self) -> Dict[str, Any]:
172172
def _default_instrumentor_builder(project_id: str):
173173
from vertexai.agent_engines import _utils
174174

175-
cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
176-
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
175+
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
177176
opentelemetry = _utils._import_opentelemetry_or_warn()
178177
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
178+
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
179179
if all(
180180
(
181-
cloud_trace_exporter,
182-
cloud_trace_v2,
181+
otlp_trace_exporter,
183182
opentelemetry,
184183
opentelemetry_sdk_trace,
184+
opentelemetry_sdk_resources,
185185
)
186186
):
187187
import google.auth
188+
import os
189+
190+
SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
191+
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
192+
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")
188193

189194
credentials, _ = google.auth.default()
190-
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
191-
project_id=project_id,
192-
client=cloud_trace_v2.TraceServiceClient(
193-
credentials=credentials.with_quota_project(project_id),
194-
),
195+
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
196+
session=google.auth.transport.requests.AuthorizedSession(credentials),
197+
endpoint="https://telemetry.googleapis.com/v1/traces",
195198
)
196199
span_processor = opentelemetry_sdk_trace.export.BatchSpanProcessor(
197200
span_exporter=span_exporter,
198201
)
202+
resource = opentelemetry_sdk_trace.Resource.create(
203+
attributes={
204+
"gcp.project_id": project_id,
205+
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
206+
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
207+
}
208+
)
199209
tracer_provider = opentelemetry.trace.get_tracer_provider()
200210
# Get the appropriate tracer provider:
201211
# 1. If _TRACER_PROVIDER is already set, use that.
@@ -204,7 +214,7 @@ def _default_instrumentor_builder(project_id: str):
204214
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
205215
# If none of the above is set, we log a warning, and
206216
# create a tracer provider.
207-
if not tracer_provider:
217+
if AGENT_ENGINE_ID or not tracer_provider:
208218
from google.cloud.aiplatform import base
209219

210220
_LOGGER = base.Logger(__name__)
@@ -214,13 +224,17 @@ def _default_instrumentor_builder(project_id: str):
214224
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
215225
"or _PROXY_TRACER_PROVIDER."
216226
)
217-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
227+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
228+
resource=resource,
229+
)
218230
opentelemetry.trace.set_tracer_provider(tracer_provider)
219231
# Avoids AttributeError:
220232
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
221233
# attribute 'add_span_processor'.
222234
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
223-
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
235+
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
236+
resource=resource,
237+
)
224238
opentelemetry.trace.set_tracer_provider(tracer_provider)
225239
# Avoids OpenTelemetry client already exists error.
226240
_override_active_span_processor(

0 commit comments

Comments
 (0)