Skip to content

feat: Add Resource to tracer provider in Agent Engine templates #5357

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add Resource to tracer provider in Agent Engine templates #5357

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,20 +146,18 @@

reasoning_engine_extra_require = [
"cloudpickle >= 3.0, < 4.0",
"google-cloud-trace < 2",
"opentelemetry-sdk < 2",
"opentelemetry-exporter-gcp-trace < 2",
"opentelemetry-exporter-otlp < 2",
"pydantic >= 2.11.1, < 3",
"typing_extensions",
]

agent_engines_extra_require = [
"packaging >= 24.0",
"cloudpickle >= 3.0, < 4.0",
"google-cloud-trace < 2",
"google-cloud-logging < 4",
"opentelemetry-sdk < 2",
"opentelemetry-exporter-gcp-trace < 2",
"opentelemetry-exporter-otlp < 2",
"pydantic >= 2.11.1, < 3",
"typing_extensions",
]
Expand Down
28 changes: 28 additions & 0 deletions vertexai/agent_engines/_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -685,6 +685,20 @@ def _import_opentelemetry_or_warn() -> Optional[types.ModuleType]:
return None


def _import_opentelemetry_sdk_resources_or_warn() -> Optional[types.ModuleType]:
"""Tries to import the opentelemetry.sdk.trace module."""
try:
import opentelemetry.sdk.resources # noqa:F401

return opentelemetry.sdk.resources
except ImportError:
LOGGER.warning(
"Failed to import opentelemetry.sdk.resources. Please call "
"'pip install google-cloud-aiplatform[agent_engines]'."
)
return None


def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]:
"""Tries to import the opentelemetry.sdk.trace module."""
try:
Expand Down Expand Up @@ -713,6 +727,20 @@ def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]:
return None


def _import_otlp_trace_exporter_or_warn() -> Optional[types.ModuleType]:
"""Tries to import the google.cloud.trace_v2 module."""
try:
import opentelemetry.exporter.otlp.proto.http.trace_exporter

return opentelemetry.exporter.otlp.proto.http.trace_exporter
except ImportError:
LOGGER.warning(
"opentelemetry-exporter-otlp is not installed. Please "
"call 'pip install google-cloud-aiplatform[agent_engines]'."
)
return None


def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]:
"""Tries to import the opentelemetry.exporter.cloud_trace module."""
try:
Expand Down
38 changes: 26 additions & 12 deletions vertexai/agent_engines/templates/ag2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,34 +91,44 @@ def _default_runnable_builder(
def _default_instrumentor_builder(project_id: str):
from vertexai.agent_engines import _utils

cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
openinference_autogen = _utils._import_openinference_autogen_or_warn()
opentelemetry = _utils._import_opentelemetry_or_warn()
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
if all(
(
cloud_trace_exporter,
cloud_trace_v2,
otlp_trace_exporter,
openinference_autogen,
opentelemetry,
opentelemetry_sdk_trace,
opentelemetry_sdk_resources,
)
):
import google.auth
import os

SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")

credentials, _ = google.auth.default()
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
project_id=project_id,
client=cloud_trace_v2.TraceServiceClient(
credentials=credentials.with_quota_project(project_id),
),
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(credentials),
endpoint="https://telemetry.googleapis.com/v1/traces",
)
span_processor: SpanProcessor = (
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
span_exporter=span_exporter,
)
)
resource = opentelemetry_sdk_trace.Resource.create(
attributes={
"gcp.project_id": project_id,
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
}
)
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
# Get the appropriate tracer provider:
# 1. If _TRACER_PROVIDER is already set, use that.
Expand All @@ -127,7 +137,7 @@ def _default_instrumentor_builder(project_id: str):
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
# If none of the above is set, we log a warning, and
# create a tracer provider.
if not tracer_provider:
if AGENT_ENGINE_ID or not tracer_provider:
from google.cloud.aiplatform import base

