Implement uninstrument for opentelemetry-insturmentation-vertexai

Author: aabmass

Diff for: 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 for: 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 for: 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__")