More robust tracer setup. (#1445)

Author: nickolai-dr

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

@@ -126,15 +126,22 @@ def __init__(self, model_dir, target_type=None):
             # 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
+                msg = """Instrumentation for requests is not loaded, make sure appropriate
                 packages are installed:
 
                 pip install opentelemetry-instrumentation-requests
+                """
+            try:
+                from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
+
+                AioHttpClientInstrumentor().instrument()
+            except (ImportError, ModuleNotFoundError):
+                msg = """Instrumentation for aiottp is not loaded, make sure appropriate
+                packages are installed:
+
                 pip install opentelemetry-instrumentation-aiohttp-client
                 """
                 self._logger.warning(msg)

custom_model_runner/datarobot_drum/drum/common.py

@@ -135,37 +135,62 @@ def make_otel_endpoint(datarobot_endpoint):
     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")
+def setup_tracer(runtime_parameters, options):
+    """Setups OTEL tracer if not configured externally.
+
+    It is possible to provied OTEL compliant OTEL_EXPORTER_OTLP_ENDPOINT
+    or OTEL_EXPORTER_OTLP_TRACES_ENDPOINT and override trace collector
+    to anydd
+
+
+    Parameters
+    ----------
+    runtime_parameters: Type[RuntimeParameters] class handles runtime parameters for custom modes
+        used to check if OTEL configuration from user.
+    options: argparse.Namespace: object obtained from argparser filled with user supplied
+        command argumetns
+    Returns
+    -------
+    None
+    """
+    if os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") or os.environ.get(
+        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"
     ):
-        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)
+        # OTEL configured externaly via env vars
+        resource = Resource.create()
+    else:
+        gpu_predictor = getattr(options, "gpu_predictor", None)
+
+        # if explicitly asked enable/disable otel
+        if runtime_parameters.has("OTEL_SDK_ENABLED"):
+            enable_otel = runtime_parameters.get("OTEL_SDK_ENABLED")
+        # if not expliciety specified, enable otel by default for gpu models
+        elif gpu_predictor:
+            enable_otel = True
+        else:
+            enable_otel = False
+        # 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 enable_otel or not deployment_id:
+            return
+
+        resource = Resource.create(
+            {
+                "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)
+        h = f"X-DataRobot-Api-Key={key},X-DataRobot-Entity-Id=deployment-{deployment_id}"
+        os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint
+        os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = h
+
+    otlp_exporter = OTLPSpanExporter()
     provider = TracerProvider(resource=resource)
     provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
     trace.set_tracer_provider(provider)

custom_model_runner/datarobot_drum/drum/main.py

@@ -92,7 +92,8 @@ def signal_handler(sig, frame):
         _setup_required_environment_variables(options)
         # Env vars may setup OTEL configuration, lets setup
         # tracer after all env vars updated
-        setup_tracer(RuntimeParameters)
+
+        setup_tracer(RuntimeParameters, options)
         if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
             options.max_workers = RuntimeParameters.get("CUSTOM_MODEL_WORKERS")
         runtime.options = options

custom_model_runner/requirements.txt

@@ -28,3 +28,5 @@ datarobot>=3.1.0,<4
 opentelemetry-api
 opentelemetry-sdk
 opentelemetry-exporter-otlp-proto-http
+opentelemetry-instrumentation-aiohttp-client
+opentelemetry-instrumentation-requests