_LOGGER = base.Logger(__name__)
Expand All @@ -137,13 +147,17 @@ def _default_instrumentor_builder(project_id: str):
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
"or _PROXY_TRACER_PROVIDER."
)
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids AttributeError:
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
# attribute 'add_span_processor'.
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids OpenTelemetry client already exists error.
_override_active_span_processor(
Expand Down
38 changes: 26 additions & 12 deletions vertexai/agent_engines/templates/langchain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,34 +170,44 @@ def _default_runnable_builder(
def _default_instrumentor_builder(project_id: str):
from vertexai.agent_engines import _utils

cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
openinference_langchain = _utils._import_openinference_langchain_or_warn()
opentelemetry = _utils._import_opentelemetry_or_warn()
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
if all(
(
cloud_trace_exporter,
cloud_trace_v2,
otlp_trace_exporter,
openinference_langchain,
opentelemetry,
opentelemetry_sdk_trace,
opentelemetry_sdk_resources,
)
):
import google.auth
import os

SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")

credentials, _ = google.auth.default()
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
project_id=project_id,
client=cloud_trace_v2.TraceServiceClient(
credentials=credentials.with_quota_project(project_id),
),
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(credentials),
endpoint="https://telemetry.googleapis.com/v1/traces",
)
span_processor: SpanProcessor = (
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
span_exporter=span_exporter,
)
)
resource = opentelemetry_sdk_trace.Resource.create(
attributes={
"gcp.project_id": project_id,
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
}
)
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
# Get the appropriate tracer provider:
# 1. If _TRACER_PROVIDER is already set, use that.
Expand All @@ -206,7 +216,7 @@ def _default_instrumentor_builder(project_id: str):
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
# If none of the above is set, we log a warning, and
# create a tracer provider.
if not tracer_provider:
if AGENT_ENGINE_ID or not tracer_provider:
from google.cloud.aiplatform import base

