Is there any plan for OpenTelemetry instrumentation?

[febus982]

Having tracing, metrics and logs supported for OpenTelemetry would be a great addition.

[s3rius]

Hello. I think it would be great to have a taskiq instrumentation, but as I understand taskiq already has a simple integration with otlp out of the box. Since we use mainstream libraries for brokers, you can enable instrumentation for those libraries to get all the traces.

[nickderobertis]

We are currently configuring tracing throughout our application, instrumenting the libraries powering brokers helps but the main thing missing here from our perspective is root spans that lasts the lifecycle of a task. Theoretically it should be possible to have the producer of the task start a trace and it be connected to the worker handling the task in a single trace. It seems like the instrumentation for Celery accomplishes this.

Beyond cross-service tracing, not having a root span at the framework level seems to be causing issues for us. We are using FastAPI as our API framework and the instrumentation for that sets a root span for every request automatically. Without having this at the framework level, traces end up without root spans and so we have some traces lasting close to an hour with TaskIQ. Maybe there is a way to fix this on the otel/Tempo side (we are not experts with otel) but everything works as expected for our FastAPI application.

As a workaround for now we've added spans to task workers by wrapping broker.task in our own custom decorator. This seems to allow us to trace our task execution in the worker accurately but we still end up with traces without root spans containing many unrelated spans over a long duration.

[s3rius]

Okay, I got the point. We will look into it.

[ipfans]

It would be better if logfire could be easily integrated. Thank you guys

[febus982]

Logfire is an OpenTelemetry implementation. Instrumenting OpenTelemetry would automatically work on logfire.

[karansharmapaypay]

Hi, any updates here? I was considering using taskiq for our usecase (currently celery+fastapi). This seems to be a deal breaker for us since we rely heavily on end to end traces for a lot of our monitoring.

[febus982]

I ended up using another package :/

[s3rius]

There is a plan, but right now I don't have much resources to spare on that. I would be really glad if you could help with that. Right now we're in the middle of preparing some big updates.

[ipfans]

@s3rius Could you share how to implement it? I would like to share the idea if possible, and the code can be implemented by the community.

[ktnlvr]

Hey! What’s the current roadblock with integrating OpenTelemetry reporting? From my brief look, it seems that simply supplying the trace ID to an executing task would be enough.

Regarding the issue brought up by @nickderobertis: from what I know, OpenTelemetry does not allow suspending spans. So having one span last several hours is actually expected. You can always scope down to a more specific subspan when analyzing. Also, if spans are implemented, the loggers should automatically pick up on that. That way, we get two features for the price of one.

I chose TaskIQ over Celery because it allows for easy async. However, not having telemetry in a distributed system could be a dealbreaker. I’ll need it for a personal project anyway, so I can start working on the basic instrumentation design if that would be helpful.

[ZhymabekRoman]

Hi guys,
I wrote my own solution to add tracing support to taskiq tasks:

logfire + opentelemetry

import asyncio
import logging
from taskiq import InMemoryBroker, TaskiqMiddleware
from taskiq.message import TaskiqMessage
from taskiq.result import TaskiqResult

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.context import attach, detach
from opentelemetry.trace import Status, StatusCode

from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, ConsoleLogExporter
from opentelemetry._logs import set_logger_provider
import logfire

logfire.configure(token="TOKEN")


class OtelMiddleware(TaskiqMiddleware):
    def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
        self.tracer = trace.get_tracer(__name__)
        self.span = self.tracer.start_span(name=message.task_name)
        self.span.set_attribute("taskiq.task_id", message.task_id)
        self.span.set_attribute("taskiq.task_name", message.task_name)
        self.token = attach(trace.set_span_in_context(self.span))
        return message

    def post_execute(self, message: TaskiqMessage, result: TaskiqResult) -> None:
        self.span.set_status(Status(StatusCode.OK))
        self.span.end()
        detach(self.token)

    def on_error(self, message: TaskiqMessage, exception: BaseException) -> None:
        self.span.set_status(Status(StatusCode.ERROR, description=str(exception)))
        self.span.record_exception(exception)
        self.span.end()
        detach(self.token)


provider = TracerProvider()
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)

logger_provider = LoggerProvider()
set_logger_provider(logger_provider)
logger_provider.add_log_record_processor(BatchLogRecordProcessor(ConsoleLogExporter()))
otel_handler = LoggingHandler()
logging.getLogger().addHandler(otel_handler)
logging.getLogger().setLevel(logging.INFO)

broker = InMemoryBroker().with_middlewares(OtelMiddleware())


@broker.task
async def add_one(value: int) -> int:
    with logfire.span("add_one"):
        logfire.info(f"Starting add_one with value: {value}")
        result = value + 1
        logfire.info(f"Result of add_one: {result}")
    return result


