From fdb681a37ef3b8e3c8958ed94616451c8b06e026 Mon Sep 17 00:00:00 2001 From: Aaron Abbott Date: Tue, 4 Mar 2025 21:09:57 +0000 Subject: [PATCH 1/2] Implement uninstrument for `opentelemetry-insturmentation-vertexai` --- .../CHANGELOG.md | 3 ++ .../instrumentation/vertexai/__init__.py | 43 ++++++++++++------- .../tests/test_instrumentor.py | 37 ++++++++++++++++ 3 files changed, 68 insertions(+), 15 deletions(-) create mode 100644 instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md b/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md index 1ac2d13a82..4432139a19 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +- Implement uninstrument for `opentelemetry-instrumentation-vertexai` + ([#3328](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3328)) + ## Version 2.0b0 (2025-02-24) - Added Vertex AI spans for request parameters diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py index 40d1cb48ac..0f8fcae03d 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/__init__.py @@ -47,6 +47,7 @@ from opentelemetry._events import get_event_logger from opentelemetry.instrumentation.instrumentor import BaseInstrumentor +from opentelemetry.instrumentation.utils import unwrap from opentelemetry.instrumentation.vertexai.package import _instruments from opentelemetry.instrumentation.vertexai.patch import ( generate_content_create, @@ -56,6 +57,23 @@ from opentelemetry.trace import get_tracer +def _client_classes(): + # This import is very slow, do it lazily in case instrument() is not called + + # pylint: disable=import-outside-toplevel + from google.cloud.aiplatform_v1.services.prediction_service import ( + client, + ) + from google.cloud.aiplatform_v1beta1.services.prediction_service import ( + client as client_v1beta1, + ) + + return ( + client.PredictionServiceClient, + client_v1beta1.PredictionServiceClient, + ) + + class VertexAIInstrumentor(BaseInstrumentor): def instrumentation_dependencies(self) -> Collection[str]: return _instruments @@ -77,20 +95,15 @@ def _instrument(self, **kwargs: Any): event_logger_provider=event_logger_provider, ) - wrap_function_wrapper( - module="google.cloud.aiplatform_v1beta1.services.prediction_service.client", - name="PredictionServiceClient.generate_content", - wrapper=generate_content_create( - tracer, event_logger, is_content_enabled() - ), - ) - wrap_function_wrapper( - module="google.cloud.aiplatform_v1.services.prediction_service.client", - name="PredictionServiceClient.generate_content", - wrapper=generate_content_create( - tracer, event_logger, is_content_enabled() - ), - ) + for client_class in _client_classes(): + wrap_function_wrapper( + client_class, + name="generate_content", + wrapper=generate_content_create( + tracer, event_logger, is_content_enabled() + ), + ) def _uninstrument(self, **kwargs: Any) -> None: - """TODO: implemented in later PR""" + for client_class in _client_classes(): + unwrap(client_class, "generate_content") diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py new file mode 100644 index 0000000000..5a110e0a8f --- /dev/null +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py @@ -0,0 +1,37 @@ +import pytest +from google.cloud.aiplatform_v1.services.prediction_service import client +from google.cloud.aiplatform_v1beta1.services.prediction_service import ( + client as client_v1beta1, +) + +from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor + + +@pytest.fixture(name="instrumentor") +def fixture_instrumentor(): + instrumentor = VertexAIInstrumentor() + instrumentor.instrument() + yield instrumentor + + if instrumentor.is_instrumented_by_opentelemetry: + instrumentor.uninstrument() + + +@pytest.fixture( + name="client_class", + params=[ + pytest.param(client.PredictionServiceClient, id="v1"), + pytest.param(client_v1beta1.PredictionServiceClient, id="v1beta1"), + ], +) +def fixture_client_class(request): + return request.param + + +def test_instruments(instrumentor: VertexAIInstrumentor, client_class): + assert hasattr(client_class.generate_content, "__wrapped__") + + +def test_uninstruments(instrumentor: VertexAIInstrumentor, client_class): + instrumentor.uninstrument() + assert not hasattr(client_class.generate_content, "__wrapped__") From 4c999544541202d461c19856cf74724caa1ed711 Mon Sep 17 00:00:00 2001 From: Aaron Abbott Date: Thu, 6 Mar 2025 15:31:58 +0000 Subject: [PATCH 2/2] Reuse fixture from conftest.py --- .../tests/conftest.py | 6 ++- .../tests/test_instrumentor.py | 38 ++++++++++++------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py index b76a108805..fb0956f0b5 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/conftest.py @@ -111,7 +111,8 @@ def instrument_no_content( yield instrumentor os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None) - instrumentor.uninstrument() + if instrumentor.is_instrumented_by_opentelemetry: + instrumentor.uninstrument() @pytest.fixture @@ -130,7 +131,8 @@ def instrument_with_content( yield instrumentor os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None) - instrumentor.uninstrument() + if instrumentor.is_instrumented_by_opentelemetry: + instrumentor.uninstrument() @pytest.fixture(scope="module") diff --git a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py index 5a110e0a8f..e701426d2e 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py +++ b/instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py @@ -1,3 +1,19 @@ +# Copyright The OpenTelemetry Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + import pytest from google.cloud.aiplatform_v1.services.prediction_service import client from google.cloud.aiplatform_v1beta1.services.prediction_service import ( @@ -7,16 +23,6 @@ from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor -@pytest.fixture(name="instrumentor") -def fixture_instrumentor(): - instrumentor = VertexAIInstrumentor() - instrumentor.instrument() - yield instrumentor - - if instrumentor.is_instrumented_by_opentelemetry: - instrumentor.uninstrument() - - @pytest.fixture( name="client_class", params=[ @@ -24,14 +30,18 @@ def fixture_instrumentor(): pytest.param(client_v1beta1.PredictionServiceClient, id="v1beta1"), ], ) -def fixture_client_class(request): +def fixture_client_class(request: pytest.FixtureRequest): return request.param -def test_instruments(instrumentor: VertexAIInstrumentor, client_class): +def test_instruments( + instrument_with_content: VertexAIInstrumentor, client_class +): assert hasattr(client_class.generate_content, "__wrapped__") -def test_uninstruments(instrumentor: VertexAIInstrumentor, client_class): - instrumentor.uninstrument() +def test_uninstruments( + instrument_with_content: VertexAIInstrumentor, client_class +): + instrument_with_content.uninstrument() assert not hasattr(client_class.generate_content, "__wrapped__")