Fix fokring and instrument OpenAI (#1498)

Author: nickolai-dr

custom_model_runner/datarobot_drum/drum/common.py

@@ -26,6 +26,7 @@
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 
@@ -135,7 +136,7 @@ def make_otel_endpoint(datarobot_endpoint):
     return result
 
 
-def setup_tracer(runtime_parameters):
+def setup_tracer(runtime_parameters, options):
     """Setups OTEL tracer.
 
     OTEL is configured by OTEL_EXPORTER_OTLP_ENDPOINT and
@@ -145,6 +146,8 @@ def setup_tracer(runtime_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
@@ -166,7 +169,19 @@ def setup_tracer(runtime_parameters):
     resource = Resource.create()
     otlp_exporter = OTLPSpanExporter()
     provider = TracerProvider(resource=resource)
-    provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+
+    # In case of NIM flask server is configured to run in multiprocessing
+    # mode that uses fork. Since BatchSpanProcessor start background thread
+    # with bunch of locks, OTEL simply deadlocks and does not offlooad any
+    # spans. Even if we start BatchSpanProcessor per fork, batches often
+    # missing due to process exits before all data offloaded. In forking
+    # case we use SimpleSpanProcessor (mostly NIMs) otherwise BatchSpanProcessor
+    # (most frequent case)
+    if options.max_workers > 1:
+        provider.add_span_processor(SimpleSpanProcessor(otlp_exporter))
+    else:
+        provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+
     trace.set_tracer_provider(provider)
 
     endpoint = main_endpoint or trace_endpoint

custom_model_runner/datarobot_drum/drum/gpu_predictors/base.py

@@ -48,13 +48,31 @@
 from datarobot_drum.drum.server import HTTP_513_DRUM_PIPELINE_ERROR
 from datarobot_drum.drum.root_predictors.chat_helpers import is_openai_model
 
+
+logger = logging.getLogger(__name__)
+
 # OpenAI client isn't a required dependency for DRUM, so we need to check if it's available
 try:
     from openai import OpenAI
     from openai.resources.chat.completions import Completions
 
     COMPLETIONS_CREATE_SIGNATURE = inspect.signature(Completions.create)
     _HAS_OPENAI = True
+
+    try:
+        # disable "anonymous" tracking from traceloop
+        os.environ["TRACELOOP_TRACE_CONTENT"] = "false"
+        from opentelemetry.instrumentation.openai import OpenAIInstrumentor
+
+        OpenAIInstrumentor().instrument()
+    except (ImportError, ModuleNotFoundError):
+        msg = """Instrumentation for openai is not loaded, make sure appropriate
+        packages are installed:
+
+        pip install opentelemetry-instrumentation-openai
+        """
+        logger.warning(msg)
+
 except ImportError:
     _HAS_OPENAI = False
 

custom_model_runner/datarobot_drum/drum/main.py

@@ -91,12 +91,11 @@ def signal_handler(sig, frame):
         CMRunnerArgsRegistry.verify_options(options)
         _setup_required_environment_variables(options)
 
-        setup_tracer(RuntimeParameters)
-
         if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
             options.max_workers = RuntimeParameters.get("CUSTOM_MODEL_WORKERS")
         runtime.options = options
 
+        setup_tracer(RuntimeParameters, options)
         signal.signal(signal.SIGINT, signal_handler)
         signal.signal(signal.SIGTERM, signal_handler)
 

tests/unit/datarobot_drum/drum/test_main.py

@@ -35,7 +35,7 @@ def test_custom_model_workers(
         runtime_params.has.return_value = False
 
     main()
-    runtime_params.has.assert_called_with("CUSTOM_MODEL_WORKERS")
+    runtime_params.has.assert_any_call("CUSTOM_MODEL_WORKERS")
     if workers_param:
-        runtime_params.get.assert_called_with("CUSTOM_MODEL_WORKERS")
+        runtime_params.get.assert_any_call("CUSTOM_MODEL_WORKERS")
     assert expected_workers == options.max_workers