Improve Everest logging with open telemetry traces

Author: HakonSohoel

src/everest/bin/main.py

@@ -16,6 +16,7 @@
 from everest.bin.monitor_script import monitor_entry
 from everest.bin.visualization_script import visualization_entry
 from everest.plugins.everest_plugin_manager import EverestPluginManager
+from everest.trace import tracer, tracer_provider
 
 
 def _build_args_parser() -> argparse.ArgumentParser:
@@ -48,7 +49,9 @@ def __init__(self, args: list[str]) -> None:
             parser.error("Unrecognized command")
 
         # Setup logging from plugins:
-        EverestPluginManager().add_log_handle_to_root()
+        plugin_manager = EverestPluginManager()
+        plugin_manager.add_log_handle_to_root()
+        plugin_manager.add_span_processor_to_trace_provider(tracer_provider)
         logger = logging.getLogger(__name__)
         logger.info(f"Started everest with {parsed_args}")
         # Use dispatch pattern to invoke method with same name
@@ -103,6 +106,7 @@ def results(self, args: list[str]) -> None:
         visualization_entry(args)
 
 
+@tracer.start_as_current_span("everest.application.start")
 def start_everest(args: list[str] | None = None) -> None:
     """Main entry point for the everest application"""
     args = args or sys.argv

src/everest/detached/__init__.py

@@ -32,6 +32,7 @@
     START_EXPERIMENT_ENDPOINT,
     STOP_ENDPOINT,
 )
+from everest.trace import get_traceparent
 
 # Specifies how many times to try a http request within the specified timeout.
 _HTTP_REQUEST_RETRY = 10
@@ -60,6 +61,8 @@ async def start_server(
             str(config.output_dir),
             "--logging-level",
             str(logging_level),
+            "--traceparent",
+            str(get_traceparent()),
         ]
         poll_task = asyncio.create_task(driver.poll(), name="poll_task")
         await driver.submit(0, "everserver", *args, name=Path(config.config_file).stem)

src/everest/detached/jobs/everserver.py

@@ -45,6 +45,7 @@
     HTTPBasic,
     HTTPBasicCredentials,
 )
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from pydantic import BaseModel
 
 from ert.config import QueueSystem
@@ -77,6 +78,7 @@
     START_EXPERIMENT_ENDPOINT,
     STOP_ENDPOINT,
 )
+from everest.trace import tracer, tracer_provider
 from everest.util import makedirs_if_needed, version_info
 
 
@@ -424,13 +426,40 @@ def make_handler_config(
     }
 
     logging.config.dictConfig(logging_config)
-    EverestPluginManager().add_log_handle_to_root()
+    plugin_manager = EverestPluginManager()
+    plugin_manager.add_log_handle_to_root()
+    plugin_manager.add_span_processor_to_trace_provider(tracer_provider)
+
+
+def get_trace_context():
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument(
+        "--traceparent",
+        type=str,
+        help="Trace parent id to be used by the storage root span",
+        default=None,
+    )
+    options = arg_parser.parse_args()
+    ctx = (
+        TraceContextTextMapPropagator().extract(
+            carrier={"traceparent": options.traceparent}
+        )
+        if options.traceparent
+        else None
+    )
+    return ctx
 
 
 def main() -> None:
     arg_parser = argparse.ArgumentParser()
     arg_parser.add_argument("--output-dir", "-o", type=str)
     arg_parser.add_argument("--logging-level", "-l", type=int, default=logging.INFO)
+    arg_parser.add_argument(
+        "--traceparent",
+        type=str,
+        help="Trace parent id to be used by the storage root span",
+        default=None,
+    )
     options = arg_parser.parse_args()
 
     output_dir = options.output_dir
@@ -441,77 +470,86 @@ def main() -> None:
     host_file = ServerConfig.get_hostfile_path(output_dir)
     msg_queue: SimpleQueue[EverestServerMsg] = SimpleQueue()
 
