nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/observability/exporter/processing_exporter.py

@@ -0,0 +1,322 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from abc import abstractmethod
+from collections.abc import Coroutine
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from aiq.builder.context import AIQContextState
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.observability.exporter.base_exporter import BaseExporter
+from aiq.observability.mixin.type_introspection_mixin import TypeIntrospectionMixin
+from aiq.observability.processor.callback_processor import CallbackProcessor
+from aiq.observability.processor.processor import Processor
+from aiq.utils.type_utils import DecomposedType
+from aiq.utils.type_utils import override
+
+PipelineInputT = TypeVar("PipelineInputT")
+PipelineOutputT = TypeVar("PipelineOutputT")
+
+logger = logging.getLogger(__name__)
+
+
+class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter, TypeIntrospectionMixin):
+    """A base class for telemetry exporters with processing pipeline support.
+
+    This class extends BaseExporter to add processor pipeline functionality.
+    It manages a chain of processors that can transform items before export.
+
+    The generic types work as follows:
+    - PipelineInputT: The type of items that enter the processing pipeline (e.g., Span)
+    - PipelineOutputT: The type of items after processing through the pipeline (e.g., converted format)
+
+    Key Features:
+    - Processor pipeline management (add, remove, clear)
+    - Type compatibility validation between processors
+    - Pipeline processing with error handling
+    - Automatic type validation before export
+    """
+
+    def __init__(self, context_state: AIQContextState | None = None):
+        """Initialize the processing exporter.
+
+        Args:
+            context_state: The context state to use for the exporter.
+        """
+        super().__init__(context_state)
+        self._processors: list[Processor] = []  # List of processors that implement process(item) -> item
+
+    def add_processor(self, processor: Processor) -> None:
+        """Add a processor to the processing pipeline.
+
+        Processors are executed in the order they are added.
+        Processors can transform between any types (T -> U).
+
+        Args:
+            processor: The processor to add to the pipeline
+        """
+
+        # Check if the processor is compatible with the last processor in the pipeline
+        if len(self._processors) > 0:
+            try:
+                if not issubclass(processor.input_class, self._processors[-1].output_class):
+                    raise ValueError(f"Processor {processor.__class__.__name__} input type {processor.input_type} "
+                                     f"is not compatible with the {self._processors[-1].__class__.__name__} "
+                                     f"output type {self._processors[-1].output_type}")
+            except TypeError:
+                # Handle cases where input_class or output_class are generic types that can't be used with issubclass
+                # Fall back to type comparison for generic types
+                logger.warning(
+                    "Cannot use issubclass() for type compatibility check between "
+                    "%s (%s) and %s (%s). Skipping compatibility check.",
+                    processor.__class__.__name__,
+                    processor.input_type,
+                    self._processors[-1].__class__.__name__,
+                    self._processors[-1].output_type)
+        self._processors.append(processor)
+
+        # Set up pipeline continuation callback for processors that support it
+        if isinstance(processor, CallbackProcessor):
+            # Create a callback that continues processing through the rest of the pipeline
+            async def pipeline_callback(item):
+                await self._continue_pipeline_after(processor, item)
+
+            processor.set_done_callback(pipeline_callback)
+
+    def remove_processor(self, processor: Processor) -> None:
+        """Remove a processor from the processing pipeline.
+
+        Args:
+            processor: The processor to remove from the pipeline
+        """
+        if processor in self._processors:
+            self._processors.remove(processor)
+
+    def clear_processors(self) -> None:
+        """Clear all processors from the pipeline."""
+        self._processors.clear()
+
+    async def _pre_start(self) -> None:
+        if len(self._processors) > 0:
+            first_processor = self._processors[0]
+            last_processor = self._processors[-1]
+
+            # validate that the first processor's input type is compatible with the exporter's input type
+            try:
+                if not issubclass(first_processor.input_class, self.input_class):
+                    raise ValueError(f"Processor {first_processor.__class__.__name__} input type "
+                                     f"{first_processor.input_type} is not compatible with the "
+                                     f"{self.input_type} input type")
+            except TypeError as e:
+                # Handle cases where classes are generic types that can't be used with issubclass
+                logger.warning(
+                    "Cannot validate type compatibility between %s (%s) "
+                    "and exporter (%s): %s. Skipping validation.",
+                    first_processor.__class__.__name__,
+                    first_processor.input_type,
+                    self.input_type,
+                    e)
+
+            # Validate that the last processor's output type is compatible with the exporter's output type
+            try:
+                if not DecomposedType.is_type_compatible(last_processor.output_type, self.output_type):
+                    raise ValueError(f"Processor {last_processor.__class__.__name__} output type "
+                                     f"{last_processor.output_type} is not compatible with the "
+                                     f"{self.output_type} output type")
+            except TypeError as e:
+                # Handle cases where classes are generic types that can't be used with issubclass
+                logger.warning(
+                    "Cannot validate type compatibility between %s (%s) "
+                    "and exporter (%s): %s. Skipping validation.",
+                    last_processor.__class__.__name__,
+                    last_processor.output_type,
+                    self.output_type,
+                    e)
+
+    async def _process_pipeline(self, item: PipelineInputT) -> PipelineOutputT:
+        """Process item through all registered processors.
+
+        Args:
+            item (PipelineInputT): The item to process (starts as PipelineInputT, can transform to PipelineOutputT)
+
+        Returns:
+            PipelineOutputT: The processed item after running through all processors
+        """
+        return await self._process_through_processors(self._processors, item)  # type: ignore
+
+    async def _process_through_processors(self, processors: list[Processor], item: Any) -> Any:
+        """Process an item through a list of processors.
+
+        Args:
+            processors (list[Processor]): List of processors to run the item through
+            item (Any): The item to process
+
+        Returns:
+            The processed item after running through all processors
+        """
+        processed_item = item
+        for processor in processors:
+            try:
+                processed_item = await processor.process(processed_item)
+            except Exception as e:
+                logger.error("Error in processor %s: %s", processor.__class__.__name__, e, exc_info=True)
+                # Continue with unprocessed item rather than failing
+        return processed_item
+
+    async def _export_final_item(self, processed_item: Any, raise_on_invalid: bool = False) -> None:
+        """Export a processed item with proper type handling.
+
+        Args:
+            processed_item (Any): The item to export
+            raise_on_invalid (bool): If True, raise ValueError for invalid types instead of logging warning
+        """
+        if isinstance(processed_item, list):
+            if len(processed_item) > 0:
+                await self.export_processed(processed_item)
+            else:
+                logger.debug("Skipping export of empty batch")
+        elif isinstance(processed_item, self.output_class):
+            await self.export_processed(processed_item)
+        else:
+            if raise_on_invalid:
+                raise ValueError(f"Processed item {processed_item} is not a valid output type. "
+                                 f"Expected {self.output_class} or list[{self.output_class}]")
+            logger.warning("Processed item %s is not a valid output type for export", processed_item)
+
+    async def _continue_pipeline_after(self, source_processor: Processor, item: Any) -> None:
+        """Continue processing an item through the pipeline after a specific processor.
+
+        This is used when processors (like BatchingProcessor) need to inject items
+        back into the pipeline flow to continue through downstream processors.
+
+        Args:
+            source_processor (Processor): The processor that generated the item
+            item (Any): The item to continue processing through the remaining pipeline
+        """
+        try:
+            # Find the source processor's position
+            try:
+                source_index = self._processors.index(source_processor)
+            except ValueError:
+                logger.error("Source processor %s not found in pipeline", source_processor.__class__.__name__)
+                return
+
+            # Process through remaining processors (skip the source processor)
+            remaining_processors = self._processors[source_index + 1:]
+            processed_item = await self._process_through_processors(remaining_processors, item)
+
+            # Export the final result
+            await self._export_final_item(processed_item)
+
+        except Exception as e:
+            logger.error("Failed to continue pipeline processing after %s: %s",
+                         source_processor.__class__.__name__,
+                         e,
+                         exc_info=True)
+
+    async def _export_with_processing(self, item: PipelineInputT) -> None:
+        """Export an item after processing it through the pipeline.
+
+        Args:
+            item: The item to export
+        """
+        try:
+            # Then, run through the processor pipeline
+            final_item: PipelineOutputT = await self._process_pipeline(item)
+
+            # Handle different output types from batch processors
+            if isinstance(final_item, list) and len(final_item) == 0:
+                logger.debug("Skipping export of empty batch from processor pipeline")
+                return
+
+            await self._export_final_item(final_item, raise_on_invalid=True)
+
+        except Exception as e:
+            logger.error("Failed to export item '%s': %s", item, e, exc_info=True)
+            raise
+
+    @override
+    def export(self, event: IntermediateStep) -> None:
+        """Export an IntermediateStep event through the processing pipeline.
+
+        This method converts the IntermediateStep to the expected PipelineInputT type,
+        processes it through the pipeline, and exports the result.
+
+        Args:
+            event (IntermediateStep): The event to be exported.
+        """
+        # Convert IntermediateStep to PipelineInputT and create export task
+        if isinstance(event, self.input_class):
+            input_item: PipelineInputT = event  # type: ignore
+            coro = self._export_with_processing(input_item)
+            self._create_export_task(coro)
+        else:
+            logger.warning("Event %s is not compatible with input type %s", event, self.input_type)
+
+    @abstractmethod
+    async def export_processed(self, item: PipelineOutputT | list[PipelineOutputT]) -> None:
+        """Export the processed item.
+
+        This method must be implemented by concrete exporters to handle
+        the actual export logic after the item has been processed through the pipeline.
+
+        Args:
+            item: The processed item to export (PipelineOutputT type)
+        """
+        pass
+
+    def _create_export_task(self, coro: Coroutine):
+        """Create task with minimal overhead but proper tracking."""
+        if not self._running:
+            logger.warning("%s: Attempted to create export task while not running", self.name)
+            return
+
+        try:
+            task = asyncio.create_task(coro)
+            self._tasks.add(task)
+            task.add_done_callback(self._tasks.discard)
+
+        except Exception as e:
+            logger.error("%s: Failed to create task: %s", self.name, e, exc_info=True)
+            raise
+
+    @override
+    async def _cleanup(self):
+        """Enhanced cleanup that shuts down all shutdown-aware processors.
+
+        Each processor is responsible for its own cleanup, including routing
+        any final batches through the remaining pipeline via their done callbacks.
+        """
+        # Shutdown all processors that support it
+        shutdown_tasks = []
+        for processor in getattr(self, '_processors', []):
+            shutdown_method = getattr(processor, 'shutdown', None)
+            if shutdown_method:
+                logger.debug("Shutting down processor: %s", processor.__class__.__name__)
+                shutdown_tasks.append(shutdown_method())
+
+        if shutdown_tasks:
+            try:
+                await asyncio.gather(*shutdown_tasks, return_exceptions=True)
+                logger.debug("Successfully shut down %d processors", len(shutdown_tasks))
+            except Exception as e:
+                logger.error("Error shutting down processors: %s", e, exc_info=True)
+
+        # Call parent cleanup
+        await super()._cleanup()