async def main() -> None:
    with logfire.span("main"):
        await broker.startup()
        task = await add_one.kiq(1)
        result = await task.wait_result(timeout=2)
        print(f"Task execution took: {result.execution_time} seconds.")
        if not result.is_err:
            print(f"Returned value: {result.return_value}")
        else:
            print("Error found while executing task.")
        await broker.shutdown()


if __name__ == "__main__":
    asyncio.run(main())

pure opentelemetry

import asyncio
import logging
from taskiq import InMemoryBroker, TaskiqMiddleware
from taskiq.message import TaskiqMessage
from taskiq.result import TaskiqResult

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.context import attach, detach
from opentelemetry.trace import Status, StatusCode

from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, ConsoleLogExporter
from opentelemetry._logs import set_logger_provider


class OtelMiddleware(TaskiqMiddleware):
    def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
        self.tracer = trace.get_tracer(__name__)
        self.span = self.tracer.start_span(name=message.task_name)
        self.span.set_attribute("taskiq.task_id", message.task_id)
        self.span.set_attribute("taskiq.task_name", message.task_name)
        self.token = attach(trace.set_span_in_context(self.span))
        return message

    def post_execute(self, message: TaskiqMessage, result: TaskiqResult) -> None:
        self.span.set_status(Status(StatusCode.OK))
        self.span.end()
        detach(self.token)

    def on_error(self, message: TaskiqMessage, exception: BaseException) -> None:
        self.span.set_status(Status(StatusCode.ERROR, description=str(exception)))
        self.span.record_exception(exception)
        self.span.end()
        detach(self.token)


provider = TracerProvider()
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)

logger_provider = LoggerProvider()
set_logger_provider(logger_provider)
logger_provider.add_log_record_processor(BatchLogRecordProcessor(ConsoleLogExporter()))
otel_handler = LoggingHandler()
logging.getLogger().addHandler(otel_handler)
logging.getLogger().setLevel(logging.INFO)

broker = InMemoryBroker().with_middlewares(OtelMiddleware())


@broker.task
async def add_one(value: int) -> int:
    logging.info(f"Starting add_one with value: {value}")
    result = value + 1
    logging.info(f"Result of add_one: {result}")
    return result


async def main() -> None:
    tracer = trace.get_tracer(__name__)
    with tracer.start_as_current_span("main"):
        await broker.startup()
        task = await add_one.kiq(1)
        result = await task.wait_result(timeout=2)
        print(f"Task execution took: {result.execution_time} seconds.")
        if not result.is_err:
            print(f"Returned value: {result.return_value}")
        else:
            print("Error found while executing task.")
        await broker.shutdown()


if __name__ == "__main__":
    asyncio.run(main())

[danfimov]

I saw this implementation of middleware for Opentelemtry if you need it right now:

Details

import typing as tp
from collections.abc import Coroutine

from opentelemetry import context as context_api
from opentelemetry import trace
from opentelemetry.sdk.trace import Span, TracerProvider
from opentelemetry.trace.propagation import _SPAN_KEY
from taskiq import TaskiqMessage, TaskiqMiddleware, TaskiqResult