-    try:
-        _configure_loggers(
-            detached_dir=Path(ServerConfig.get_detached_node_dir(output_dir)),
-            log_dir=Path(output_dir) / OPTIMIZATION_LOG_DIR,
-            logging_level=logging_level,
+    ctx = (
+        TraceContextTextMapPropagator().extract(
+            carrier={"traceparent": options.traceparent}
         )
+        if options.traceparent
+        else None
+    )
 
-        update_everserver_status(status_path, ServerStatus.starting)
-        logging.getLogger(EVEREST).info(version_info())
-        logging.getLogger(EVEREST).info(f"Output directory: {output_dir}")
+    with tracer.start_as_current_span("everest.server", context=ctx):
+        try:
+            _configure_loggers(
+                detached_dir=Path(ServerConfig.get_detached_node_dir(output_dir)),
+                log_dir=Path(output_dir) / OPTIMIZATION_LOG_DIR,
+                logging_level=logging_level,
+            )
 
-        authentication = _generate_authentication()
-        cert_path, key_path, key_pw = _generate_certificate(
-            ServerConfig.get_certificate_dir(output_dir)
-        )
-        host = _get_machine_name()
-        port = _find_open_port(host, lower=5000, upper=5800)
-        _write_hostfile(host_file, host, port, cert_path, authentication)
-
-        shared_data = {
-            STOP_ENDPOINT: False,
-            "started": False,
-            "events": [],
-            "subscribers": {},
-        }
+            update_everserver_status(status_path, ServerStatus.starting)
+            logging.getLogger(EVEREST).info(version_info())
+            logging.getLogger(EVEREST).info(f"Output directory: {output_dir}")
 
-        server_config = {
-            "optimization_output_dir": optimization_output_dir,
-            "port": port,
-            "cert_path": cert_path,
-            "key_path": key_path,
-            "key_passwd": key_pw,
-            "authentication": authentication,
-        }
-        # Starting the server
-        everserver_instance = threading.Thread(
-            target=_everserver_thread,
-            args=(shared_data, server_config, msg_queue),
-        )
-        everserver_instance.daemon = True
-        everserver_instance.start()
+            authentication = _generate_authentication()
+            cert_path, key_path, key_pw = _generate_certificate(
+                ServerConfig.get_certificate_dir(output_dir)
+            )
+            host = _get_machine_name()
+            port = _find_open_port(host, lower=5000, upper=5800)
+            _write_hostfile(host_file, host, port, cert_path, authentication)
+
+            shared_data = {
+                STOP_ENDPOINT: False,
+                "started": False,
+                "events": [],
+                "subscribers": {},
+            }
+
+            server_config = {
+                "optimization_output_dir": optimization_output_dir,
+                "port": port,
+                "cert_path": cert_path,
+                "key_path": key_path,
+                "key_passwd": key_pw,
+                "authentication": authentication,
+            }
+            # Starting the server
+            everserver_instance = threading.Thread(
+                target=_everserver_thread,
+                args=(shared_data, server_config, msg_queue),
+            )
+            everserver_instance.daemon = True
+            everserver_instance.start()
 
-        # Monitoring the server
-        while True:
-            try:
-                item = msg_queue.get(timeout=1)  # Wait for data
-                match item:
-                    case ServerStarted():
-                        update_everserver_status(status_path, ServerStatus.running)
-                    case ServerStopped():
-                        update_everserver_status(status_path, ServerStatus.stopped)
-                        return
-                    case ExperimentFailed():
-                        update_everserver_status(
-                            status_path, ServerStatus.failed, item.msg
-                        )
-                        return
-                    case ExperimentComplete():
-                        status, message = _get_optimization_status(
-                            item.exit_code, item.data
-                        )
-                        update_everserver_status(status_path, status, message)
-                        return
-            except Empty:
-                continue
-    except:
-        update_everserver_status(
-            status_path,
-            ServerStatus.failed,
-            message=traceback.format_exc(),
-        )
+            # Monitoring the server
+            while True:
+                try:
+                    item = msg_queue.get(timeout=1)  # Wait for data
+                    match item:
+                        case ServerStarted():
+                            update_everserver_status(status_path, ServerStatus.running)
+                        case ServerStopped():
+                            update_everserver_status(status_path, ServerStatus.stopped)
+                            return
+                        case ExperimentFailed():
+                            update_everserver_status(
+                                status_path, ServerStatus.failed, item.msg
+                            )
+                            return
+                        case ExperimentComplete():
+                            status, message = _get_optimization_status(
+                                item.exit_code, item.data
+                            )
+                            update_everserver_status(status_path, status, message)
+                            return
+                except Empty:
+                    continue
+        except:
+            update_everserver_status(
+                status_path,
+                ServerStatus.failed,
+                message=traceback.format_exc(),
+            )
 
 
 def _get_optimization_status(

src/everest/plugins/everest_plugin_manager.py

@@ -2,6 +2,7 @@
 from typing import Any
 
 import pluggy
+from opentelemetry.sdk.trace import TracerProvider
 
 from everest.plugins import hook_impl, hook_specs
 from everest.strings import EVEREST
@@ -26,3 +27,10 @@ def add_log_handle_to_root(self) -> None:
         root_logger = logging.getLogger()
         for handler in self.hook.add_log_handle_to_root():
             root_logger.addHandler(handler)
+
+    def add_span_processor_to_trace_provider(
+        self, trace_provider: TracerProvider
+    ) -> None:
+        span_processors = self.hook.add_span_processor()
+        for span_processor in span_processors:
+            trace_provider.add_span_processor(span_processor)

src/everest/plugins/hook_specs.py

@@ -2,6 +2,7 @@
 from collections.abc import Sequence
 from typing import Any
 
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from pydantic import BaseModel
 
 from everest.plugins import hookspec
@@ -102,6 +103,16 @@ def add_log_handle_to_root() -> logging.Handler:  # type: ignore[empty-body]
     """
 
 
+@hookspec
+def add_span_processor() -> BatchSpanProcessor:  # type: ignore
+    """
+    Create a BatchSpanProcessor which will be added to the trace provider
+    in ert.
+
+    :return: A BatchSpanProcessor that will be added to the trace provider in everest
+    """
+
+
 @hookspec
 def get_forward_model_documentations() -> dict[str, Any]:  # type: ignore[empty-body]
     """ """

src/everest/trace.py

@@ -0,0 +1,22 @@
+from opentelemetry import trace
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import SpanLimits, TracerProvider
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+resource = Resource(attributes={SERVICE_NAME: "everest"})
+tracer_provider = TracerProvider(
+    resource=resource, span_limits=SpanLimits(max_events=128 * 16)
+)
+
+tracer = trace.get_tracer("everest.main", tracer_provider=tracer_provider)
+
+
+def get_trace_id() -> str:
+    return trace.format_trace_id(trace.get_current_span().get_span_context().trace_id)
+
+
+def get_traceparent() -> str | None:
+    carrier: dict[str, str] = {}
+    # Write the current context into the carrier.
+    TraceContextTextMapPropagator().inject(carrier)
+    return carrier.get("traceparent")