Implement uninstrument for opentelemetry-insturmentation-vertexai

Author: aabmass

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-insturmentation-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

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,
@@ -55,6 +56,11 @@
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
+_MODULES = (
+    "google.cloud.aiplatform_v1beta1.services.prediction_service.client",
+    "google.cloud.aiplatform_v1.services.prediction_service.client",
+)
+
 
 class VertexAIInstrumentor(BaseInstrumentor):
     def instrumentation_dependencies(self) -> Collection[str]:
@@ -77,20 +83,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 module in _MODULES:
+            wrap_function_wrapper(
+                module=module,
+                name="PredictionServiceClient.generate_content",
+                wrapper=generate_content_create(
+                    tracer, event_logger, is_content_enabled()
+                ),
+            )
 
     def _uninstrument(self, **kwargs: Any) -> None:
-        """TODO: implemented in later PR"""
+        for module in _MODULES:
+            unwrap(f"{module}.PredictionServiceClient", "generate_content")

instrumentation-genai/opentelemetry-instrumentation-vertexai/tests/test_instrumentor.py

@@ -0,0 +1,49 @@
+import pytest
+from typing import Type
+
+from google.cloud.aiplatform_v1.services.prediction_service import client
+from google.cloud.aiplatform_v1.types import (
+    content,
+    prediction_service,
+)
+from google.cloud.aiplatform_v1beta1.services.prediction_service import (
+    client as client_v1beta1,
+)
+from google.cloud.aiplatform_v1beta1.types import (
+    content as content_v1beta1,
+)
+from google.cloud.aiplatform_v1beta1.types import (
+    prediction_service as prediction_service_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__")