Coverage for src/dataknobs_fsm/execution/stream.py: 79%

1"""Stream executor for chunk-based processing."""

2

3 import time

4 from dataclasses import dataclass, field

5 from typing import Any, Callable, Dict, List, Tuple, Union

6

7 from dataknobs_fsm.core.fsm import FSM

8 from dataknobs_fsm.core.modes import ProcessingMode, TransactionMode

9 from dataknobs_fsm.execution.context import ExecutionContext

10 from dataknobs_fsm.execution.engine import ExecutionEngine

11 from dataknobs_fsm.streaming.core import (

12 IStreamSink,

13 IStreamSource,

14 StreamChunk,

15 StreamConfig,

16 StreamContext,

17)

18

19

20@dataclass

21 class StreamPipeline:

22 """Pipeline configuration for stream processing."""

23 source: IStreamSource

24 sink: IStreamSink | None = None

25 transformations: List[Callable] = field(default_factory=list)

26 chunk_processors: List[Callable] = field(default_factory=list)

27

28

29@dataclass

30 class StreamProgress:

31 """Progress tracking for stream processing."""

32 chunks_processed: int = 0

33 records_processed: int = 0

34 bytes_processed: int = 0

35 errors: List[Tuple[int, Exception]] = field(default_factory=list)

36 start_time: float = field(default_factory=time.time)

37 last_chunk_time: float = field(default_factory=time.time)

38

39 @property

40 def elapsed_time(self) -> float:

41 """Get elapsed time."""

42 return time.time() - self.start_time

43

44 @property

45 def chunks_per_second(self) -> float:

46 """Get chunk processing rate."""

47 elapsed = self.elapsed_time

48 if elapsed == 0:

49 return 0.0

50 return self.chunks_processed / elapsed

51

52 @property

53 def records_per_second(self) -> float:

54 """Get record processing rate."""

55 elapsed = self.elapsed_time

56 if elapsed == 0:

57 return 0.0

58 return self.records_processed / elapsed

59

60

61 class StreamExecutor:

62 """Executor for stream-based processing.

63

64 This executor handles:

65 - Chunk-based processing with backpressure

66 - Pipeline coordination

67 - Memory management

68 - Progress reporting

69 - Stream transformations

70 """

71

72 def __init__(

73 self,

74 fsm: FSM,

75 stream_config: StreamConfig | None = None,

76 enable_backpressure: bool = True,

77 progress_callback: Union[Callable, None] = None

78 ):

79 """Initialize stream executor.

80

81 Args:

82 fsm: FSM to execute.

83 stream_config: Stream configuration.

84 enable_backpressure: Enable backpressure handling.

85 progress_callback: Callback for progress updates.

86 """

87 self.fsm = fsm

88 self.stream_config = stream_config or StreamConfig()

89 self.enable_backpressure = enable_backpressure

90 self.progress_callback = progress_callback

91

92 # Create execution engine

93 self.engine = ExecutionEngine(fsm)

94

95 # Memory management

96 self._memory_usage = 0

97 self._memory_limit = self.stream_config.memory_limit_mb * 1024 * 1024

98

99 # Backpressure management

100 self._pending_chunks = 0

101 self._backpressure_threshold = self.stream_config.backpressure_threshold

102

103 def execute_stream(

104 self,

105 pipeline: StreamPipeline,

106 context_template: ExecutionContext | None = None,

107 max_transitions: int = 1000

108 ) -> Dict[str, Any]:

109 """Execute stream processing pipeline.

110

111 Args:

112 pipeline: Stream pipeline configuration.

113 context_template: Template context.

114 max_transitions: Maximum transitions per record.

115

116 Returns:

117 Stream processing statistics.

118 """

119 # Create progress tracker

120 progress = StreamProgress()

121

122 # Create base context

123 if context_template is None:

124 # Use SINGLE mode since we process items individually

125 context_template = ExecutionContext(

126 data_mode=ProcessingMode.SINGLE,

127 transaction_mode=TransactionMode.NONE

128 )

129

130 # Create stream context

131 stream_context = StreamContext(config=self.stream_config)

132

133 # Set stream context in execution context

134 context_template.stream_context = stream_context

135

136 # Process stream

137 try:

138 while True:

139 # Check memory usage

140 if self._should_apply_backpressure():

141 time.sleep(0.1)

142 continue

143

144 # Read next chunk from source

145 chunk = pipeline.source.read_chunk()

146 if chunk is None:

147 break

148

149 # Apply chunk processors

150 for processor in pipeline.chunk_processors:

151 chunk = processor(chunk)

152 if chunk is None:

153 break

154

155 if chunk is None:

156 continue

157

158 # Process chunk

159 chunk_results = self._process_chunk(

160 chunk,

161 context_template,

162 pipeline.transformations,

163 max_transitions,

164 progress

165 )

166

167 # Write results to sink if provided

168 if pipeline.sink and chunk_results:

