Skip to content

Commit a2b460f

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 186e6d8 commit a2b460f

File tree

4 files changed

+76
-21
lines changed

4 files changed

+76
-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 = {
769-
"telemetry.sdk.language": "python",
770-
"telemetry.sdk.name": "opentelemetry",
771-
"telemetry.sdk.version": "1.36.0",
772-
"gcp.project_id": "test-project",
773-
"cloud.account.id": "test-project",
781+
"cloud.account.id": _TEST_PROJECT_ID,
774782
"cloud.platform": "gcp.agent_engine",
775-
"service.name": "test_agent_id",
776-
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id",
777-
"service.instance.id": "12345678123456781234567812345678-123123123",
778783
"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",
779788
"some-attribute": "some-value",
789+
"telemetry.sdk.language": "python",
790+
"telemetry.sdk.name": "opentelemetry",
791+
"telemetry.sdk.version": "1.36.0",
780792
}
781793

782794
otlp_span_exporter_mock.assert_called_once_with(
@@ -785,6 +797,8 @@ def test_tracing_setup(
785797
headers=mock.ANY,
786798
)
787799

800+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT)
801+
788802
user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
789803
assert (
790804
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 = {
711-
"telemetry.sdk.language": "python",
712-
"telemetry.sdk.name": "opentelemetry",
713-
"telemetry.sdk.version": "1.36.0",
714-
"gcp.project_id": "test-project",
715-
"cloud.account.id": "test-project",
723+
"cloud.account.id": _TEST_PROJECT_ID,
716724
"cloud.platform": "gcp.agent_engine",
717-
"service.name": "test_agent_id",
718-
"cloud.resource_id": "//aiplatform.googleapis.com/projects/test-project/locations/us-central1/reasoningEngines/test_agent_id",
719-
"service.instance.id": "12345678123456781234567812345678-123123123",
720725
"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",
721730
"some-attribute": "some-value",
731+
"telemetry.sdk.language": "python",
732+
"telemetry.sdk.name": "opentelemetry",
733+
"telemetry.sdk.version": "1.36.0",
722734
}
723735

724736
otlp_span_exporter_mock.assert_called_once_with(
@@ -727,6 +739,8 @@ def test_tracing_setup(
727739
headers=mock.ANY,
728740
)
729741

742+
get_project_id_mock.assert_called_once_with(_TEST_PROJECT)
743+
730744
user_agent = otlp_span_exporter_mock.call_args.kwargs["headers"]["User-Agent"]
731745
assert (
732746
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,19 @@ 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+
293298
import os
294299

295300
def _warn_missing_dependency(
@@ -737,6 +742,7 @@ def set_up(self):
737742
project = self._tmpl_attrs.get("project")
738743
if project:
739744
os.environ["GOOGLE_CLOUD_PROJECT"] = project
745+
os.environ["GOOGLE_CLOUD_PROJECT_ID"] = str(self.project_id())
740746
location = self._tmpl_attrs.get("location")
741747
if location:
742748
os.environ["GOOGLE_CLOUD_LOCATION"] = location
@@ -790,11 +796,11 @@ def set_up(self):
790796
)
791797

792798
if custom_instrumentor and self._tracing_enabled():
793-
self._tmpl_attrs["instrumentor"] = custom_instrumentor(project)
799+
self._tmpl_attrs["instrumentor"] = custom_instrumentor(self.project_id())
794800

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

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -284,14 +284,19 @@ 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+
295300
import os
296301

297302
def _warn_missing_dependency(
@@ -711,7 +716,7 @@ def set_up(self):
711716
enable_logging = bool(self._telemetry_enabled())
712717

713718
self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder(
714-
project,
719+
self.project_id(),
715720
enable_tracing=self._tracing_enabled(),
716721
enable_logging=enable_logging,
717722
)
@@ -1539,3 +1544,11 @@ def _warn_if_telemetry_api_disabled(self):
15391544
r = session.post("https://telemetry.googleapis.com/v1/traces", data=None)
15401545
if "Telemetry API has not been used in project" in r.text:
15411546
_warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))
1547+
1548+
def project_id(self) -> Optional[str]:
1549+
if project := self._tmpl_attrs.get("project"):
1550+
from google.cloud.aiplatform.utils import resource_manager_utils
1551+
1552+
return resource_manager_utils.get_project_id(project)
1553+
1554+
return None

0 commit comments

Comments
 (0)