feat: Add enable_telemetry to the initializer

PiperOrigin-RevId: 703574132

Author: yeesian

google/cloud/aiplatform/initializer.py

@@ -151,6 +151,7 @@ def init(
         api_key: Optional[str] = None,
         api_transport: Optional[str] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = None,
+        enable_telemetry: Optional[bool] = False,
     ):
         """Updates common initialization parameters with provided options.
 
@@ -218,6 +219,9 @@ def init(
                 beta state (preview).
             request_metadata:
                 Optional. Additional gRPC metadata to send with every client request.
+            enable_telemetry:
+                Optional. Whether to enable telemetry. If True, traces will be
+                sent to Cloud Trace (under the given project). Default is False.
         Raises:
             ValueError:
                 If experiment_description is provided but experiment is not.
@@ -298,6 +302,30 @@ def init(
                 backing_tensorboard=experiment_tensorboard,
             )
 
+        if enable_telemetry:
+            logging.info("Telemetry is enabled.")
+            import google.cloud.trace_v2
+            from opentelemetry import trace
+            from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
+            from opentelemetry.sdk.trace import TracerProvider
+            from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+            credentials, _ = google.auth.default()
+            span_exporter = CloudTraceSpanExporter(
+                project_id=self.project,
+                client=trace_v2.TraceServiceClient(
+                    credentials=credentials.with_quota_project(self.project),
+                ),
+            )
+            span_processor = SimpleSpanProcessor(span_exporter=span_exporter)
+            tracer_provider = trace.get_tracer_provider()
+            trace.set_tracer_provider(tracer_provider)
+            cloud_trace_exporter = CloudTraceSpanExporter(project_id=self.project)
+            trace.get_tracer_provider().add_span_processor(
+                SimpleSpanProcessor(cloud_trace_exporter)
+            )
+            telemetry.VertexAIInstrumentor().instrument()
+
     def get_encryption_spec(
         self,
         encryption_spec_key_name: Optional[str],