Coverage for llm_dataset_engine/orchestration/pipeline_executor.py: 0%

1"""

2 Pipeline executor for orchestrating stage execution.

3

4 Implements the complete execution flow with state machine management

5 as specified in the design document.

6"""

7

8 from datetime import datetime

9 from enum import Enum

10 from typing import Any, List

11 from uuid import UUID, uuid4

12

13 import pandas as pd

14

15 from llm_dataset_engine.core.models import ExecutionResult

16 from llm_dataset_engine.orchestration.execution_context import ExecutionContext

17 from llm_dataset_engine.orchestration.observers import ExecutionObserver

18 from llm_dataset_engine.orchestration.state_manager import StateManager

19 from llm_dataset_engine.stages.pipeline_stage import PipelineStage

20 from llm_dataset_engine.utils import get_logger

21

22 logger = get_logger(__name__)

23

24

25 class ExecutionState(str, Enum):

26 """Pipeline execution states."""

27

28 IDLE = "idle"

29 INITIALIZING = "initializing"

30 EXECUTING = "executing"

31 PAUSED = "paused"

32 COMPLETED = "completed"

33 FAILED = "failed"

34

35

36 class PipelineExecutor:

37 """

38 Orchestrates pipeline execution with state management.

39

40 Implements Command and Mediator patterns for coordinating

41 stages, observers, and state management.

42

43 State Machine:

44 IDLE → INITIALIZING → EXECUTING → [PAUSED ↔ EXECUTING] → COMPLETED

45 ↓

46 FAILED

47 """

48

49 def __init__(

50 self,

51 stages: List[PipelineStage],

52 state_manager: StateManager,

53 observers: List[ExecutionObserver] | None = None,

54 ):

55 """

56 Initialize pipeline executor.

57

58 Args:

59 stages: Ordered list of processing stages

60 state_manager: State manager for checkpointing

61 observers: Optional execution observers

62 """

63 self.execution_id = uuid4()

64 self.stages = stages

65 self.state_manager = state_manager

66 self.observers = observers or []

67 self.state = ExecutionState.IDLE

68 self.context: ExecutionContext | None = None

69 self.logger = get_logger(f"{__name__}.{self.execution_id}")

70

71 def add_observer(self, observer: ExecutionObserver) -> "PipelineExecutor":

72 """

73 Add execution observer.

74

75 Args:

76 observer: Observer to add

77

78 Returns:

79 Self for chaining

80 """

81 self.observers.append(observer)

82 return self

83

84 def execute(self, pipeline: Any) -> ExecutionResult:

85 """

86 Execute pipeline end-to-end.

87

88 Args:

89 pipeline: Pipeline instance to execute

90

91 Returns:

92 ExecutionResult with data and metrics

93

94 Raises:

95 RuntimeError: If pipeline in invalid state

96 """

97 if self.state not in [ExecutionState.IDLE, ExecutionState.FAILED]:

98 raise RuntimeError(

99 f"Cannot execute from state: {self.state}"

100 )

101

102 try:

103 # Initialize

104 self.state = ExecutionState.INITIALIZING

105 self.context = self._initialize_context()

106

107 # Check for existing checkpoint

108 if self.state_manager.can_resume(self.context.session_id):

109 self.logger.info("Found existing checkpoint, resuming...")

110 self.context = self.state_manager.load_checkpoint(

111 self.context.session_id

112 )

113

114 # Notify observers

115 self._notify_pipeline_start(pipeline)

116

117 # Execute stages

118 self.state = ExecutionState.EXECUTING

119 result_data = self._execute_all_stages(pipeline)

120

121 # Mark completion

122 self.state = ExecutionState.COMPLETED

123 self.context.end_time = datetime.now()

124

125 # Create result

126 result = self._create_execution_result(result_data)

127

128 # Cleanup checkpoints

129 self.state_manager.cleanup_checkpoints(self.context.session_id)

130

131 # Notify observers

132 self._notify_pipeline_complete(result)

133

134 return result

135

136 except Exception as e:

137 self.state = ExecutionState.FAILED

138 self._notify_pipeline_error(e)

139

140 # Save checkpoint on failure

