[MMM-19334] Add support for OTEL. (#1411)

* Add support for OTEL.

* Set OTEL_EXPORTER_OTLP_ENDPOINT

* Add extra header or DR permissions.

* Instrument moderation proactivly.

Author: nickolai-dr

custom_model_runner/datarobot_drum/drum/adapters/model_adapters/python_model_adapter.py

@@ -27,7 +27,6 @@
 from datarobot_drum.drum.artifact_predictors.xgboost_predictor import XGBoostPredictor
 from datarobot_drum.drum.artifact_predictors.onnx_predictor import ONNXPredictor
 from datarobot_drum.drum.root_predictors.chat_helpers import is_openai_model
-
 from datarobot_drum.drum.common import (
     reroute_stdout_to_stderr,
     SupportedPayloadFormats,
@@ -123,6 +122,22 @@ def __init__(self, model_dir, target_type=None):
                     "Unexpected empty target name for text generation, "
                     "vector database, or agentic workflow target."
                 )
+            # Instrument http clients in order to get nice traces from moderation library. We are
+            # doing this here because moderation library is loaded before custom.py.
+            try:
+                from opentelemetry.instrumentation.requests import RequestsInstrumentor
+                from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
+
+                RequestsInstrumentor().instrument()
+                AioHttpClientInstrumentor().instrument()
+            except (ImportError, ModuleNotFoundError):
+                msg = """Instrumentation for requests or aiottp is not loaded, make sure appropriate
+                packages are installed:
+
+                pip install opentelemetry-instrumentation-requests
+                pip install opentelemetry-instrumentation-aiohttp-client
+                """
+                self._logger.warning(msg)
             self._load_moderation_hooks()
         else:
             self._target_name = None

custom_model_runner/datarobot_drum/drum/common.py

@@ -9,6 +9,7 @@
 import sys
 from contextvars import ContextVar
 from distutils.util import strtobool
+from urllib.parse import urlparse, urlunparse
 
 from contextlib import contextmanager
 from pathlib import Path
@@ -20,6 +21,12 @@
     PayloadFormat,
 )
 from datarobot_drum.drum.exceptions import DrumCommonException
+from opentelemetry import trace, context
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 
 ctx_request_id = ContextVar("request_id")
@@ -119,3 +126,57 @@ def to_bool(value):
 
 
 FIT_METADATA_FILENAME = "fit_runtime_data.json"
+
+
+def make_otel_endpoint(datarobot_endpoint):
+    parsed_url = urlparse(datarobot_endpoint)
+    stripped_url = (parsed_url.scheme, parsed_url.netloc, "otel", "", "", "")
+    result = urlunparse(stripped_url)
+    return result
+
+
+def setup_tracer(runtime_parameters):
+    # OTEL disabled by default for now.
+    if not (
+        runtime_parameters.has("OTEL_SDK_ENABLED") and runtime_parameters.get("OTEL_SDK_ENABLED")
+    ):
+        return
+    # if deployment_id is not found, most likely this is custom model
+    # testing
+    deployment_id = os.environ.get("MLOPS_DEPLOYMENT_ID", os.environ.get("DEPLOYMENT_ID"))
+    if not deployment_id:
+        return
+
+    service_name = f"deployment-{deployment_id}"
+    resource = Resource.create(
+        {
+            "service.name": service_name,
+            "datarobot.deployment_id": deployment_id,
+        }
+    )
+    key = os.environ.get("DATAROBOT_API_TOKEN")
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    if not key or not datarobot_endpoint:
+        return
+    endpoint = make_otel_endpoint(datarobot_endpoint)
+
+    os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "X-DataRobot-Entity-Id": f"entity=deployment; id={deployment_id};",
+    }
+    otlp_exporter = OTLPSpanExporter(headers=headers)
+    provider = TracerProvider(resource=resource)
+    provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+    trace.set_tracer_provider(provider)
+
+
+@contextmanager
+def otel_context(tracer, span_name, carrier):
+    ctx = TraceContextTextMapPropagator().extract(carrier=carrier)
+    token = context.attach(ctx)
+    try:
+        with tracer.start_as_current_span(span_name) as span:
+            yield span
+    finally:
+        context.detach(token)

custom_model_runner/datarobot_drum/drum/main.py

@@ -43,7 +43,7 @@
 import sys
 
 from datarobot_drum.drum.args_parser import CMRunnerArgsRegistry
-from datarobot_drum.drum.common import config_logging
+from datarobot_drum.drum.common import config_logging, setup_tracer
 from datarobot_drum.drum.enum import RunMode
 from datarobot_drum.drum.enum import ExitCodes
 from datarobot_drum.drum.exceptions import DrumSchemaValidationException
@@ -90,6 +90,9 @@ def signal_handler(sig, frame):
         options = arg_parser.parse_args()
         CMRunnerArgsRegistry.verify_options(options)
         _setup_required_environment_variables(options)
+        # Env vars may setup OTEL configuration, lets setup
+        # tracer after all env vars updated
+        setup_tracer(RuntimeParameters)
         if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
             options.max_workers = RuntimeParameters.get("CUSTOM_MODEL_WORKERS")
         runtime.options = options

custom_model_runner/datarobot_drum/drum/root_predictors/prediction_server.py

@@ -12,6 +12,7 @@
 from flask import Response, jsonify, request
 from werkzeug.exceptions import HTTPException
 
