nvidia-nat 1.2.0a20250813__py3-none-any.whl

nat/observability/mixin/type_introspection_mixin.py

@@ -0,0 +1,183 @@
+# 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.
+
+from functools import lru_cache
+from typing import Any
+from typing import get_args
+from typing import get_origin
+
+
+class TypeIntrospectionMixin:
+    """Mixin class providing type introspection capabilities for generic classes.
+
+    This mixin extracts type information from generic class definitions,
+    allowing classes to determine their InputT and OutputT types at runtime.
+    """
+
+    def _find_generic_types(self) -> tuple[type[Any], type[Any]] | None:
+        """
+        Recursively search through the inheritance hierarchy to find generic type parameters.
+
+        This method handles cases where a class inherits from a generic parent class,
+        resolving the concrete types through the inheritance chain.
+
+        Returns:
+            tuple[type[Any], type[Any]] | None: (input_type, output_type) if found, None otherwise
+        """
+        # First, try to find types directly in this class's __orig_bases__
+        for base_cls in getattr(self.__class__, '__orig_bases__', []):
+            base_cls_args = get_args(base_cls)
+
+            # Direct case: MyClass[InputT, OutputT]
+            if len(base_cls_args) >= 2:
+                return base_cls_args[0], base_cls_args[1]
+
+            # Indirect case: MyClass[SomeGeneric[ConcreteType]]
+            # Need to resolve the generic parent's types
+            if len(base_cls_args) == 1:
+                base_origin = get_origin(base_cls)
+                if base_origin and hasattr(base_origin, '__orig_bases__'):
+                    # Look at the parent's generic definition
+                    for parent_base in getattr(base_origin, '__orig_bases__', []):
+                        parent_args = get_args(parent_base)
+                        if len(parent_args) >= 2:
+                            # Found the pattern: ParentClass[T, list[T]]
+                            # Substitute T with our concrete type
+                            concrete_type = base_cls_args[0]
+                            input_type = self._substitute_type_var(parent_args[0], concrete_type)
+                            output_type = self._substitute_type_var(parent_args[1], concrete_type)
+                            return input_type, output_type
+
+        return None
+
+    def _substitute_type_var(self, type_expr: Any, concrete_type: type) -> type[Any]:
+        """
+        Substitute TypeVar in a type expression with a concrete type.
+
+        Args:
+            type_expr: The type expression potentially containing TypeVars
+            concrete_type: The concrete type to substitute
+
+        Returns:
+            The type expression with TypeVars substituted
+        """
+        from typing import TypeVar
+
+        # If it's a TypeVar, substitute it
+        if isinstance(type_expr, TypeVar):
+            return concrete_type
+
+        # If it's a generic type like list[T], substitute the args
+        origin = get_origin(type_expr)
+        args = get_args(type_expr)
+
+        if origin and args:
+            # Recursively substitute in the arguments
+            new_args = tuple(self._substitute_type_var(arg, concrete_type) for arg in args)
+            # Reconstruct the generic type
+            return origin[new_args]
+
+        # Otherwise, return as-is
+        return type_expr
+
+    @property
+    @lru_cache
+    def input_type(self) -> type[Any]:
+        """
+        Get the input type of the class. The input type is determined by the generic parameters of the class.
+
+        For example, if a class is defined as `MyClass[list[int], str]`, the `input_type` is `list[int]`.
+
+        Returns
+        -------
+        type[Any]
+            The input type specified in the generic parameters
+
+        Raises
+        ------
+        ValueError
+            If the input type cannot be determined from the class definition
+        """
+        types = self._find_generic_types()
+        if types:
+            return types[0]
+
+        raise ValueError(f"Could not find input type for {self.__class__.__name__}")
+
+    @property
+    @lru_cache
+    def output_type(self) -> type[Any]:
+        """
+        Get the output type of the class. The output type is determined by the generic parameters of the class.
+
+        For example, if a class is defined as `MyClass[list[int], str]`, the `output_type` is `str`.
+
+        Returns
+        -------
+        type[Any]
+            The output type specified in the generic parameters
+
+        Raises
+        ------
+        ValueError
+            If the output type cannot be determined from the class definition
+        """
+        types = self._find_generic_types()
+        if types:
+            return types[1]
+
+        raise ValueError(f"Could not find output type for {self.__class__.__name__}")
+
+    @property
+    @lru_cache
+    def input_class(self) -> type:
+        """
+        Get the python class of the input type. This is the class that can be used to check if a value is an
+        instance of the input type. It removes any generic or annotation information from the input type.
+
+        For example, if the input type is `list[int]`, the `input_class` is `list`.
+
+        Returns
+        -------
+        type
+            The python type of the input type
+        """
+        input_origin = get_origin(self.input_type)
+
+        if input_origin is None:
+            return self.input_type
+
+        return input_origin
+
+    @property
+    @lru_cache
+    def output_class(self) -> type:
+        """
+        Get the python class of the output type. This is the class that can be used to check if a value is an
+        instance of the output type. It removes any generic or annotation information from the output type.
+
+        For example, if the output type is `list[int]`, the `output_class` is `list`.
+
+        Returns
+        -------
+        type
+            The python type of the output type
+        """
+        output_origin = get_origin(self.output_type)
+
+        if output_origin is None:
+            return self.output_type
+
+        return output_origin

