From 359645655056f1087f1932e1d99c78201fc30990 Mon Sep 17 00:00:00 2001
From: Vincent Emonet <vincent.emonet@gmail.com>
Date: Mon, 18 Dec 2023 17:06:51 +0100
Subject: [PATCH] add opentelemetry

---
 pyproject.toml |  7 ++++++-
 src/api.py     | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 4920fc1..c19cc12 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,12 @@ dependencies = [
     "PubChemPy",
     "typer",
     # NOTE: bug in 1.6.1
-    "linkml-runtime ==1.6.0"
+    "linkml-runtime ==1.6.0",
+
+    "opentelemetry-sdk",
+    "opentelemetry-exporter-otlp-proto-http",
+    "opentelemetry-instrumentation-fastapi",
+    "opentelemetry-instrumentation-httpx",
 ]
 
 [project.optional-dependencies]
diff --git a/src/api.py b/src/api.py
index 54191a1..5d59112 100644
--- a/src/api.py
+++ b/src/api.py
@@ -1,4 +1,5 @@
 import logging
+import os
 
 from trapi_predict_kit import TRAPI, settings
 
@@ -91,3 +92,48 @@ def post_train(api_key: str):
     # return scores_df.to_dict(orient="records")
 
     return scores
+
+
+
+
+def configure_otel(app):
+    # open telemetry https://github.com/ranking-agent/aragorn/blob/main/src/otel_config.py#L4
+    # https://ncatstranslator.github.io/TranslatorTechnicalDocumentation/deployment-guide/monitoring/
+    # https://github.com/TranslatorSRI/Jaeger-demo
+    if not os.environ.get("NO_JAEGER"):
+        logging.info("Starting up jaeger telemetry")
+        import warnings
+        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+        from opentelemetry import trace
+        from opentelemetry.exporter.jaeger.thrift import JaegerExporter
+        from opentelemetry.sdk.resources import SERVICE_NAME as telemetery_service_name_key, Resource
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+        from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+
+        service_name = os.environ.get("OTEL_SERVICE_NAME", "OPENPREDICT")
+        # httpx connections need to be open a little longer by the otel decorators
+        # but some libs display warnings of resource being unclosed.
+        # these supresses such warnings.
+        logging.captureWarnings(capture=True)
+        warnings.filterwarnings("ignore",category=ResourceWarning)
+        trace.set_tracer_provider(
+            TracerProvider(
+                resource=Resource.create({telemetery_service_name_key: service_name})
+            )
+        )
+        jaeger_host = os.environ.get('JAEGER_HOST', 'jaeger-otel-agent.sri')
+        jaeger_port = int(os.environ.get('JAEGER_PORT', '6831'))
+        jaeger_exporter = JaegerExporter(
+            agent_host_name=jaeger_host,
+            agent_port=jaeger_port,
+        )
+        trace.get_tracer_provider().add_span_processor(
+            BatchSpanProcessor(jaeger_exporter)
+        )
+        # tracer = trace.get_tracer(__name__)
+        FastAPIInstrumentor.instrument_app(app, tracer_provider=trace, excluded_urls="docs,openapi.json")
+        HTTPXClientInstrumentor().instrument()
+
+# Configure open telemetry if enabled
+configure_otel(app)