169 result_chunk = StreamChunk(

170 data=chunk_results,

171 sequence_number=chunk.sequence_number,

172 metadata=chunk.metadata,

173 is_last=chunk.is_last

174 )

175 pipeline.sink.write_chunk(result_chunk)

176

177 # Update progress

178 progress.chunks_processed += 1

179 progress.records_processed += len(chunk.data)

180 progress.last_chunk_time = time.time()

181

182 # Fire progress callback

183 if self.progress_callback:

184 self.progress_callback(progress)

185

186 # Check if last chunk

187 if chunk.is_last:

188 break

189

190 finally:

191 # Clean up

192 if hasattr(pipeline.source, 'aclose') or hasattr(pipeline.source, 'close'):

193 pipeline.source.close()

194

195 if pipeline.sink:

196 pipeline.sink.flush()

197 pipeline.sink.close()

198

199 return self._generate_stats(progress)

200

201 def _process_chunk(

202 self,

203 chunk: StreamChunk,

204 context_template: ExecutionContext,

205 transformations: List[Callable],

206 max_transitions: int,

207 progress: StreamProgress

208 ) -> List[Any]:

209 """Process a single chunk.

210

211 Args:

212 chunk: Chunk to process.

213 context_template: Template context.

214 transformations: Transformations to apply.

215 max_transitions: Maximum transitions.

216 progress: Progress tracker.

217

218 Returns:

219 List of processed results.

220 """

221 results = []

222 self._pending_chunks += 1

223

224 try:

225 for i, record in enumerate(chunk.data):

226 # Apply transformations

227 transformed = record

228 for transform in transformations:

229 transformed = transform(transformed)

230 if transformed is None:

231 break

232

233 if transformed is None:

234 continue

235

236 # Create context for this record

237 context = context_template.clone()

238 context.data = transformed

239 context.set_stream_chunk(chunk)

240

241 # Reset to initial state

242 initial_state = self._find_initial_state()

243 if initial_state:

244 context.set_state(initial_state)

245

246 # Execute FSM

247 try:

248 success, result = self.engine.execute(

249 context,

250 transformed,

251 max_transitions

252 )

253

254 if success:

255 results.append(result)

256 else:

257 # FSM failed, but still pass the data through

258 results.append(transformed)

259 progress.errors.append((

260 progress.records_processed + i,

261 Exception(result)

262 ))

263 except Exception as e:

264 # On error, pass the data through

265 results.append(transformed)

266 progress.errors.append((

267 progress.records_processed + i,

268 e

269 ))

270 else:

271 # No FSM configured, just pass data through

272 results.append(transformed)

273

274 finally:

275 self._pending_chunks -= 1

276

277 return results

278

279 def _should_apply_backpressure(self) -> bool:

280 """Check if backpressure should be applied.

281

282 Returns:

283 True if backpressure needed.

284 """

285 if not self.enable_backpressure:

286 return False

287

288 # Check pending chunks

289 if self._pending_chunks >= self._backpressure_threshold:

290 return True

291

292 # Check memory usage

293 if self._memory_usage >= self._memory_limit:

294 return True

295

296 return False

297

298 def _find_initial_state(self) -> str | None:

299 """Find initial state in FSM.

300

301 Returns:

302 Initial state name or None.

303 """

304 # Get main network

305 if self.fsm.name in self.fsm.networks:

306 network = self.fsm.networks[self.fsm.name]

307 if network.initial_states:

308 return next(iter(network.initial_states))

309 return None

310

311 def _generate_stats(self, progress: StreamProgress) -> Dict[str, Any]:

312 """Generate stream processing statistics.

313

314 Args:

315 progress: Progress tracker.

316

317 Returns:

318 Processing statistics.

319 """

320 return {

321 'total_processed': progress.records_processed,

322 'successful': progress.records_processed - len(progress.errors),

323 'failed': len(progress.errors),

324 'duration': progress.elapsed_time,

325 'throughput': progress.records_per_second,

326 # Additional details

327 'chunks_processed': progress.chunks_processed,

328 'bytes_processed': progress.bytes_processed,

329 'error_details': progress.errors[:10] # First 10 errors

330 }

331

332 def create_multi_stage_pipeline(

333 self,

334 stages: List[Dict[str, Any]]

335 ) -> StreamPipeline:

336 """Create a multi-stage processing pipeline.

337

338 Args:

339 stages: List of stage configurations.

340

341 Returns:

342 Configured pipeline.

343 """

344 # Build pipeline from stages

345 transformations = []

346 chunk_processors = []

347

348 for stage in stages:

349 stage_type = stage.get('type')

350

351 if stage_type == 'transform':

352 transformations.append(stage['function'])

353 elif stage_type == 'chunk_processor':

354 chunk_processors.append(stage['function'])

355

356 return StreamPipeline(

357 source=stages[0].get('source'),

358 sink=stages[-1].get('sink'),

359 transformations=transformations,

360 chunk_processors=chunk_processors

361 )