nat/observability/processor/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/observability/processor/batching_processor.py

@@ -0,0 +1,310 @@
+# 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
+import time
+from collections import deque
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from nat.observability.processor.callback_processor import CallbackProcessor
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar('T')
+
+
+class BatchingProcessor(CallbackProcessor[T, list[T]], Generic[T]):
+    """Pass-through batching processor that accumulates items and outputs batched lists.
+
+    This processor extends CallbackProcessor[T, List[T]] to provide batching functionality.
+    It accumulates individual items and outputs them as batches when size or time thresholds
+    are met. The batched output continues through the processing pipeline.
+
+    CRITICAL: Implements proper cleanup to ensure NO ITEMS ARE LOST during shutdown.
+    The ProcessingExporter._cleanup() method calls shutdown() on all processors.
+
+    Key Features:
+    - Pass-through design: Processor[T, List[T]]
+    - Size-based and time-based batching
+    - Pipeline flow: batches continue through downstream processors
+    - GUARANTEED: No items lost during cleanup
+    - Comprehensive statistics and monitoring
+    - Proper cleanup and shutdown handling
+    - High-performance async implementation
+    - Back-pressure handling with queue limits
+
+    Pipeline Flow:
+        Normal processing: Individual items → BatchingProcessor → List[items] → downstream processors → export
+        Time-based flush: Scheduled batches automatically continue through remaining pipeline
+        Shutdown: Final batch immediately routed through remaining pipeline
+
+    Cleanup Guarantee:
+        When shutdown() is called, this processor:
+        1. Stops accepting new items
+        2. Creates final batch from all queued items
+        3. Immediately routes final batch through remaining pipeline via callback
+        4. Ensures zero data loss with no external coordination needed
+
+    Usage in Pipeline:
+        ```python
+        # Individual spans → Batched spans → Continue through downstream processors
+        exporter.add_processor(BatchingProcessor[Span](batch_size=100))  # Auto-wired with pipeline callback
+        exporter.add_processor(FilterProcessor())  # Processes List[Span] from batching
+        exporter.add_processor(TransformProcessor())  # Further processing
+        ```
+
+    Args:
+        batch_size: Maximum items per batch (default: 100)
+        flush_interval: Max seconds to wait before flushing (default: 5.0)
+        max_queue_size: Maximum items to queue before blocking (default: 1000)
+        drop_on_overflow: If True, drop items when queue is full (default: False)
+        shutdown_timeout: Max seconds to wait for final batch processing (default: 10.0)
+
+    Note:
+        The done_callback for pipeline integration is automatically set by ProcessingExporter
+        when the processor is added to a pipeline. For standalone usage, call set_done_callback().
+    """
+
+    def __init__(self,
+                 batch_size: int = 100,
+                 flush_interval: float = 5.0,
+                 max_queue_size: int = 1000,
+                 drop_on_overflow: bool = False,
+                 shutdown_timeout: float = 10.0):
+        self._batch_size = batch_size
+        self._flush_interval = flush_interval
+        self._max_queue_size = max_queue_size
+        self._drop_on_overflow = drop_on_overflow
+        self._shutdown_timeout = shutdown_timeout
+        self._done_callback: Callable[[list[T]], Awaitable[None]] | None = None
+
+        # Batching state
+        self._batch_queue: deque[T] = deque()
+        self._last_flush_time = time.time()
+        self._flush_task: asyncio.Task | None = None
+        self._batch_lock = asyncio.Lock()
+        self._shutdown_requested = False
+        self._shutdown_complete = False
+        self._shutdown_complete_event = asyncio.Event()
+
+        # Callback for immediate export of scheduled batches
+        self._done = None
+
+        # Statistics
+        self._batches_created = 0
+        self._items_processed = 0
+        self._items_dropped = 0
+        self._queue_overflows = 0
+        self._shutdown_batches = 0
+
+    async def process(self, item: T) -> list[T]:
+        """Process an item by adding it to the batch queue.
+
+        Returns a batch when batching conditions are met, otherwise returns empty list.
+        This maintains the Processor[T, List[T]] contract while handling batching logic.
+
+        During shutdown, immediately returns items as single-item batches to ensure
+        no data loss.
+
+        Args:
+            item: The item to add to the current batch
+
+        Returns:
+            List[T]: A batch of items when ready, empty list otherwise
+        """
+        if self._shutdown_requested:
+            # During shutdown, return item immediately as single-item batch
+            # This ensures no items are lost even if shutdown is in progress
+            self._items_processed += 1
+            self._shutdown_batches += 1
+            logger.debug("Shutdown mode: returning single-item batch for item %s", item)
+            return [item]
+
+        async with self._batch_lock:
+            # Handle queue overflow
+            if len(self._batch_queue) >= self._max_queue_size:
+                self._queue_overflows += 1
+
+                if self._drop_on_overflow:
+                    # Drop the item and return empty
+                    self._items_dropped += 1
+                    logger.warning("Dropping item due to queue overflow (dropped: %d)", self._items_dropped)
+                    return []
+                # Force flush to make space, then add item
+                logger.warning("Queue overflow, forcing flush of %d items", len(self._batch_queue))
+                forced_batch = await self._create_batch()
+                if forced_batch:
+                    # Add current item to queue and return the forced batch
+                    self._batch_queue.append(item)
+                    self._items_processed += 1
+                    return forced_batch
+
+            # Add item to batch queue
+            self._batch_queue.append(item)
+            self._items_processed += 1
+
+            # Check flush conditions
+            should_flush = (len(self._batch_queue) >= self._batch_size
+                            or (time.time() - self._last_flush_time) >= self._flush_interval)
+
+            if should_flush:
+                return await self._create_batch()
+            # Schedule a time-based flush if not already scheduled
+            if self._flush_task is None or self._flush_task.done():
+                self._flush_task = asyncio.create_task(self._schedule_flush())
+            return []
+
+    def set_done_callback(self, callback: Callable[[list[T]], Awaitable[None]]):
+        """Set callback function for routing batches through the remaining pipeline.
+
+        This is automatically set by ProcessingExporter.add_processor() to continue
+        batches through downstream processors before final export.
+        """
+        self._done_callback = callback
+
+    async def _schedule_flush(self):
+        """Schedule a flush after the flush interval."""
+        try:
+            await asyncio.sleep(self._flush_interval)
+            async with self._batch_lock:
+                if not self._shutdown_requested and self._batch_queue:
+                    batch = await self._create_batch()
+                    if batch:
+                        # Route scheduled batches through pipeline via callback
+                        if self._done_callback is not None:
+                            try:
+                                await self._done_callback(batch)
+                                logger.debug("Scheduled flush routed batch of %d items through pipeline", len(batch))
+                            except Exception as e:
+                                logger.error("Error routing scheduled batch through pipeline: %s", e, exc_info=True)
+                        else:
+                            logger.warning("Scheduled flush created batch of %d items but no pipeline callback set",
+                                           len(batch))
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.error("Error in scheduled flush: %s", e, exc_info=True)
+
+    async def _create_batch(self) -> list[T]:
+        """Create a batch from the current queue."""
+        if not self._batch_queue:
+            return []
+
+        batch = list(self._batch_queue)
+        self._batch_queue.clear()
+        self._last_flush_time = time.time()
+        self._batches_created += 1
+
+        logger.debug("Created batch of %d items (total: %d items in %d batches)",
+                     len(batch),
+                     self._items_processed,
+                     self._batches_created)
+
+        return batch
+
+    async def force_flush(self) -> list[T]:
+        """Force an immediate flush of all queued items.
+
+        Returns:
+            List[T]: The current batch, empty list if no items queued
+        """
+        async with self._batch_lock:
+            return await self._create_batch()
+
+    async def shutdown(self) -> None:
+        """Shutdown the processor and ensure all items are processed.
+
+        CRITICAL: This method is called by ProcessingExporter._cleanup() to ensure
+        no items are lost during shutdown. It immediately routes any remaining
+        items as a final batch through the rest of the processing pipeline.
+        """
+        if self._shutdown_requested:
+            logger.debug("Shutdown already requested, waiting for completion")
+            # Wait for shutdown to complete using event instead of polling
+            try:
+                await asyncio.wait_for(self._shutdown_complete_event.wait(), timeout=self._shutdown_timeout)
+                logger.debug("Shutdown completion detected via event")
+            except asyncio.TimeoutError:
+                logger.warning("Shutdown completion timeout exceeded (%s seconds)", self._shutdown_timeout)
+            return
+
+        logger.debug("Starting shutdown of BatchingProcessor (queue size: %d)", len(self._batch_queue))
+        self._shutdown_requested = True
+
+        try:
+            # Cancel scheduled flush task
+            if self._flush_task and not self._flush_task.done():
+                self._flush_task.cancel()
+                try:
+                    await self._flush_task
+                except asyncio.CancelledError:
+                    pass
+
+            # Create and route final batch through pipeline
+            async with self._batch_lock:
+                if self._batch_queue:
+                    final_batch = await self._create_batch()
+                    logger.debug("Created final batch of %d items during shutdown", len(final_batch))
+
+                    # Route final batch through pipeline via callback
+                    if self._done_callback is not None:
+                        try:
+                            await self._done_callback(final_batch)
+                            logger.debug(
+                                "Successfully flushed final batch of %d items through pipeline during shutdown",
+                                len(final_batch))
+                        except Exception as e:
+                            logger.error("Error routing final batch through pipeline during shutdown: %s",
+                                         e,
+                                         exc_info=True)
+                    else:
+                        logger.warning("Final batch of %d items created during shutdown but no pipeline callback set",
+                                       len(final_batch))
+                else:
+                    logger.debug("No items remaining during shutdown")
+
+            self._shutdown_complete = True
+            self._shutdown_complete_event.set()
+            logger.debug("BatchingProcessor shutdown completed successfully")
+
+        except Exception as e:
+            logger.error("Error during BatchingProcessor shutdown: %s", e, exc_info=True)
+            self._shutdown_complete = True
+            self._shutdown_complete_event.set()
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get comprehensive batching statistics."""
+        return {
+            "current_queue_size": len(self._batch_queue),
+            "batch_size_limit": self._batch_size,
+            "flush_interval": self._flush_interval,
+            "max_queue_size": self._max_queue_size,
+            "drop_on_overflow": self._drop_on_overflow,
+            "shutdown_timeout": self._shutdown_timeout,
+            "batches_created": self._batches_created,
+            "items_processed": self._items_processed,
+            "items_dropped": self._items_dropped,
+            "queue_overflows": self._queue_overflows,
+            "shutdown_batches": self._shutdown_batches,
+            "shutdown_requested": self._shutdown_requested,
+            "shutdown_complete": self._shutdown_complete,
+            "avg_items_per_batch": self._items_processed / max(1, self._batches_created),
+            "drop_rate": self._items_dropped / max(1, self._items_processed) * 100 if self._items_processed > 0 else 0
+        }

nat/observability/processor/callback_processor.py

@@ -0,0 +1,42 @@
+# 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.
+
+from abc import abstractmethod
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+from typing import TypeVar
+
+from nat.observability.processor.processor import Processor
+
+InputT = TypeVar('InputT')
+OutputT = TypeVar('OutputT')
+
+
+class CallbackProcessor(Processor[InputT, OutputT]):
+    """Abstract base class for processors that support done callbacks.
+
+    Processors inheriting from this class can register callbacks that are
+    invoked when items are ready for further processing or export.
+    """
+
+    @abstractmethod
+    def set_done_callback(self, callback: Callable[[Any], Awaitable[None]]) -> None:
+        """Set a callback function to be invoked when items are processed.
+
+        Args:
+            callback (Callable[[Any], Awaitable[None]]): Function to call with processed items
+        """
+        pass

nat/observability/processor/intermediate_step_serializer.py

@@ -0,0 +1,28 @@
+# 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.
+
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.observability.mixin.serialize_mixin import SerializeMixin
+from nat.observability.processor.processor import Processor
+from nat.utils.type_utils import override
+
+
+class IntermediateStepSerializer(SerializeMixin, Processor[IntermediateStep, str]):
+    """A File processor that exports telemetry traces to a local file."""
+
+    @override
+    async def process(self, item: IntermediateStep) -> str:
+        serialized_payload, _ = self._serialize_payload(item)
+        return serialized_payload

nat/observability/processor/processor.py

@@ -0,0 +1,71 @@
+# 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.
+
+from abc import ABC
+from abc import abstractmethod
+from typing import Generic
+from typing import TypeVar
+
+from nat.observability.mixin.type_introspection_mixin import TypeIntrospectionMixin
+
+InputT = TypeVar('InputT')
+OutputT = TypeVar('OutputT')
+
+
+class Processor(Generic[InputT, OutputT], TypeIntrospectionMixin, ABC):
+    """Generic protocol for processors that can convert between types in export pipelines.
+
+    Processors are the building blocks of processing pipelines in exporters. They can
+    transform data from one type to another, enabling flexible data processing chains.
+
+    The generic types work as follows:
+    - InputT: The type of items that this processor accepts
+    - OutputT: The type of items that this processor produces
+
+    Key Features:
+    - Type-safe transformations through generics
+    - Type introspection capabilities via TypeIntrospectionMixin
+    - Async processing support
+    - Chainable in processing pipelines
+
+    Inheritance Structure:
+    - Inherits from TypeIntrospectionMixin for type introspection capabilities
+    - Implements Generic[InputT, OutputT] for type safety
+    - Abstract base class requiring implementation of process()
+
+    Example:
+        .. code-block:: python
+
+            class SpanToOtelProcessor(Processor[Span, OtelSpan]):
+                async def process(self, item: Span) -> OtelSpan:
+                    return convert_span_to_otel(item)
+
+    Note:
+        Processors are typically added to ProcessingExporter instances to create
+        transformation pipelines. The exporter validates type compatibility between
+        chained processors.
+    """
+
+    @abstractmethod
+    async def process(self, item: InputT) -> OutputT:
+        """Process an item and return a potentially different type.
+
+        Args:
+            item (InputT): The item to process
+
+        Returns:
+            OutputT: The processed item
+        """
+        pass