Coverage for llm_dataset_engine/orchestration/observers.py: 52%

1"""

2 Execution observers for monitoring and logging.

3

4 Implements Observer pattern for decoupled event notification.

5"""

6

7 from abc import ABC, abstractmethod

8 from typing import Any

9

10 from tqdm import tqdm

11

12 from llm_dataset_engine.core.models import ExecutionResult

13 from llm_dataset_engine.orchestration.execution_context import ExecutionContext

14 from llm_dataset_engine.stages.pipeline_stage import PipelineStage

15 from llm_dataset_engine.utils import get_logger

16

17

18 class ExecutionObserver(ABC):

19 """

20 Abstract base for execution observers.

21

22 Observers can monitor pipeline execution without coupling

23 to the executor implementation.

24 """

25

26 @abstractmethod

27 def on_pipeline_start(

28 self, pipeline: Any, context: ExecutionContext

29 ) -> None:

30 """Called before first stage execution."""

31 pass

32

33 @abstractmethod

34 def on_stage_start(

35 self, stage: PipelineStage, context: ExecutionContext

36 ) -> None:

37 """Called before each stage."""

38 pass

39

40 @abstractmethod

41 def on_stage_complete(

42 self, stage: PipelineStage, context: ExecutionContext, result: Any

43 ) -> None:

44 """Called after successful stage completion."""

45 pass

46

47 @abstractmethod

48 def on_stage_error(

49 self, stage: PipelineStage, context: ExecutionContext, error: Exception

50 ) -> None:

51 """Called on stage failure."""

52 pass

53

54 @abstractmethod

55 def on_pipeline_complete(

56 self, context: ExecutionContext, result: ExecutionResult

57 ) -> None:

58 """Called after all stages complete."""

59 pass

60

61 @abstractmethod

62 def on_pipeline_error(

63 self, context: ExecutionContext, error: Exception

64 ) -> None:

65 """Called on fatal pipeline failure."""

66 pass

67

68

69 class ProgressBarObserver(ExecutionObserver):

70 """Observer that displays progress bar with tqdm."""

71

72 def __init__(self):

73 """Initialize progress bar observer."""

74 self.progress_bar: tqdm | None = None

75

76 def on_pipeline_start(

77 self, pipeline: Any, context: ExecutionContext

78 ) -> None:

79 """Initialize progress bar."""

80 if context.total_rows > 0:

81 self.progress_bar = tqdm(

82 total=context.total_rows,

83 desc="Processing",

84 unit="rows",

85 )

86

87 def on_stage_start(

88 self, stage: PipelineStage, context: ExecutionContext

89 ) -> None:

90 """Update progress bar description."""

91 if self.progress_bar:

92 self.progress_bar.set_description(f"Stage: {stage.name}")

93

94 def on_stage_complete(

95 self, stage: PipelineStage, context: ExecutionContext, result: Any

96 ) -> None:

97 """Update progress bar."""

98 if self.progress_bar:

99 self.progress_bar.n = context.last_processed_row

100 self.progress_bar.refresh()

101

102 def on_stage_error(

103 self, stage: PipelineStage, context: ExecutionContext, error: Exception

104 ) -> None:

105 """Handle error in progress bar."""

106 if self.progress_bar:

107 self.progress_bar.set_description(f"Error in {stage.name}")

108

109 def on_pipeline_complete(

110 self, context: ExecutionContext, result: ExecutionResult

111 ) -> None:

112 """Close progress bar."""

113 if self.progress_bar:

114 self.progress_bar.close()

115 self.progress_bar = None

116

117 def on_pipeline_error(

118 self, context: ExecutionContext, error: Exception

119 ) -> None:

120 """Close progress bar on error."""

121 if self.progress_bar:

122 self.progress_bar.close()

123 self.progress_bar = None

124

125

126 class LoggingObserver(ExecutionObserver):

127 """Observer that logs execution events."""

128

129 def __init__(self):

130 """Initialize logging observer."""

131 self.logger = get_logger(__name__)

132