+from opentelemetry import trace
 from datarobot_drum.drum.description import version as drum_version
 from datarobot_drum.drum.enum import (
     FLASK_EXT_FILE_NAME,
@@ -38,10 +39,14 @@
     get_flask_app,
 )
 from datarobot_drum.profiler.stats_collector import StatsCollector, StatsOperation
+from datarobot_drum.drum.common import otel_context
 
 logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + __name__)
 
 
+tracer = trace.get_tracer(__name__)
+
+
 class PredictionServer(PredictMixin):
     def __init__(self, params: dict):
         self._params = params
@@ -157,54 +162,50 @@ def health():
         @model_api.route("/invocations", methods=["POST"])
         def predict():
             logger.debug("Entering predict() endpoint")
-
-            self._pre_predict_and_transform()
-            try:
-                response, response_status = self.do_predict_structured(logger=logger)
-            finally:
-                self._post_predict_and_transform()
+            with otel_context(tracer, "drum.invocations", request.headers):
+                self._pre_predict_and_transform()
+                try:
+                    response, response_status = self.do_predict_structured(logger=logger)
+                finally:
+                    self._post_predict_and_transform()
 
             return response, response_status
 
         @model_api.route("/transform/", methods=["POST"])
         def transform():
             logger.debug("Entering transform() endpoint")
-
-            self._pre_predict_and_transform()
-
-            try:
-                response, response_status = self.do_transform(logger=logger)
-            finally:
-                self._post_predict_and_transform()
+            with otel_context(tracer, "drum.transform", request.headers):
+                self._pre_predict_and_transform()
+                try:
+                    response, response_status = self.do_transform(logger=logger)
+                finally:
+                    self._post_predict_and_transform()
 
             return response, response_status
 
         @model_api.route("/predictionsUnstructured/", methods=["POST"])
         @model_api.route("/predictUnstructured/", methods=["POST"])
         def predict_unstructured():
             logger.debug("Entering predict() endpoint")
-
-            self._pre_predict_and_transform()
-
-            try:
-                response, response_status = self.do_predict_unstructured(logger=logger)
-            finally:
-                self._post_predict_and_transform()
-
+            with otel_context(tracer, "drum.predictUnstructured", request.headers):
+                self._pre_predict_and_transform()
+                try:
+                    response, response_status = self.do_predict_unstructured(logger=logger)
+                finally:
+                    self._post_predict_and_transform()
             return (response, response_status)
 
         # Chat routes are defined without trailing slash because this is required by the OpenAI python client.
         @model_api.route("/chat/completions", methods=["POST"])
         @model_api.route("/v1/chat/completions", methods=["POST"])
         def chat():
             logger.debug("Entering chat endpoint")
-
-            self._pre_predict_and_transform()
-
-            try:
-                response, response_status = self.do_chat(logger=logger)
-            finally:
-                self._post_predict_and_transform()
+            with otel_context(tracer, "drum.chat.completions", request.headers):
+                self._pre_predict_and_transform()
+                try:
+                    response, response_status = self.do_chat(logger=logger)
+                finally:
+                    self._post_predict_and_transform()
 
             return response, response_status
 
@@ -226,24 +227,25 @@ def get_supported_llm_models():
         @model_api.route("/directAccess/<path:path>", methods=["GET", "POST", "PUT"])
         @model_api.route("/nim/<path:path>", methods=["GET", "POST", "PUT"])
         def forward_request(path):
-            if not hasattr(self._predictor, "openai_host") or not hasattr(
-                self._predictor, "openai_port"
-            ):
-                return {
-                    "message": "This endpoint is only supported by OpenAI based predictors"
-                }, HTTP_400_BAD_REQUEST
-
-            openai_host = self._predictor.openai_host
-            openai_port = self._predictor.openai_port
-
-            resp = requests.request(
-                method=request.method,
-                url=f"http://{openai_host}:{openai_port}/{path.rstrip('/')}",
-                headers=request.headers,
-                params=request.args,
-                data=request.get_data(),
-                allow_redirects=False,
-            )
+            with otel_context(tracer, "drum.directAccess", request.headers) as span:
+                if not hasattr(self._predictor, "openai_host") or not hasattr(
+                    self._predictor, "openai_port"
+                ):
+                    msg = "This endpoint is only supported by OpenAI based predictors"
+                    span.set_status(StatusCode.ERROR, msg)
+                    return {"message": msg}, HTTP_400_BAD_REQUEST
+
+                openai_host = self._predictor.openai_host
+                openai_port = self._predictor.openai_port
+
+                resp = requests.request(
+                    method=request.method,
+                    url=f"http://{openai_host}:{openai_port}/{path.rstrip('/')}",
+                    headers=request.headers,
+                    params=request.args,
+                    data=request.get_data(),
+                    allow_redirects=False,
+                )
 
             return Response(resp.content, status=resp.status_code, headers=dict(resp.headers))
 

custom_model_runner/requirements.txt

@@ -24,3 +24,7 @@ pydantic
 datarobot-storage
 datarobot-mlops>=10.2.0  # Required for the 'set_api_spooler' with arugments
 datarobot>=3.1.0,<4
+# otel
+opentelemetry-api
+opentelemetry-sdk
+opentelemetry-exporter-otlp-proto-http