Function async_wrapper

tensorrt_llm/llmapi/tracing.py:201–223

(*args, **kwargs)

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            span_name = name if name is not None else func.__name__
            if global_otlp_tracer() is None:
                return await func(*args, **kwargs)

            trace_headers = None
            for arg in list(args) + list(kwargs.values()):
                if hasattr(arg, 'headers'):
                    trace_headers = extract_trace_context(arg.headers)
                    break

            with global_otlp_tracer().start_as_current_span(
                    span_name, kind=SpanKind.SERVER,
                    context=trace_headers) as span:
                try:
                    result = await func(*args, **kwargs)
                    span.set_status(Status(StatusCode.OK))
                    return result
                except Exception as e:
                    span.record_exception(e)
                    span.set_status(
                        Status(StatusCode.ERROR, f"An error occurred: {e}"))
                    raise e

        return async_wrapper