_LOGGER = base.Logger(__name__)
Expand All @@ -216,13 +226,17 @@ def _default_instrumentor_builder(project_id: str):
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
"or _PROXY_TRACER_PROVIDER."
)
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids AttributeError:
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
# attribute 'add_span_processor'.
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids OpenTelemetry client already exists error.
_override_active_span_processor(
Expand Down
38 changes: 26 additions & 12 deletions vertexai/agent_engines/templates/langgraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,34 +161,44 @@ def _default_runnable_builder(
def _default_instrumentor_builder(project_id: str):
from vertexai.agent_engines import _utils

cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
openinference_langchain = _utils._import_openinference_langchain_or_warn()
opentelemetry = _utils._import_opentelemetry_or_warn()
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
if all(
(
cloud_trace_exporter,
cloud_trace_v2,
otlp_trace_exporter,
openinference_langchain,
opentelemetry,
opentelemetry_sdk_trace,
opentelemetry_sdk_resources,
)
):
import google.auth
import os

SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")

credentials, _ = google.auth.default()
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
project_id=project_id,
client=cloud_trace_v2.TraceServiceClient(
credentials=credentials.with_quota_project(project_id),
),
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(credentials),
endpoint="https://telemetry.googleapis.com/v1/traces",
)
span_processor: SpanProcessor = (
opentelemetry_sdk_trace.export.SimpleSpanProcessor(
span_exporter=span_exporter,
)
)
resource = opentelemetry_sdk_trace.Resource.create(
attributes={
"gcp.project_id": project_id,
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
}
)
tracer_provider: TracerProvider = opentelemetry.trace.get_tracer_provider()
# Get the appropriate tracer provider:
# 1. If _TRACER_PROVIDER is already set, use that.
Expand All @@ -197,7 +207,7 @@ def _default_instrumentor_builder(project_id: str):
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
# If none of the above is set, we log a warning, and
# create a tracer provider.
if not tracer_provider:
if AGENT_ENGINE_ID or not tracer_provider:
from google.cloud.aiplatform import base

base.Logger(__name__).warning(
Expand All @@ -206,13 +216,17 @@ def _default_instrumentor_builder(project_id: str):
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
"or _PROXY_TRACER_PROVIDER."
)
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids AttributeError:
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
# attribute 'add_span_processor'.
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids OpenTelemetry client already exists error.
_override_active_span_processor(
Expand Down
38 changes: 26 additions & 12 deletions vertexai/preview/reasoning_engines/templates/adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,30 +172,40 @@ def dump(self) -> Dict[str, Any]:
def _default_instrumentor_builder(project_id: str):
from vertexai.agent_engines import _utils

cloud_trace_exporter = _utils._import_cloud_trace_exporter_or_warn()
cloud_trace_v2 = _utils._import_cloud_trace_v2_or_warn()
otlp_trace_exporter = _utils._import_otlp_trace_exporter_or_warn()
opentelemetry = _utils._import_opentelemetry_or_warn()
opentelemetry_sdk_trace = _utils._import_opentelemetry_sdk_trace_or_warn()
opentelemetry_sdk_resources = _utils._import_opentelemetry_sdk_resources_or_warn()
if all(
(
cloud_trace_exporter,
cloud_trace_v2,
otlp_trace_exporter,
opentelemetry,
opentelemetry_sdk_trace,
opentelemetry_sdk_resources,
)
):
import google.auth
import os

SERVICE_INSTANCE_ID = opentelemetry_sdk_resources.SERVICE_INSTANCE_ID
SERVICE_NAME = opentelemetry_sdk_resources.SERVICE_NAME
AGENT_ENGINE_ID = os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID", "")

credentials, _ = google.auth.default()
span_exporter = cloud_trace_exporter.CloudTraceSpanExporter(
project_id=project_id,
client=cloud_trace_v2.TraceServiceClient(
credentials=credentials.with_quota_project(project_id),
),
span_exporter = otlp_trace_exporter.OTLPSpanExporter(
session=google.auth.transport.requests.AuthorizedSession(credentials),
endpoint="https://telemetry.googleapis.com/v1/traces",
)
span_processor = opentelemetry_sdk_trace.export.BatchSpanProcessor(
span_exporter=span_exporter,
)
resource = opentelemetry_sdk_trace.Resource.create(
attributes={
"gcp.project_id": project_id,
SERVICE_NAME: "aiplatform.googleapis.com/ReasoningEngine",
SERVICE_INSTANCE_ID: AGENT_ENGINE_ID,
}
)
tracer_provider = opentelemetry.trace.get_tracer_provider()
# Get the appropriate tracer provider:
# 1. If _TRACER_PROVIDER is already set, use that.
Expand All @@ -204,7 +214,7 @@ def _default_instrumentor_builder(project_id: str):
# 3. As a final fallback, use _PROXY_TRACER_PROVIDER.
# If none of the above is set, we log a warning, and
# create a tracer provider.
if not tracer_provider:
if AGENT_ENGINE_ID or not tracer_provider:
from google.cloud.aiplatform import base

_LOGGER = base.Logger(__name__)
Expand All @@ -214,13 +224,17 @@ def _default_instrumentor_builder(project_id: str):
"OTEL_PYTHON_TRACER_PROVIDER, _TRACER_PROVIDER, "
"or _PROXY_TRACER_PROVIDER."
)
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids AttributeError:
# 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no
# attribute 'add_span_processor'.
if _utils.is_noop_or_proxy_tracer_provider(tracer_provider):
tracer_provider = opentelemetry_sdk_trace.TracerProvider()
tracer_provider = opentelemetry_sdk_trace.TracerProvider(
resource=resource,
)
opentelemetry.trace.set_tracer_provider(tracer_provider)
# Avoids OpenTelemetry client already exists error.
_override_active_span_processor(
Expand Down
Loading