From bd76b46ce296409d929ae69c5c43347c73e7b365 Mon Sep 17 00:00:00 2001 From: Max Ind Date: Mon, 6 Oct 2025 05:13:02 -0700 Subject: [PATCH] feat(otel): Switch CloudTraceSpanExporter to telemetry.googleapis.com PiperOrigin-RevId: 815675872 --- pyproject.toml | 2 +- src/google/adk/cli/adk_web_server.py | 22 ++++--- src/google/adk/telemetry/google_cloud.py | 63 ++++++++++++++----- src/google/adk/telemetry/setup.py | 6 +- .../unittests/telemetry/test_google_cloud.py | 34 ++++++++++ 5 files changed, 98 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8dc92258..14756309 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,8 +52,8 @@ dependencies = [ "python-dateutil>=2.9.0.post0, <3.0.0", # For Vertext AI Session Service "python-dotenv>=1.0.0, <2.0.0", # To manage environment variables "requests>=2.32.4, <3.0.0", - "sqlalchemy>=2.0, <3.0.0", # SQL database ORM "sqlalchemy-spanner>=1.14.0", # Spanner database session service + "sqlalchemy>=2.0, <3.0.0", # SQL database ORM "starlette>=0.46.2, <1.0.0", # For FastAPI CLI "tenacity>=8.0.0, <9.0.0", # For Retry management "typing-extensions>=4.5, <5", diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py index 2a2dbc05..d503e762 100644 --- a/src/google/adk/cli/adk_web_server.py +++ b/src/google/adk/cli/adk_web_server.py @@ -292,8 +292,9 @@ def _setup_telemetry( else: # Old logic - to be removed when above leaves experimental. tracer_provider = TracerProvider() - for exporter in internal_exporters: - tracer_provider.add_span_processor(exporter) + if internal_exporters is not None: + for exporter in internal_exporters: + tracer_provider.add_span_processor(exporter) trace.set_tracer_provider(tracer_provider=tracer_provider) @@ -312,10 +313,10 @@ def _otel_env_vars_enabled() -> bool: def _setup_gcp_telemetry_experimental( internal_exporters: list[SpanProcessor] = None, ): - from ..telemetry.setup import maybe_set_otel_providers + if typing.TYPE_CHECKING: + from ..telemetry.setup import OTelHooks - otel_hooks_to_add = [] - otel_resource = None + otel_hooks_to_add: list[OTelHooks] = [] if internal_exporters: from ..telemetry.setup import OTelHooks @@ -323,8 +324,13 @@ def _setup_gcp_telemetry_experimental( # Register ADK-specific exporters in trace provider. otel_hooks_to_add.append(OTelHooks(span_processors=internal_exporters)) + import google.auth + from ..telemetry.google_cloud import get_gcp_exporters from ..telemetry.google_cloud import get_gcp_resource + from ..telemetry.setup import maybe_set_otel_providers + + credentials, project_id = google.auth.default() otel_hooks_to_add.append( get_gcp_exporters( @@ -334,12 +340,14 @@ def _setup_gcp_telemetry_experimental( # TODO - reenable metrics once errors during shutdown are fixed. enable_cloud_metrics=False, enable_cloud_logging=True, + google_auth=(credentials, project_id), ) ) - otel_resource = get_gcp_resource() + otel_resource = get_gcp_resource(project_id) maybe_set_otel_providers( - otel_hooks_to_setup=otel_hooks_to_add, otel_resource=otel_resource + otel_hooks_to_setup=otel_hooks_to_add, + otel_resource=otel_resource, ) _setup_instrumentation_lib_if_installed() diff --git a/src/google/adk/telemetry/google_cloud.py b/src/google/adk/telemetry/google_cloud.py index ec78cf16..64016d64 100644 --- a/src/google/adk/telemetry/google_cloud.py +++ b/src/google/adk/telemetry/google_cloud.py @@ -15,6 +15,9 @@ from __future__ import annotations import logging +from typing import cast +from typing import Optional +from typing import TYPE_CHECKING import google.auth from opentelemetry.sdk._logs import LogRecordProcessor @@ -29,6 +32,9 @@ from opentelemetry.sdk.trace.export import BatchSpanProcessor from ..utils.feature_decorator import experimental from .setup import OTelHooks +if TYPE_CHECKING: + from google.auth.credentials import Credentials + logger = logging.getLogger('google_adk.' + __name__) @@ -37,6 +43,7 @@ def get_gcp_exporters( enable_cloud_tracing: bool = False, enable_cloud_metrics: bool = False, enable_cloud_logging: bool = False, + google_auth: Optional[tuple[Credentials, str]] = None, ) -> OTelHooks: """Returns GCP OTel exporters to be used in the app. @@ -44,8 +51,16 @@ def get_gcp_exporters( enable_tracing: whether to enable tracing to Cloud Trace. enable_metrics: whether to enable raporting metrics to Cloud Monitoring. enable_logging: whether to enable sending logs to Cloud Logging. + google_auth: optional custom credentials and project_id. google.auth.default() used when this is omitted. """ - _, project_id = google.auth.default() + + credentials, project_id = ( + google_auth if google_auth is not None else google.auth.default() + ) + if TYPE_CHECKING: + credentials = cast(Credentials, credentials) + project_id = cast(str, project_id) + if not project_id: logger.warning( 'Cannot determine GCP Project. OTel GCP Exporters cannot be set up.' @@ -53,18 +68,18 @@ def get_gcp_exporters( ) return OTelHooks() - span_processors = [] + span_processors: list[SpanProcessor] = [] if enable_cloud_tracing: - exporter = _get_gcp_span_exporter(project_id) + exporter = _get_gcp_span_exporter(credentials) span_processors.append(exporter) - metric_readers = [] + metric_readers: list[MetricReader] = [] if enable_cloud_metrics: exporter = _get_gcp_metrics_exporter(project_id) if exporter: metric_readers.append(exporter) - log_record_processors = [] + log_record_processors: list[LogRecordProcessor] = [] if enable_cloud_logging: exporter = _get_gcp_logs_exporter(project_id) if exporter: @@ -77,10 +92,18 @@ def get_gcp_exporters( ) -def _get_gcp_span_exporter(project_id: str) -> SpanProcessor: - from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter +def _get_gcp_span_exporter(credentials: Credentials) -> SpanProcessor: + """Adds OTEL span exporter to telemetry.googleapis.com""" - return BatchSpanProcessor(CloudTraceSpanExporter(project_id=project_id)) + from google.auth.transport.requests import AuthorizedSession + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + + return BatchSpanProcessor( + OTLPSpanExporter( + session=AuthorizedSession(credentials=credentials), + endpoint='https://telemetry.googleapis.com/v1/traces', + ) + ) def _get_gcp_metrics_exporter(project_id: str) -> MetricReader: @@ -101,15 +124,22 @@ def _get_gcp_logs_exporter(project_id: str) -> LogRecordProcessor: ) -def get_gcp_resource() -> Resource: - # The OTELResourceDetector populates resource labels from - # environment variables like OTEL_SERVICE_NAME and OTEL_RESOURCE_ATTRIBUTES. - # Then the GCP detector adds attributes corresponding to a correct - # monitored resource if ADK runs on one of supported platforms - # (e.g. GCE, GKE, CloudRun). - - resource = OTELResourceDetector().detect() +def get_gcp_resource(project_id: Optional[str] = None) -> Resource: + """Returns OTEL with attributes specified in the following order (attributes specified later, overwrite those specified earlier): + 1. Populates gcp.project_id attribute from the project_id argument if present. + 2. OTELResourceDetector populates resource labels from environment variables like OTEL_SERVICE_NAME and OTEL_RESOURCE_ATTRIBUTES. + 3. GCP detector adds attributes corresponding to a correct monitored resource if ADK runs on one of supported platforms (e.g. GCE, GKE, CloudRun). + Args: + project_id: project id to fill out as `gcp.project_id` on the OTEL resource. + This may be overwritten by OTELResourceDetector, if `gcp.project_id` is present in `OTEL_RESOURCE_ATTRIBUTES` env var. + """ + resource = Resource( + attributes={'gcp.project_id': project_id} + if project_id is not None + else {} + ) + resource = resource.merge(OTELResourceDetector().detect()) try: from opentelemetry.resourcedetector.gcp_resource_detector import GoogleCloudResourceDetector @@ -121,5 +151,4 @@ def get_gcp_resource() -> Resource: 'Cloud not import opentelemetry.resourcedetector.gcp_resource_detector' ' GCE, GKE or CloudRun related resource attributes may be missing' ) - return resource diff --git a/src/google/adk/telemetry/setup.py b/src/google/adk/telemetry/setup.py index fd059ce0..eac5c961 100644 --- a/src/google/adk/telemetry/setup.py +++ b/src/google/adk/telemetry/setup.py @@ -73,10 +73,8 @@ def maybe_set_otel_providers( otel_resource: OTel resource to use in providers. If empty - default OTel resource detection will be used. """ - if otel_hooks_to_setup is None: - otel_hooks_to_setup = [] - if otel_resource is None: - otel_resource = _get_otel_resource() + otel_hooks_to_setup = otel_hooks_to_setup or [] + otel_resource = otel_resource or _get_otel_resource() # Add generic OTel exporters based on OTel env variables. otel_hooks_to_setup.append(_get_otel_exporters()) diff --git a/tests/unittests/telemetry/test_google_cloud.py b/tests/unittests/telemetry/test_google_cloud.py index c861c070..318be630 100644 --- a/tests/unittests/telemetry/test_google_cloud.py +++ b/tests/unittests/telemetry/test_google_cloud.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +from typing import Optional from unittest import mock from google.adk.telemetry.google_cloud import get_gcp_exporters +from google.adk.telemetry.google_cloud import get_gcp_resource import pytest @@ -55,3 +58,34 @@ def test_get_gcp_exporters( assert len(otel_hooks.log_record_processors) == ( 1 if enable_cloud_logging else 0 ) + + +@pytest.mark.parametrize("project_id_in_arg", ["project_id_in_arg", None]) +@pytest.mark.parametrize("project_id_on_env", ["project_id_on_env", None]) +def test_get_gcp_resource( + project_id_in_arg: Optional[str], + project_id_on_env: Optional[str], + monkeypatch: pytest.MonkeyPatch, +): + # Arrange. + if project_id_on_env is not None: + monkeypatch.setenv( + "OTEL_RESOURCE_ATTRIBUTES", f"gcp.project_id={project_id_on_env}" + ) + + # Act. + otel_resource = get_gcp_resource(project_id_in_arg) + + # Assert. + expected_project_id = ( + project_id_on_env + if project_id_on_env is not None + else project_id_in_arg + if project_id_in_arg is not None + else None + ) + assert otel_resource is not None + assert ( + otel_resource.attributes.get("gcp.project_id", None) + == expected_project_id + )