aiq/observability/exporter/raw_exporter.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import abstractmethod
+from typing import TypeVar
+
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.observability.exporter.processing_exporter import ProcessingExporter
+from aiq.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+InputT = TypeVar("InputT")
+OutputT = TypeVar("OutputT")
+
+
+class RawExporter(ProcessingExporter[InputT, OutputT]):
+    """A base class for exporting raw intermediate steps.
+
+    This class provides a base implementation for telemetry exporters that
+    work directly with IntermediateStep objects. It can optionally process
+    them through a pipeline before export.
+
+    The flow is: IntermediateStep -> [Processing Pipeline] -> OutputT -> Export
+
+    Args:
+        context_state (AIQContextState, optional): The context state to use for the exporter. Defaults to None.
+    """
+
+    @abstractmethod
+    async def export_processed(self, item: OutputT):
+        pass
+
+    @override
+    def export(self, event: IntermediateStep):
+        if not isinstance(event, IntermediateStep):
+            return
+
+        self._create_export_task(self._export_with_processing(event))  # type: ignore

aiq/observability/exporter/span_exporter.py

@@ -0,0 +1,265 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import re
+from abc import abstractmethod
+from typing import TypeVar
+
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepState
+from aiq.data_models.intermediate_step import TraceMetadata
+from aiq.data_models.span import MimeTypes
+from aiq.data_models.span import Span
+from aiq.data_models.span import SpanAttributes
+from aiq.data_models.span import SpanContext
+from aiq.data_models.span import event_type_to_span_kind
+from aiq.observability.exporter.base_exporter import IsolatedAttribute
+from aiq.observability.exporter.processing_exporter import ProcessingExporter
+from aiq.observability.mixin.serialize_mixin import SerializeMixin
+from aiq.observability.utils.dict_utils import merge_dicts
+from aiq.observability.utils.time_utils import ns_timestamp
+from aiq.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+InputSpanT = TypeVar("InputSpanT")
+OutputSpanT = TypeVar("OutputSpanT")
+
+
+class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
+    """Abstract base class for span exporters with processing pipeline support.
+
+    This class specializes ProcessingExporter for span-based telemetry export. It converts
+    IntermediateStep events into Span objects and supports processing pipelines for
+    span transformation before export.
+
+    The generic types work as follows:
+    - InputSpanT: The type of spans that enter the processing pipeline (typically Span)
+    - OutputSpanT: The type of spans after processing through the pipeline (e.g., OtelSpan)
+
+    Key Features:
+    - Automatic span creation from IntermediateStep events
+    - Span lifecycle management (start/end event tracking)
+    - Processing pipeline support via ProcessingExporter
+    - Metadata and attribute handling
+    - Usage information tracking
+    - Automatic isolation of mutable state for concurrent execution using descriptors
+
+    Inheritance Hierarchy:
+    - BaseExporter: Core event subscription and lifecycle management + DescriptorIsolationMixin
+    - ProcessingExporter: Adds processor pipeline functionality
+    - SpanExporter: Specializes for span creation and export
+
+    Event Processing Flow:
+    1. IntermediateStep (START) → Create Span → Add to tracking
+    2. IntermediateStep (END) → Complete Span → Process through pipeline → Export
+
+    Args:
+        context_state (AIQContextState, optional): The context state to use for the exporter. Defaults to None.
+    """
+
+    # Use descriptors for automatic isolation of span-specific state
+    _outstanding_spans: IsolatedAttribute[dict] = IsolatedAttribute(dict)
+    _span_stack: IsolatedAttribute[dict] = IsolatedAttribute(dict)
+    _metadata_stack: IsolatedAttribute[dict] = IsolatedAttribute(dict)
+
+    @abstractmethod
+    async def export_processed(self, item: OutputSpanT) -> None:
+        """Export the processed span.
+
+        Args:
+            item (OutputSpanT): The processed span to export.
+        """
+        pass
+
+    @override
+    def export(self, event: IntermediateStep) -> None:
+        """The main logic that reacts to each IntermediateStep.
+
+        Args:
+            event (IntermediateStep): The event to process.
+        """
+        if not isinstance(event, IntermediateStep):
+            return
+
+        if (event.event_state == IntermediateStepState.START):
+            self._process_start_event(event)
+        elif (event.event_state == IntermediateStepState.END):
+            self._process_end_event(event)
+
+    def _process_start_event(self, event: IntermediateStep):
+        """Process the start event of an intermediate step.
+
+        Args:
+            event (IntermediateStep): The event to process.
+        """
+
+        parent_span = None
+        span_ctx = None
+
+        # Look up the parent span to establish hierarchy
+        # event.parent_id is the UUID of the last START step with a different UUID from current step
+        # This maintains proper parent-child relationships in the span tree
+        # Skip lookup if parent_id is "root" (indicates this is a top-level span)
+        if len(self._span_stack) > 0 and event.parent_id and event.parent_id != "root":
+
+            parent_span = self._span_stack.get(event.parent_id, None)
+            if parent_span is None:
+                logger.warning("No parent span found for step %s", event.UUID)
+                return
+
+            parent_span = parent_span.model_copy() if isinstance(parent_span, Span) else None
+            if parent_span and parent_span.context:
+                span_ctx = SpanContext(trace_id=parent_span.context.trace_id)
+
+        # Extract start/end times from the step
+        # By convention, `span_event_timestamp` is the time we started, `event_timestamp` is the time we ended.
+        # If span_event_timestamp is missing, we default to event_timestamp (meaning zero-length).
+        s_ts = event.payload.span_event_timestamp or event.payload.event_timestamp
+        start_ns = ns_timestamp(s_ts)
+
+        # Optional: embed the LLM/tool name if present
+        if event.payload.name:
+            sub_span_name = f"{event.payload.name}"
+        else:
+            sub_span_name = f"{event.payload.event_type}"
+
+        sub_span = Span(
+            name=sub_span_name,
+            parent=parent_span,
+            context=span_ctx,
+            attributes={
+                "aiq.event_type": event.payload.event_type.value,
+                "aiq.function.id": event.function_ancestry.function_id if event.function_ancestry else "unknown",
+                "aiq.function.name": event.function_ancestry.function_name if event.function_ancestry else "unknown",
+                "aiq.subspan.name": event.payload.name or "",
+                "aiq.event_timestamp": event.event_timestamp,
+                "aiq.framework": event.payload.framework.value if event.payload.framework else "unknown",
+            },
+            start_time=start_ns)
+
+        span_kind = event_type_to_span_kind(event.event_type)
+        sub_span.set_attribute("aiq.span.kind", span_kind.value)
+
+        if event.payload.data and event.payload.data.input:
+            match = re.search(r"Human:\s*Question:\s*(.*)", str(event.payload.data.input))
+            if match:
+                human_question = match.group(1).strip()
+                sub_span.set_attribute(SpanAttributes.INPUT_VALUE.value, human_question)
+            else:
+                serialized_input, is_json = self._serialize_payload(event.payload.data.input)
+                sub_span.set_attribute(SpanAttributes.INPUT_VALUE.value, serialized_input)
+                sub_span.set_attribute(SpanAttributes.INPUT_MIME_TYPE.value,
+                                       MimeTypes.JSON.value if is_json else MimeTypes.TEXT.value)
+
+        # Add metadata to the metadata stack
+        start_metadata = event.payload.metadata or {}
+
+        if isinstance(start_metadata, dict):
+            self._metadata_stack[event.UUID] = start_metadata  # type: ignore
+        elif isinstance(start_metadata, TraceMetadata):
+            self._metadata_stack[event.UUID] = start_metadata.model_dump()  # type: ignore
+        else:
+            logger.warning("Invalid metadata type for step %s", event.UUID)
+            return
+
+        self._span_stack[event.UUID] = sub_span  # type: ignore
+        self._outstanding_spans[event.UUID] = sub_span  # type: ignore
+
+        logger.debug(
+            "Added span to tracking (outstanding: %d, stack: %d, event_id: %s)",
+            len(self._outstanding_spans),  # type: ignore
+            len(self._span_stack),  # type: ignore
+            event.UUID)
+
+    def _process_end_event(self, event: IntermediateStep):
+        """Process the end event of an intermediate step.
+
+        Args:
+            event (IntermediateStep): The event to process.
+        """
+
+        # Find the subspan that was created in the start event
+        sub_span: Span | None = self._outstanding_spans.pop(event.UUID, None)  # type: ignore
+
+        if sub_span is None:
+            logger.warning("No subspan found for step %s", event.UUID)
+            return
+
+        self._span_stack.pop(event.UUID, None)  # type: ignore
+
+        # Optionally add more attributes from usage_info or data
+        usage_info = event.payload.usage_info
+        if usage_info:
+            sub_span.set_attribute(SpanAttributes.AIQ_USAGE_NUM_LLM_CALLS.value,
+                                   usage_info.num_llm_calls if usage_info.num_llm_calls else 0)
+            sub_span.set_attribute(SpanAttributes.AIQ_USAGE_SECONDS_BETWEEN_CALLS.value,
+                                   usage_info.seconds_between_calls if usage_info.seconds_between_calls else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT.value,
+                                   usage_info.token_usage.prompt_tokens if usage_info.token_usage else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.value,
+                                   usage_info.token_usage.completion_tokens if usage_info.token_usage else 0)
+            sub_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL.value,
+                                   usage_info.token_usage.total_tokens if usage_info.token_usage else 0)
+
+        if event.payload.data and event.payload.data.output is not None:
+            serialized_output, is_json = self._serialize_payload(event.payload.data.output)
+            sub_span.set_attribute(SpanAttributes.OUTPUT_VALUE.value, serialized_output)
+            sub_span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE.value,
+                                   MimeTypes.JSON.value if is_json else MimeTypes.TEXT.value)
+
+        # Merge metadata from start event with end event metadata
+        start_metadata = self._metadata_stack.pop(event.UUID)  # type: ignore
+
+        if start_metadata is None:
+            logger.warning("No metadata found for step %s", event.UUID)
+            return
+
+        end_metadata = event.payload.metadata or {}
+
+        if not isinstance(end_metadata, (dict, TraceMetadata)):
+            logger.warning("Invalid metadata type for step %s", event.UUID)
+            return
+
+        if isinstance(end_metadata, TraceMetadata):
+            end_metadata = end_metadata.model_dump()
+
+        merged_metadata = merge_dicts(start_metadata, end_metadata)
+        serialized_metadata, is_json = self._serialize_payload(merged_metadata)
+        sub_span.set_attribute("aiq.metadata", serialized_metadata)
+        sub_span.set_attribute("aiq.metadata.mime_type", MimeTypes.JSON.value if is_json else MimeTypes.TEXT.value)
+
+        end_ns = ns_timestamp(event.payload.event_timestamp)
+
+        # End the subspan
+        sub_span.end(end_time=end_ns)
+
+        # Export the span with processing pipeline
+        self._create_export_task(self._export_with_processing(sub_span))  # type: ignore
+
+    @override
+    async def _cleanup(self):
+        """Clean up any remaining spans."""
+        if self._outstanding_spans:  # type: ignore
+            logger.warning("Not all spans were closed. Remaining: %s", self._outstanding_spans)  # type: ignore
+
+            for span_info in self._outstanding_spans.values():  # type: ignore
+                span_info.end()
+
+        self._outstanding_spans.clear()  # type: ignore
+        self._span_stack.clear()  # type: ignore
+        self._metadata_stack.clear()  # type: ignore
+        await super()._cleanup()