133 def on_pipeline_start(

134 self, pipeline: Any, context: ExecutionContext

135 ) -> None:

136 """Log pipeline start."""

137 self.logger.info(

138 f"Pipeline execution started (session: {context.session_id})"

139 )

140

141 def on_stage_start(

142 self, stage: PipelineStage, context: ExecutionContext

143 ) -> None:

144 """Log stage start."""

145 self.logger.info(

146 f"Starting stage: {stage.name} "

147 f"(progress: {context.get_progress():.1f}%)"

148 )

149

150 def on_stage_complete(

151 self, stage: PipelineStage, context: ExecutionContext, result: Any

152 ) -> None:

153 """Log stage completion."""

154 self.logger.info(

155 f"Completed stage: {stage.name} "

156 f"(cost: ${context.accumulated_cost:.4f})"

157 )

158

159 def on_stage_error(

160 self, stage: PipelineStage, context: ExecutionContext, error: Exception

161 ) -> None:

162 """Log stage error."""

163 self.logger.error(f"Stage {stage.name} failed: {error}")

164

165 def on_pipeline_complete(

166 self, context: ExecutionContext, result: ExecutionResult

167 ) -> None:

168 """Log pipeline completion."""

169 self.logger.info(

170 f"Pipeline execution completed successfully\n"

171 f" Processed: {result.metrics.processed_rows} rows\n"

172 f" Duration: {result.metrics.total_duration_seconds:.2f}s\n"

173 f" Total cost: ${result.costs.total_cost:.4f}\n"

174 f" Errors: {result.metrics.failed_rows}"

175 )

176

177 def on_pipeline_error(

178 self, context: ExecutionContext, error: Exception

179 ) -> None:

180 """Log pipeline error."""

181 self.logger.error(f"Pipeline execution failed: {error}")

182

183

184 class CostTrackingObserver(ExecutionObserver):

185 """Observer that tracks and warns about costs."""

186

187 def __init__(self, warning_threshold: float = 0.75):

188 """

189 Initialize cost tracking observer.

190

191 Args:

192 warning_threshold: Warn when this fraction of budget used

193 """

194 self.logger = get_logger(__name__)

195 self.warning_threshold = warning_threshold

196 self.max_budget: float | None = None

197

198 def on_pipeline_start(

199 self, pipeline: Any, context: ExecutionContext

200 ) -> None:

201 """Set max budget if available."""

202 # Could extract from pipeline specs

203 pass

204

205 def on_stage_start(

206 self, stage: PipelineStage, context: ExecutionContext

207 ) -> None:

208 """No action on stage start."""

209 pass

210

211 def on_stage_complete(

212 self, stage: PipelineStage, context: ExecutionContext, result: Any

213 ) -> None:

214 """Check cost after stage completion."""

215 if self.max_budget:

216 usage_ratio = float(context.accumulated_cost) / self.max_budget

217

218 if usage_ratio >= self.warning_threshold:

219 self.logger.warning(

220 f"Cost warning: {usage_ratio * 100:.1f}% of budget used "

221 f"(${context.accumulated_cost:.4f} / ${self.max_budget:.2f})"

222 )

223

224 def on_stage_error(

225 self, stage: PipelineStage, context: ExecutionContext, error: Exception

226 ) -> None:

227 """No action on error."""

228 pass

229

230 def on_pipeline_complete(

231 self, context: ExecutionContext, result: ExecutionResult

232 ) -> None:

233 """Log final cost summary."""

234 self.logger.info(

235 f"Cost summary:\n"

236 f" Total: ${result.costs.total_cost:.4f}\n"

237 f" Input tokens: {result.costs.input_tokens:,}\n"

238 f" Output tokens: {result.costs.output_tokens:,}\n"

239 f" Cost per row: ${float(result.costs.total_cost) / result.metrics.total_rows:.6f}"

240 )

241

242 def on_pipeline_error(

243 self, context: ExecutionContext, error: Exception

244 ) -> None:

245 """Log cost at failure."""

246 self.logger.info(

247 f"Cost at failure: ${context.accumulated_cost:.4f}"

248 )

249