From f1f18c8e1c6e6da7422ee09aac5776f714c2fa5a Mon Sep 17 00:00:00 2001
From: fdupont <fdupont@pop-os.localdomain>
Date: Tue, 6 Feb 2024 15:42:41 +0100
Subject: [PATCH] feat: enable opentelemetry grpc instrumentation

---
 setup.py                  |  4 ++++
 src/pythie_serving/run.py | 15 +++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/setup.py b/setup.py
index 8e0021d..c37b119 100644
--- a/setup.py
+++ b/setup.py
@@ -9,6 +9,10 @@
     "treelite_runtime~=2.2.2",
     "scikit-learn~=1.2.0",
     "cloudpickle~=2.1.0",
+    "opentelemetry-instrumentation-grpc~=0.38b0",
+    "opentelemetry-api>=1.17.0, <2.0",
+    "opentelemetry-sdk>=1.17.0, <2.0",
+    "opentelemetry-exporter-otlp>=1.17.0, <2.0",
 ]
 extras_require_test = [
     *extras_require_serving,
diff --git a/src/pythie_serving/run.py b/src/pythie_serving/run.py
index 26d5e11..7da6012 100644
--- a/src/pythie_serving/run.py
+++ b/src/pythie_serving/run.py
@@ -5,6 +5,11 @@
 from logging.config import dictConfig
 
 from google.protobuf import text_format
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.grpc import GrpcInstrumentorServer
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
 from pythie_serving import create_grpc_server
 from pythie_serving.tensorflow_proto.tensorflow_serving.config import (
@@ -12,7 +17,17 @@
 )
 
 
+def initialize_tracing():
+    tracing_collector_host = str(os.environ.get("TRACING_COLLECTOR_HOST"))
+    if tracing_collector_host is not None:
+        trace.set_tracer_provider(TracerProvider())
+        otlp_exporter = OTLPSpanExporter(endpoint=tracing_collector_host, insecure=True)
+        trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(otlp_exporter))
+        GrpcInstrumentorServer().instrument()
+
+
 def run():
+    initialize_tracing()
     model_choice_set = {"xgboost", "lightgbm", "treelite", "sklearn", "table"}
     model_choice_str = ",".join(model_choice_set)