feat: enable opentelemetry grpc instrumentation

README.md

@@ -52,6 +52,7 @@ The following models can be served by pythie-serving:
    * `--port`: Port number to listen to.
 
 #### Environment variables
+* `OPENTELEMETRY_COLLECTOR_HOST`: OpenTelemetry Collector receiver endpoint. If not defined OpenTelemtetry will not be activated. See https://opentelemetry.io/docs/what-is-opentelemetry for more details.
 For a treelite served model:
 * `TREELITE_NTHREAD`: Number of threads to use to compute predictions
 * `TREELINTE_BIND_THREADS`: Set to `0` to deactivate thread pinning. See https://treelite.readthedocs.io/en/latest/treelite-runtime-api.html

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,

src/pythie_serving/run.py

@@ -5,13 +5,27 @@
 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 (
     model_server_config_pb2,
 )
 
 
+def initialize_opentelemetry():
+    otel_collector_host = str(os.environ.get("OPENTELEMETRY_COLLECTOR_HOST"))
+    if otel_collector_host is not None:
+        trace.set_tracer_provider(TracerProvider())
+        otlp_exporter = OTLPSpanExporter(endpoint=otel_collector_host, insecure=True)
+        trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(otlp_exporter))
+        GrpcInstrumentorServer().instrument()
+
+
 def run():
     model_choice_set = {"xgboost", "lightgbm", "treelite", "sklearn", "table"}
     model_choice_str = ",".join(model_choice_set)
@@ -76,6 +90,8 @@ def run():
     with open(ns.model_config_file_path) as opened_config_file:
         text_format.Parse(opened_config_file.read(), model_server_config)
 
+    initialize_opentelemetry()
+
     maximum_concurrent_rpcs = ns.maximum_concurrent_rpcs
     if maximum_concurrent_rpcs < 0:
         maximum_concurrent_rpcs = None  # grpc.server takes None to accept unlimited amount of connections