Is it possible to get open telemetry tracing in python to record spans that happen within a ThreadPoolExecutor?

Question

I have open telemetry running with a python FastApi application. Traces are being sent to Jaeger and I can view them.

I have a bunch of IO intensive work being done, so I'm doing it in parallel with a ThreadPoolExecutor. Spans created in the functions executed by ThreadPoolExecutor are not making their way into Jaeger.

Can anyone point me in the right direction to get this working? At the moment I'm resorting to disabling the concurrency to record traces for performance debugging, but that won't work in production.

Answer 1

I've had a bit more time to dig into this and created this class to overcome the issue:

from opentelemetry import context as otel_context

class TracedThreadPoolExecutor(ThreadPoolExecutor):
    """Implementation of :class:`ThreadPoolExecutor` that will pass context into sub tasks."""

    def __init__(self, tracer: Tracer, *args, **kwargs):
        self.tracer = tracer
        super().__init__(*args, **kwargs)

    def with_otel_context(self, context: otel_context.Context, fn: Callable):
        otel_context.attach(context)
        return fn()

    def submit(self, fn, *args, **kwargs):
        """Submit a new task to the thread pool."""

        # get the current otel context
        context = otel_context.get_current()
        if context:
            return super().submit(
                lambda: self.with_otel_context(
                    context, lambda: fn(*args, **kwargs)
                ),
            )
        else:
            return super().submit(lambda: fn(*args, **kwargs))

This class can be used as follows:

from opentelemetry import trace
import multiprocessing

tracer = trace.get_tracer(__name__)
executor = TracedThreadPoolExecutor(tracer, max_workers=multiprocessing.cpu_count())

# from here you can use it as you would a regular ThreadPoolExecutor

result = executor.submit(some_work)
executor.shutdown(wait=True)
# do something with the result