Gracefull shutdown for OTEL tracer. (#1506)

Author: nickolai-dr

custom_model_runner/datarobot_drum/drum/common.py

@@ -150,7 +150,7 @@ def setup_tracer(runtime_parameters, options):
         command argumetns
     Returns
     -------
-    None
+    TracerProvider
     """
     log = get_drum_logger("setup_tracer")
 
@@ -186,6 +186,7 @@ def setup_tracer(runtime_parameters, options):
 
     endpoint = main_endpoint or trace_endpoint
     log.info(f"Tracing is configured with endpoint: {endpoint}")
+    return provider
 
 
 @contextmanager

custom_model_runner/datarobot_drum/drum/main.py

@@ -68,6 +68,9 @@ def signal_handler(sig, frame):
             if runtime.options and RunMode(runtime.options.subparser_name) == RunMode.SERVER:
                 if runtime.cm_runner:
                     runtime.cm_runner.terminate()
+            # Let traceer offload accumulated spans before shutdown.
+            if runtime.trace_provider is not None:
+                runtime.trace_provider.shutdown()
 
             os._exit(130)
 
@@ -95,7 +98,8 @@ def signal_handler(sig, frame):
             options.max_workers = RuntimeParameters.get("CUSTOM_MODEL_WORKERS")
         runtime.options = options
 
-        setup_tracer(RuntimeParameters, options)
+        runtime.trace_provider = setup_tracer(RuntimeParameters, options)
+
         signal.signal(signal.SIGINT, signal_handler)
         signal.signal(signal.SIGTERM, signal_handler)
 

custom_model_runner/datarobot_drum/drum/runtime.py

@@ -30,6 +30,7 @@ def __init__(self):
         self.initialization_succeeded = False
         self.options = None
         self.cm_runner = None
+        self.trace_provider = None
 
     def __enter__(self):
         return self