aiqtoolkit 1.2.0rc2__py3-none-any.whl → 1.2.0rc4__py3-none-any.whl

aiq/observability/exporter/processing_exporter.py

@@ -17,6 +17,7 @@ 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
 
@@ -24,6 +25,7 @@ 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
@@ -89,6 +91,14 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
                     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.
 
@@ -143,20 +153,82 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
         """Process item through all registered processors.
 
         Args:
-            item: The item to process (starts as PipelineInputT, can transform to PipelineOutputT)
+            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 self._processors:
+        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 the export
+                # 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.
 
-        return processed_item  # type: ignore
+        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.
@@ -169,20 +241,11 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
             final_item: PipelineOutputT = await self._process_pipeline(item)
 
             # Handle different output types from batch processors
-            if isinstance(final_item, list):
-                # Empty lists from batch processors should be skipped, not exported
-                if len(final_item) == 0:
-                    logger.debug("Skipping export of empty batch from processor pipeline")
-                    return
-
-                # Non-empty lists should be exported (batch processors)
-                await self.export_processed(final_item)
-            elif isinstance(final_item, self.output_class):
-                # Single items should be exported normally
-                await self.export_processed(final_item)
-            else:
-                raise ValueError(f"Processed item {final_item} is not a valid output type. "
-                                 f"Expected {self.output_class} or list[{self.output_class}]")
+            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)
@@ -235,35 +298,25 @@ class ProcessingExporter(Generic[PipelineInputT, PipelineOutputT], BaseExporter,
 
     @override
     async def _cleanup(self):
-        """Enhanced cleanup that shuts down all shutdown-aware processors."""
+        """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
-        if hasattr(self, '_processors'):
-            shutdown_tasks = []
-            for processor in getattr(self, '_processors', []):
-                if hasattr(processor, 'shutdown'):
-                    logger.debug("Shutting down processor: %s", processor.__class__.__name__)
-                    shutdown_tasks.append(processor.shutdown())
-
-            if shutdown_tasks:
-                try:
-                    await asyncio.gather(*shutdown_tasks, return_exceptions=True)
-                    logger.info("Successfully shut down %d processors", len(shutdown_tasks))
-                except Exception as e:
-                    logger.error("Error shutting down processors: %s", e, exc_info=True)
-
-            # Process final batches from batch processors
-            for processor in getattr(self, '_processors', []):
-                if hasattr(processor, 'has_final_batch') and hasattr(processor, 'get_final_batch'):
-                    if processor.has_final_batch():
-                        final_batch = processor.get_final_batch()
-                        if final_batch:
-                            logger.info("Processing final batch of %d items from %s during cleanup",
-                                        len(final_batch),
-                                        processor.__class__.__name__)
-                            try:
-                                await self.export_processed(final_batch)
-                            except Exception as e:
-                                logger.error("Error processing final batch during cleanup: %s", e, exc_info=True)
+        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.info("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/span_exporter.py

@@ -262,3 +262,4 @@ class SpanExporter(ProcessingExporter[InputSpanT, OutputSpanT], SerializeMixin):
         self._outstanding_spans.clear()  # type: ignore
         self._span_stack.clear()  # type: ignore
         self._metadata_stack.clear()  # type: ignore
+        await super()._cleanup()

aiq/observability/processor/batching_processor.py

@@ -23,17 +23,17 @@ from typing import Any
 from typing import Generic
 from typing import TypeVar
 
-from aiq.observability.processor.processor import Processor
+from aiq.observability.processor.callback_processor import CallbackProcessor
 
 logger = logging.getLogger(__name__)
 
 T = TypeVar('T')
 
 
-class BatchingProcessor(Processor[T, list[T]], Generic[T]):
+class BatchingProcessor(CallbackProcessor[T, list[T]], Generic[T]):
     """Pass-through batching processor that accumulates items and outputs batched lists.
 
-    This processor fits properly into the generics design by implementing Processor[T, List[T]].
+    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.
 
@@ -43,25 +43,31 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
     Key Features:
     - Pass-through design: Processor[T, List[T]]
     - Size-based and time-based batching
-    - Fits into generics processing pipeline design
+    - 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 ProcessingExporter._cleanup() calls shutdown(), this processor:
+        When shutdown() is called, this processor:
         1. Stops accepting new items
-        2. Processes all queued items as final batch
-        3. Returns final batch to continue through pipeline
-        4. Ensures zero data loss during shutdown
+        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 processing
-        exporter.add_processor(BatchingProcessor[Span](batch_size=100))
-        exporter.add_processor(BatchedSpanProcessor())  # Processes List[Span]
+        # 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:
@@ -70,6 +76,10 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
         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,
@@ -77,14 +87,13 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
                  flush_interval: float = 5.0,
                  max_queue_size: int = 1000,
                  drop_on_overflow: bool = False,
-                 shutdown_timeout: float = 10.0,
-                 done_callback: Callable[[list[T]], Awaitable[None]] | None = None):
+                 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 = done_callback
+        self._done_callback: Callable[[list[T]], Awaitable[None]] | None = None
 
         # Batching state
         self._batch_queue: deque[T] = deque()
@@ -93,11 +102,7 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
         self._batch_lock = asyncio.Lock()
         self._shutdown_requested = False
         self._shutdown_complete = False
-        self._shutdown_complete_event: asyncio.Event | None = None
-
-        # Final batch handling for cleanup
-        self._final_batch: list[T] | None = None
-        self._final_batch_processed = False
+        self._shutdown_complete_event = asyncio.Event()
 
         # Callback for immediate export of scheduled batches
         self._done = None
@@ -167,7 +172,11 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
             return []
 
     def set_done_callback(self, callback: Callable[[list[T]], Awaitable[None]]):
-        """Set callback function for immediate export of scheduled batches."""
+        """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):
@@ -178,15 +187,15 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
                 if not self._shutdown_requested and self._batch_queue:
                     batch = await self._create_batch()
                     if batch:
