Skip to content

Commit 69b1ac5

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: correlate traces with logs in Cloud Trace panel on adk deploy agent_engine
PiperOrigin-RevId: 836220348
1 parent f1df7a2 commit 69b1ac5

File tree

4 files changed

+77
-21
lines changed

4 files changed

+77
-21
lines changed

‎tests/unit/vertex_adk/test_agent_engine_templates_adk.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ def __init__(self, name: str, model: str):
5555

5656
_TEST_LOCATION = "us-central1"
5757
_TEST_PROJECT = "test-project"
58+
_TEST_PROJECT_ID = "test-project-id"
5859
_TEST_API_KEY = "test-api-key"
5960
_TEST_MODEL = "gemini-2.0-flash"
6061
_TEST_USER_ID = "test_user_id"
@@ -224,6 +225,15 @@ def adk_version_mock():
224225
yield adk_version_mock
225226

226227

228+
@pytest.fixture
229+
def get_project_id_mock():
230+
with mock.patch(
231+
"google.cloud.aiplatform.aiplatform.utils.resource_manager_utils.get_project_id"
232+
) as get_project_id_mock:
233+
get_project_id_mock.return_value = _TEST_PROJECT_ID
234+
yield get_project_id_mock
235+
236+
227237
class _MockRunner:
228238
def run(self, *args, **kwargs):
229239
from google.adk.events import event
@@ -757,26 +767,28 @@ def test_tracing_setup(
757767
monkeypatch,
758768
trace_provider_mock: mock.Mock,
759769
otlp_span_exporter_mock: mock.Mock,
770+
get_project_id_mock: mock.Mock,
760771
):
761772
monkeypatch.setattr(
762773
"uuid.uuid4", lambda: uuid.UUID("12345678123456781234567812345678")
763774
)
764775
monkeypatch.setattr("os.getpid", lambda: 123123123)
765776
app = agent_engines.AdkApp(agent=_TEST_AGENT, enable_tracing=True)
777+
app._warn_if_telemetry_api_disabled = lambda: None
766778
app.set_up()
767779

768780
expected_attributes = {
781+
"cloud.account.id": _TEST_PROJECT_ID,
782+
"cloud.platform": "gcp.agent_engine",
783+
"cloud.region": "us-central1",
784+
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project-id/locations/us-central1/reasoningEngines/test_agent_id",
785+
"gcp.project_id": _TEST_PROJECT_ID,
786+
"service.instance.id": "12345678123456781234567812345678-123123123",
787+
"service.name": "test_agent_id",
788+
"some-attribute": "some-value",
769789
"telemetry.sdk.language": "python",
770790
"telemetry.sdk.name": "opentelemetry",
771791
"telemetry.sdk.version": "1.36.0",
772-
"gcp.project_id": "test-project",
773-
"cloud.account.id": "test-project",
774-
"cloud.provider": "gcp",
775-
"cloud.platform": "gcp.agent_engine",
776-
"service.name": "test_agent_id",
777-
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id",
778-
"service.instance.id": "12345678123456781234567812345678-123123123",
779-
"cloud.region": "us-central1",
780792
"some-attribute": "some-value",
781793
}
782794

@@ -786,6 +798,8 @@ def test_tracing_setup(
786798
headers=mock.ANY,
787799
)
788800

