Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/pipeline/executor.py: 16%
75 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Pipeline executor for orchestrating stage execution."""
3from __future__ import annotations
5# to annotate failures and continue per configured error strategies rather than crash the runner
6import asyncio
8from lexigram.ai.rag.pipeline.base import PipelineStageProtocol
9from lexigram.ai.rag.pipeline.types import (
10 ErrorStrategy,
11 PipelineContext,
12 StageStatus,
13)
14from lexigram.logging import (
15 get_logger,
16)
18logger = get_logger(__name__)
21class PipelineExecutor:
22 """Executes pipeline stages in sequence with error handling.
24 The executor orchestrates the execution of pipeline stages,
25 handling errors according to configured strategies and
26 tracking metrics throughout execution.
27 """
29 def __init__(
30 self,
31 stages: list[PipelineStageProtocol],
32 default_error_strategy: ErrorStrategy = ErrorStrategy.GRACEFUL,
33 max_retries: int = 3,
34 retry_delay: float = 1.0,
35 ):
36 """Initialize the pipeline executor.
38 Args:
39 stages: List of pipeline stages to execute
40 default_error_strategy: Default error handling strategy
41 max_retries: Maximum number of retries for RETRY strategy
42 retry_delay: Initial delay between retries (exponential backoff)
43 """
44 self.stages = stages
45 self.default_error_strategy = default_error_strategy
46 self.max_retries = max_retries
47 self.retry_delay = retry_delay
49 async def execute(self, context: PipelineContext) -> PipelineContext:
50 """Execute all pipeline stages in sequence.
52 Args:
53 context: Initial pipeline context
55 Returns:
56 Final pipeline context after all stages
58 Raises:
59 Exception: If a stage fails with FAIL_FAST error strategy
60 """
61 logger.info(
62 "Starting pipeline execution",
63 extra={
64 "request_id": context.request_id,
65 "query": context.query,
66 "num_stages": len(self.stages),
67 },
68 )
70 for stage in self.stages:
71 try:
72 context = await self._execute_stage(stage, context)
73 except Exception as e:
74 logger.exception(
75 "Pipeline execution failed",
76 extra={
77 "request_id": context.request_id,
78 "stage": stage.name,
79 "error": str(e),
80 },
81 )
82 # If we reach here, error strategy was FAIL_FAST
83 context.complete_stage(
84 stage.name,
85 status=StageStatus.FAILED,
86 )
87 raise
89 logger.info(
90 "Pipeline execution completed",
91 extra={
92 "request_id": context.request_id,
93 "duration_ms": context.total_duration_ms,
94 "completed_stages": len(context.completed_stages),
95 "failed_stages": len(context.failed_stages),
96 "has_errors": context.has_errors,
97 },
98 )
100 return context
102 async def _execute_stage(
103 self,
104 stage: PipelineStageProtocol,
105 context: PipelineContext,
106 retry_count: int = 0,
107 ) -> PipelineContext:
108 """Execute a single pipeline stage with error handling.
110 Args:
111 stage: The stage to execute
112 context: Current pipeline context
113 retry_count: Current retry attempt count
115 Returns:
116 Updated pipeline context
118 Raises:
119 Exception: If error strategy is FAIL_FAST
120 """
121 stage_name = stage.name
123 logger.debug(
124 "Starting stage execution",
125 extra={
126 "request_id": context.request_id,
127 "stage": stage_name,
128 "retry_count": retry_count,
129 },
130 )
132 context.start_stage(stage_name)
134 try:
135 # Execute the stage
136 updated_context = await stage.process(context)
138 # Mark as completed
139 updated_context.complete_stage(stage_name, status=StageStatus.COMPLETED)
141 logger.info(
142 "Stage execution completed",
143 extra={
144 "request_id": context.request_id,
145 "stage": stage_name,
146 "duration_ms": [
147 m.duration_ms
148 for m in filter(
149 lambda m: m.stage_name == stage_name,
150 updated_context.stage_metrics,
151 )
152 ][-1],
153 },
154 )
156 except (RuntimeError, ValueError, TypeError, OSError) as e:
157 logger.warning(
158 "Stage execution failed",
159 extra={
160 "request_id": context.request_id,
161 "stage": stage_name,
162 "error": str(e),
163 "retry_count": retry_count,
164 },
165 )
167 # Add error to context
168 context.add_error(
169 stage=stage_name,
170 error=e,
171 recoverable=retry_count < self.max_retries,
172 retry_count=retry_count,
173 )
175 # Apply error strategy
176 return await self._handle_error(stage, context, e, retry_count)
177 else:
178 return updated_context
180 async def _handle_error(
181 self,
182 stage: PipelineStageProtocol,
183 context: PipelineContext,
184 error: Exception,
185 retry_count: int,
186 ) -> PipelineContext:
187 """Handle stage execution error based on error strategy.
189 Args:
190 stage: The stage that failed
191 context: Current pipeline context
192 error: The exception that was raised
193 retry_count: Current retry attempt count
195 Returns:
196 Updated pipeline context
198 Raises:
199 Exception: If error strategy is FAIL_FAST
200 """
201 strategy = self.default_error_strategy
203 if strategy == ErrorStrategy.FAIL_FAST:
204 logger.error(
205 "Failing fast due to stage error",
206 extra={
207 "request_id": context.request_id,
208 "stage": stage.name,
209 "error": str(error),
210 },
211 )
212 context.complete_stage(stage.name, status=StageStatus.FAILED)
213 raise error
215 if strategy == ErrorStrategy.RETRY and retry_count < self.max_retries:
216 # Exponential backoff
217 delay = self.retry_delay * (2**retry_count)
218 logger.info(
219 "Retrying stage after delay",
220 extra={
221 "request_id": context.request_id,
222 "stage": stage.name,
223 "retry_count": retry_count + 1,
224 "delay_seconds": delay,
225 },
226 )
228 await asyncio.sleep(delay)
230 # Complete the failed attempt
231 context.complete_stage(stage.name, status=StageStatus.FAILED)
233 # Retry
234 return await self._execute_stage(stage, context, retry_count + 1)
236 if strategy == ErrorStrategy.SKIP:
237 logger.warning(
238 "Skipping stage due to error",
239 extra={
240 "request_id": context.request_id,
241 "stage": stage.name,
242 "error": str(error),
243 },
244 )
245 context.add_warning(f"Skipped stage {stage.name} due to error: {error}")
246 context.complete_stage(stage.name, status=StageStatus.SKIPPED)
247 return context
249 if strategy == ErrorStrategy.GRACEFUL:
250 logger.warning(
251 "Continuing with partial results due to stage error",
252 extra={
253 "request_id": context.request_id,
254 "stage": stage.name,
255 "error": str(error),
256 },
257 )
258 context.add_warning(
259 f"Stage {stage.name} failed but continuing: {error}",
260 )
261 context.complete_stage(stage.name, status=StageStatus.FAILED)
262 return context
264 # Default to graceful degradation
265 logger.warning(
266 "Unknown error strategy, using graceful degradation",
267 extra={
268 "request_id": context.request_id,
269 "stage": stage.name,
270 "strategy": strategy,
271 },
272 )
273 context.add_warning(
274 f"Stage {stage.name} failed but continuing: {error}",
275 )
276 context.complete_stage(stage.name, status=StageStatus.FAILED)
277 return context
279 async def execute_parallel(
280 self,
281 context: PipelineContext,
282 stages: list[PipelineStageProtocol] | None = None,
283 ) -> PipelineContext:
284 """Execute multiple stages in parallel.
286 Args:
287 context: Pipeline context
288 stages: Stages to execute in parallel (default: all stages)
290 Returns:
291 Updated pipeline context with results from all parallel stages
292 """
293 if stages is None:
294 stages = self.stages
296 logger.info(
297 "Starting parallel stage execution",
298 extra={
299 "request_id": context.request_id,
300 "num_stages": len(stages),
301 },
302 )
304 # Execute all stages in parallel
305 tasks = [self._execute_stage(stage, context) for stage in stages]
306 results = await asyncio.gather(*tasks, return_exceptions=True)
308 # Merge results
309 for i, result in enumerate(results):
310 if isinstance(result, Exception):
311 logger.error(
312 "Parallel stage execution failed",
313 extra={
314 "request_id": context.request_id,
315 "stage": stages[i].name,
316 "error": str(result),
317 },
318 )
319 context.add_error(
320 stage=stages[i].name,
321 error=result,
322 recoverable=False,
323 )
324 elif isinstance(result, PipelineContext):
325 # Merge context results
326 # Note: This is simplified - real merging would be more sophisticated
327 context = result
329 return context