class OpenTelemetryTaskiqMiddleware(TaskiqMiddleware):
    INSTRUMENTATION_NAME = 'taskiq'
    TASK_ARGS_LENGTH_LIMIT = 1000
    MSG_SPAN_ATTRIBUTE = '__span'
    MSG_CTX_TOKEN_ATTRIBUTE = '__otlp_ctx_token'  # noqa: S105

    def __init__(self, tracer_provider: TracerProvider | None = None) -> None:
        if not tracer_provider:
            tracer_provider = tp.cast('TracerProvider', trace.get_tracer_provider())
        self.tracer_provider = tracer_provider
        self.tracer = tracer_provider.get_tracer(self.INSTRUMENTATION_NAME, '0.1.0')
        super().__init__()

    def __serialize_args_and_kwargs(self, message: TaskiqMessage) -> tuple[str, str]:
        task_args_str = str(message.args)
        if len(task_args_str) > self.TASK_ARGS_LENGTH_LIMIT:
            task_args_str = f'{task_args_str[: self.TASK_ARGS_LENGTH_LIMIT]}...'
        task_kwargs_str = str(message.kwargs)
        if len(task_kwargs_str) > self.TASK_ARGS_LENGTH_LIMIT:
            task_kwargs_str = f'{task_kwargs_str[: self.TASK_ARGS_LENGTH_LIMIT]}...'
        return task_args_str, task_kwargs_str

    def __serialize_labels(self, message: TaskiqMessage) -> str:
        labels_str = str(message.labels)
        if len(labels_str) > self.TASK_ARGS_LENGTH_LIMIT:
            labels_str = f'{labels_str[: self.TASK_ARGS_LENGTH_LIMIT]}...'
        return labels_str

    def pre_send(self, message: TaskiqMessage) -> TaskiqMessage:
        """Pre-send hook to attach task name to Taskiq message."""
        task_args_str, task_kwargs_str = self.__serialize_args_and_kwargs(message)
        span = self.tracer.start_span(
            f'Post Task: {message.task_name}',
            kind=trace.SpanKind.PRODUCER,
            attributes={
                'task.name': message.task_name,
                'task.id': message.task_id,
                'task.args': task_args_str,
                'task.kwargs': task_kwargs_str,
                'messaging.message.id': message.task_id,
            },
        )
        setattr(message, self.MSG_SPAN_ATTRIBUTE, span)
        ctx = span.get_span_context()
        service_name = self.tracer_provider.resource.attributes.get('service.name')
        message.labels['otlp.caller_trace_id'] = format(ctx.trace_id, '032x')
        message.labels['otlp.caller_span_id'] = format(ctx.span_id, '016x')
        message.labels['otlp.caller_service_name'] = service_name
        return message

    def post_send(self, message: TaskiqMessage) -> None:
        """Post-send hook."""
        span: Span | None = getattr(message, self.MSG_SPAN_ATTRIBUTE, None)
        if not span:
            return
        labels_str = self.__serialize_labels(message)
        span.set_attribute('task.labels', labels_str)
        span.set_status(trace.Status(trace.StatusCode.OK))
        span.end()

    def pre_execute(
        self,
        message: 'TaskiqMessage',
    ) -> TaskiqMessage | Coroutine[tp.Any, tp.Any, TaskiqMessage]:
        """Post execute hook."""
        producer_span_id = message.labels.get('otlp.caller_span_id')
        producer_trace_id = message.labels.get('otlp.caller_trace_id')
        producer_context: trace.SpanContext | None = None
        if producer_span_id and producer_trace_id:
            producer_context = trace.SpanContext(
                trace_id=int(producer_trace_id, 16),
                span_id=int(producer_span_id, 16),
                is_remote=True,
            )
        span = self.tracer.start_span(
            f'Handle: {message.task_name} / {message.task_id}',
            kind=trace.SpanKind.CONSUMER,
            links=(
                [
                    trace.Link(
                        producer_context,
                        {
                            'kind': 'producer',
                            'caller.service.name': message.labels.get('otlp.caller_service_name', ''),
                            'messaging.message.id': message.task_id,
                        },
                    ),
                ]
                if producer_context
                else None
            ),
            attributes={
                'messaging.message.id': message.task_id,
                'task.name': message.task_name,
                'task.id': message.task_id,
                'task.args': str(message.args),
                'task.kwargs': str(message.kwargs),
                'task.labels': str(message.labels),
                'caller.otlp.trace_id': message.labels.get('otlp.caller_trace_id', ''),
                'caller.otlp.span_id': message.labels.get('otlp.caller_span_id', ''),
                'caller.otlp.service.name': message.labels.get('otlp.caller_service_name', ''),
            },
        )

        trace.use_span(span, set_status_on_exception=True)
        context_token = context_api.attach(context_api.set_value(_SPAN_KEY, span))
        setattr(message, self.MSG_SPAN_ATTRIBUTE, span)
        setattr(message, self.MSG_CTX_TOKEN_ATTRIBUTE, context_token)
        return message

    def post_execute(
        self,
        message: 'TaskiqMessage',
        result: 'TaskiqResult[tp.Any]',  # noqa: ARG002
    ) -> None:
        """Post-execute hook."""
        span: Span | None = getattr(message, self.MSG_SPAN_ATTRIBUTE, None)
        ctx_token = getattr(message, self.MSG_CTX_TOKEN_ATTRIBUTE, None)
        if ctx_token:
            context_api.detach(ctx_token)
        if span:
            span.set_status(trace.Status(trace.StatusCode.OK))
            span.end()

    def on_error(
        self,
        message: 'TaskiqMessage',
        result: TaskiqResult[tp.Any],
        exception: BaseException,
    ) -> None:
        """On error hook."""
        span: Span | None = getattr(message, self.MSG_SPAN_ATTRIBUTE, None)
        if span:
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(exception)))
            span.end()
        super().on_error(message, result, exception)

Also Opentelemetry instrumentation was implemented in this MR (huge thanks for @soapun btw): #525. So I hope it will be a part of takiq in a future release

[s3rius]

@ZhymabekRoman @ktnlvr @karansharmapaypay @febus982 @ipfans @nickderobertis, opentelemetry instrumentation support just merged in master branch. Will be released on next release.

I'm closing this discussion as it was resolved.