801+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT)
802+
789803
user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
790804
assert (
791805
re.fullmatch(

‎tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ def __init__(self, name: str, model: str):
4545

4646
_TEST_LOCATION = "us-central1"
4747
_TEST_PROJECT = "test-project"
48+
_TEST_PROJECT_ID = "test-project-id"
4849
_TEST_MODEL = "gemini-2.0-flash"
4950
_TEST_USER_ID = "test_user_id"
5051
_TEST_AGENT_NAME = "test_agent"
@@ -173,6 +174,15 @@ def adk_version_mock():
173174
yield adk_version_mock
174175

175176

177+
@pytest.fixture
178+
def get_project_id_mock():
179+
with mock.patch(
180+
"google.cloud.aiplatform.aiplatform.utils.resource_manager_utils.get_project_id"
181+
) as get_project_id_mock:
182+
get_project_id_mock.return_value = _TEST_PROJECT_ID
183+
yield get_project_id_mock
184+
185+
176186
class _MockRunner:
177187
def run(self, *args, **kwargs):
178188
from google.adk.events import event
@@ -699,26 +709,28 @@ def test_tracing_setup(
699709
monkeypatch: pytest.MonkeyPatch,
700710
trace_provider_mock: mock.Mock,
701711
otlp_span_exporter_mock: mock.Mock,
712+
get_project_id_mock: mock.Mock,
702713
):
703714
monkeypatch.setattr(
704715
"uuid.uuid4", lambda: uuid.UUID("12345678123456781234567812345678")
705716
)
706717
monkeypatch.setattr("os.getpid", lambda: 123123123)
707718
app = reasoning_engines.AdkApp(agent=_TEST_AGENT, enable_tracing=True)
719+
app._warn_if_telemetry_api_disabled = lambda: None
708720
app.set_up()
709721

710722
expected_attributes = {
723+
"cloud.account.id": _TEST_PROJECT_ID,
724+
"cloud.platform": "gcp.agent_engine",
725+
"cloud.region": "us-central1",
726+
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project-id/locations/us-central1/reasoningEngines/test_agent_id",
727+
"gcp.project_id": _TEST_PROJECT_ID,
728+
"service.instance.id": "12345678123456781234567812345678-123123123",
729+
"service.name": "test_agent_id",
730+
"some-attribute": "some-value",
711731
"telemetry.sdk.language": "python",
712732
"telemetry.sdk.name": "opentelemetry",
713733
"telemetry.sdk.version": "1.36.0",
714-
"gcp.project_id": "test-project",
715-
"cloud.account.id": "test-project",
716-
"cloud.provider": "gcp",
717-
"cloud.platform": "gcp.agent_engine",
718-
"service.name": "test_agent_id",
719-
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id",
720-
"service.instance.id": "12345678123456781234567812345678-123123123",
721-
"cloud.region": "us-central1",
722734
"some-attribute": "some-value",
723735
}
724736

@@ -728,6 +740,8 @@ def test_tracing_setup(
728740
headers=mock.ANY,
729741
)
730742

743+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT)
744+
731745
user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
732746
assert (
733747
re.fullmatch(

‎vertexai/agent_engines/templates/adk.py‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -282,14 +282,20 @@ async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool):
282282

283283

284284
def _default_instrumentor_builder(
285-
project_id: str,
285+
project_id: Optional[str],
286286
*,
287287
enable_tracing: bool = False,
288288
enable_logging: bool = False,
289289
):
290290
if not enable_tracing and not enable_logging:
291291
return None
292292

293+
if project_id is None:
294+
_warn(
295+
"telemetry is only supported when project is specified, proceeding with no telemetry"
296+
)
297+
return None
298+
293299
import os
294300

295301
def _warn_missing_dependency(
@@ -791,11 +797,11 @@ def set_up(self):
791797
)
792798

793799
if custom_instrumentor and self._tracing_enabled():
794-
self._tmpl_attrs["instrumentor"] = custom_instrumentor(project)
800+
self._tmpl_attrs["instrumentor"] = custom_instrumentor(self.project_id())
795801

796802
if not custom_instrumentor:
797803
self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder(
798-
project,
804+
self.project_id(),
799805
enable_tracing=self._tracing_enabled(),
800806
enable_logging=enable_logging,
801807
)
@@ -1607,3 +1613,11 @@ def _warn_if_telemetry_api_disabled(self):
16071613
r = session.post("https://telemetry.googleapis.com/v1/traces", data=None)
16081614
if "Telemetry API has not been used in project" in r.text:
16091615
_warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))
1616+
1617+
def project_id(self) -> Optional[str]:
1618+
if project := self._tmpl_attrs.get("project"):
1619+
from google.cloud.aiplatform.utils import resource_manager_utils
1620+
1621+
return resource_manager_utils.get_project_id(project)
1622+
1623+
return None

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -284,14 +284,20 @@ async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool):
284284

285285

286286
def _default_instrumentor_builder(
287-
project_id: str,
287+
project_id: Optional[str],
288288
*,
289289
enable_tracing: bool = False,
290290
enable_logging: bool = False,
291291
):
292292
if not enable_tracing and not enable_logging:
293293
return None
294294

295+
if project_id is None:
296+
_warn(
297+
"telemetry is only supported when project is specified, proceeding with no telemetry"
298+
)
299+
return None
300+
295301
import os
296302

297303
def _warn_missing_dependency(
@@ -712,7 +718,7 @@ def set_up(self):
712718
enable_logging = bool(self._telemetry_enabled())
713719

714720
self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder(
715-
project,
721+
self.project_id(),
716722
enable_tracing=self._tracing_enabled(),
717723
enable_logging=enable_logging,
718724
)
@@ -1540,3 +1546,11 @@ def _warn_if_telemetry_api_disabled(self):
15401546
r = session.post("https://telemetry.googleapis.com/v1/traces", data=None)
15411547
if "Telemetry API has not been used in project" in r.text:
15421548
_warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))
1549+
1550+
def project_id(self) -> Optional[str]:
1551+
if project := self._tmpl_attrs.get("project"):
1552+
from google.cloud.aiplatform.utils import resource_manager_utils
1553+
1554+
return resource_manager_utils.get_project_id(project)
1555+
1556+
return None

0 commit comments

Comments
 (0)