141 if self.context:

142 self.state_manager.save_checkpoint(self.context)

143

144 raise

145

146 def pause(self) -> None:

147 """

148 Gracefully pause execution.

149

150 Finishes current batch and saves checkpoint.

151 """

152 if self.state != ExecutionState.EXECUTING:

153 raise RuntimeError(

154 f"Cannot pause from state: {self.state}"

155 )

156

157 self.logger.info("Pausing execution...")

158 self.state = ExecutionState.PAUSED

159

160 # Save checkpoint

161 if self.context:

162 self.state_manager.save_checkpoint(self.context)

163

164 def resume(self, session_id: UUID) -> ExecutionResult:

165 """

166 Resume from saved checkpoint.

167

168 Args:

169 session_id: Session ID to resume

170

171 Returns:

172 ExecutionResult

173

174 Raises:

175 ValueError: If no checkpoint found

176 """

177 if not self.state_manager.can_resume(session_id):

178 raise ValueError(f"No checkpoint found for session {session_id}")

179

180 self.context = self.state_manager.load_checkpoint(session_id)

181 if not self.context:

182 raise ValueError("Failed to load checkpoint")

183

184 self.logger.info(

185 f"Resuming from row {self.context.last_processed_row}"

186 )

187

188 # Continue execution

189 # Note: Would need to reconstruct pipeline and skip processed rows

190 raise NotImplementedError("Resume functionality coming soon")

191

192 def cancel(self) -> None:

193 """

194 Immediately stop and save checkpoint.

195 """

196 self.logger.info("Cancelling execution...")

197

198 # Save checkpoint

199 if self.context:

200 self.state_manager.save_checkpoint(self.context)

201

202 self.state = ExecutionState.IDLE

203

204 def _initialize_context(self) -> ExecutionContext:

205 """Initialize execution context."""

206 return ExecutionContext(session_id=self.execution_id)

207

208 def _execute_all_stages(self, pipeline: Any) -> pd.DataFrame:

209 """

210 Execute all pipeline stages sequentially.

211

212 Args:

213 pipeline: Pipeline instance

214

215 Returns:

216 Final DataFrame result

217 """

218 # This will be implemented with actual stage orchestration

219 # For now, delegate to pipeline's execution logic

220 # In a future iteration, we'll move all execution here

221 raise NotImplementedError(

222 "Stage orchestration implemented in Pipeline.execute() currently"

223 )

224

225 def _create_execution_result(

226 self, data: pd.DataFrame

227 ) -> ExecutionResult:

228 """

229 Create execution result from context and data.

230

231 Args:

232 data: Final processed data

233

234 Returns:

235 ExecutionResult

236 """

237 if not self.context:

238 raise RuntimeError("No execution context available")

239

240 return ExecutionResult(

241 data=data,

242 metrics=self.context.get_stats(),

243 costs=self.context.accumulated_cost,

244 execution_id=self.context.session_id,

245 start_time=self.context.start_time,

246 end_time=self.context.end_time,

247 success=True,

248 )

249

250 def _notify_pipeline_start(self, pipeline: Any) -> None:

251 """Notify all observers of pipeline start."""

252 if self.context:

253 for observer in self.observers:

254 try:

255 observer.on_pipeline_start(pipeline, self.context)

256 except Exception as e:

257 self.logger.error(

258 f"Observer {observer.__class__.__name__} failed: {e}"

259 )

260

261 def _notify_pipeline_complete(self, result: ExecutionResult) -> None:

262 """Notify all observers of pipeline completion."""

263 if self.context:

264 for observer in self.observers:

265 try:

266 observer.on_pipeline_complete(self.context, result)

267 except Exception as e:

268 self.logger.error(

269 f"Observer {observer.__class__.__name__} failed: {e}"

270 )

271

272 def _notify_pipeline_error(self, error: Exception) -> None:

273 """Notify all observers of pipeline error."""

274 if self.context:

275 for observer in self.observers:

276 try:

277 observer.on_pipeline_error(self.context, error)

278 except Exception as e:

279 self.logger.error(

280 f"Observer {observer.__class__.__name__} failed: {e}"

281 )

282