-                        # Immediately export scheduled batches via callback
+                        # Route scheduled batches through pipeline via callback
                         if self._done_callback is not None:
                             try:
                                 await self._done_callback(batch)
-                                logger.debug("Scheduled flush exported batch of %d items", len(batch))
+                                logger.debug("Scheduled flush routed batch of %d items through pipeline", len(batch))
                             except Exception as e:
-                                logger.error("Error exporting scheduled batch: %s", e, exc_info=True)
+                                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 export callback set",
+                            logger.warning("Scheduled flush created batch of %d items but no pipeline callback set",
                                            len(batch))
         except asyncio.CancelledError:
             pass
@@ -223,11 +232,8 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
         """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 creates a final batch from any
-        remaining items and stores it for processing.
-
-        The final batch will be processed by the next process() call or can be
-        retrieved via get_final_batch().
+        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")
@@ -251,13 +257,26 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
                 except asyncio.CancelledError:
                     pass
 
-            # Create final batch from remaining items
+            # Create and route final batch through pipeline
             async with self._batch_lock:
                 if self._batch_queue:
-                    self._final_batch = await self._create_batch()
-                    logger.info("Created final batch of %d items during shutdown", len(self._final_batch))
+                    final_batch = await self._create_batch()
+                    logger.info("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.info("Successfully routed 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:
-                    self._final_batch = []
                     logger.info("No items remaining during shutdown")
 
             self._shutdown_complete = True
@@ -269,30 +288,6 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
             self._shutdown_complete = True
             self._shutdown_complete_event.set()
 
-    def get_final_batch(self) -> list[T]:
-        """Get the final batch created during shutdown.
-
-        This method allows the exporter to retrieve and process any items
-        that were queued when shutdown was called.
-
-        Returns:
-            List[T]: Final batch of items, empty list if none
-        """
-        if self._final_batch is not None:
-            final_batch = self._final_batch
-            self._final_batch = None  # Clear to avoid double processing
-            self._final_batch_processed = True
-            return final_batch
-        return []
-
-    def has_final_batch(self) -> bool:
-        """Check if there's a final batch waiting to be processed.
-
-        Returns:
-            bool: True if final batch exists and hasn't been processed
-        """
-        return self._final_batch is not None and not self._final_batch_processed
-
     def get_stats(self) -> dict[str, Any]:
         """Get comprehensive batching statistics."""
         return {
@@ -309,8 +304,6 @@ class BatchingProcessor(Processor[T, list[T]], Generic[T]):
             "shutdown_batches": self._shutdown_batches,
             "shutdown_requested": self._shutdown_requested,
             "shutdown_complete": self._shutdown_complete,
-            "final_batch_size": len(self._final_batch) if self._final_batch else 0,
-            "final_batch_processed": self._final_batch_processed,
             "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
         }

aiq/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 aiq.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

aiq/observability/processor/processor.py

@@ -63,6 +63,9 @@ class Processor(Generic[InputT, OutputT], TypeIntrospectionMixin, ABC):
         """Process an item and return a potentially different type.
 
         Args:
-            item: The item to process
+            item (InputT): The item to process
+
+        Returns:
+            OutputT: The processed item
         """
         pass

aiq/profiler/calc/calc_runner.py

@@ -34,6 +34,7 @@ from aiq.profiler.calc.data_models import CalcData
 from aiq.profiler.calc.data_models import CalcRunnerConfig
 from aiq.profiler.calc.data_models import CalcRunnerOutput
 from aiq.profiler.calc.data_models import FitConfig
+from aiq.profiler.calc.data_models import FitResults
 from aiq.profiler.calc.data_models import GPUEstimates
 from aiq.profiler.calc.data_models import SizingMetricPerItem
 from aiq.profiler.calc.data_models import SizingMetrics
@@ -408,7 +409,10 @@ class CalcRunner:
         if gpu_estimates.gpu_estimate_by_llm_latency is not None:
             logger.info("GPU estimate by LLM latency: %.2f", gpu_estimates.gpu_estimate_by_llm_latency)
 
-        return CalcRunnerOutput(gpu_estimates=gpu_estimates, calc_data=calc_data)
+        return CalcRunnerOutput(gpu_estimates=gpu_estimates,
+                                calc_data=calc_data,
+                                fit_results=FitResults(llm_latency_fit=self.linear_analyzer.llm_latency_fit,
+                                                       wf_runtime_fit=self.linear_analyzer.wf_runtime_fit))
 
     def plot_concurrency_vs_time_metrics(self, output_dir: Path):
         """Plots concurrency vs. time metrics using pre-computed fits."""

aiq/profiler/calc/data_models.py

@@ -17,6 +17,7 @@ import typing
 from pathlib import Path
 
 from pydantic import BaseModel
+from pydantic import Field
 
 
 class FitConfig(BaseModel):
@@ -76,7 +77,7 @@ class CalcRunnerConfig(BaseModel):
     plot_data: bool = True
 
     # Configuration for linear fit and outlier detection
-    fit_config: FitConfig = FitConfig()
+    fit_config: FitConfig = Field(default_factory=FitConfig)
 
 
 # Sizing metrics are gathered from the evaluation runs and used as input by the calculator.
@@ -103,7 +104,7 @@ class SizingMetrics(BaseModel):
     Sizing metrics for a single concurrency.
     """
     # alerts associated with the sizing metrics
-    alerts: SizingMetricsAlerts = SizingMetricsAlerts()
+    alerts: SizingMetricsAlerts = Field(default_factory=SizingMetricsAlerts)
 
     # p95 LLM latency
     llm_latency_p95: float = 0.0
@@ -125,6 +126,14 @@ class LinearFitResult(BaseModel):
     outliers_removed: list[int]
 
 
+class FitResults(BaseModel):
+    """
+    Linear fit results for both LLM latency and workflow runtime analysis.
+    """
+    llm_latency_fit: LinearFitResult | None = None
+    wf_runtime_fit: LinearFitResult | None = None
+
+
 # GPU estimates are generated by the calculator.
 class GPUEstimates(BaseModel):
     """
@@ -158,11 +167,11 @@ class CalcData(BaseModel):
     """
     # ROUGH GPU estimates per concurrency: these are not used for the final GPU estimation
     # they are only available for information purposes
-    gpu_estimates: GPUEstimates = GPUEstimates()
+    gpu_estimates: GPUEstimates = Field(default_factory=GPUEstimates)
     # Calc runner alerts
-    alerts: CalcAlerts = CalcAlerts()
+    alerts: CalcAlerts = Field(default_factory=CalcAlerts)
     # Sizing metrics
-    sizing_metrics: SizingMetrics = SizingMetrics()
+    sizing_metrics: SizingMetrics = Field(default_factory=SizingMetrics)
 
 
 class CalcRunnerOutput(BaseModel):
@@ -170,7 +179,10 @@ class CalcRunnerOutput(BaseModel):
     Output of the calc runner.
     """
     # GPU estimates based on the slope of the time vs concurrency, calculated online or offline
-    gpu_estimates: GPUEstimates
+    gpu_estimates: GPUEstimates = Field(default_factory=GPUEstimates)
+
+    # Linear fit results for analysis and debugging
+    fit_results: FitResults = Field(default_factory=FitResults)
 
     # Per-concurrency data (GPU estimates, out-of-range runs, and sizing metrics)
     calc_data: dict[int, CalcData] = {}