File indexing completed on 2026-08-12 08:24:57
0001 """Execution logging and checkpointing utilities for workflow jobs.
0002
0003 Provides comprehensive logging with checkpoints, allowing jobs to be traced
0004 and debugged with detailed context at each execution stage.
0005
0006 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0007 Homepage: https://aid2e.github.io/aid2e-framework
0008 Repository: https://github.com/aid2e/AID2E-framework.git
0009 """
0010
0011 import json
0012 import logging
0013 import sys
0014 from datetime import datetime
0015 from pathlib import Path
0016 from typing import Any, Dict, Optional
0017 from dataclasses import dataclass, asdict
0018
0019
0020 @dataclass
0021 class Checkpoint:
0022 """Represents a checkpoint in job execution.
0023
0024 Attributes:
0025 stage: Execution stage name (e.g., "rule_resolution", "payload_validation")
0026 status: Status code ("start", "success", "warning", "error", "skipped")
0027 timestamp: ISO format timestamp
0028 message: Human-readable message
0029 context: Contextual data (job_id, payload, resolved values, etc.)
0030 details: Additional details (error messages, stack traces, etc.)
0031 """
0032 stage: str
0033 status: str
0034 timestamp: str
0035 message: str
0036 context: Dict[str, Any]
0037 details: Optional[Dict[str, Any]] = None
0038
0039 def to_dict(self) -> Dict[str, Any]:
0040 """Convert checkpoint to dict for JSON serialization."""
0041 return asdict(self)
0042
0043
0044 class ExecutionLogger:
0045 """Comprehensive logger for job execution with checkpoints.
0046
0047 Features:
0048 - Structured logging (JSON-compatible)
0049 - Checkpoint tracking at each execution stage
0050 - Context preservation across stages
0051 - File-based logging with rotation
0052 - Console logging with color formatting
0053
0054 Example:
0055 >>> logger = ExecutionLogger(
0056 ... job_name="dtlz2_evaluate",
0057 ... output_dir="/tmp/stage_output",
0058 ... log_level="DEBUG"
0059 ... )
0060 >>> logger.checkpoint("rule_resolution", "start", "Resolving rule template")
0061 >>> logger.log_info("Rule resolved successfully")
0062 >>> logger.checkpoint("rule_resolution", "success", "Rule resolved")
0063 >>> checkpoint = logger.get_last_checkpoint()
0064 """
0065
0066 def __init__(
0067 self,
0068 job_name: str,
0069 output_dir: str,
0070 log_level: str = "INFO",
0071 enable_file_logging: bool = True,
0072 enable_checkpoint_file: bool = True
0073 ):
0074 """Initialize execution logger.
0075
0076 Args:
0077 job_name: Name of the job being executed
0078 output_dir: Directory for log files
0079 log_level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
0080 enable_file_logging: Whether to log to file
0081 enable_checkpoint_file: Whether to save checkpoints to JSON file
0082 """
0083 self.job_name = job_name
0084 self.output_dir = Path(output_dir)
0085 self.log_level = getattr(logging, log_level.upper(), logging.INFO)
0086 self.enable_file_logging = enable_file_logging
0087 self.enable_checkpoint_file = enable_checkpoint_file
0088
0089
0090 self.output_dir.mkdir(parents=True, exist_ok=True)
0091
0092
0093 self.logger = logging.getLogger(f"aid2e.{job_name}")
0094 self.logger.setLevel(self.log_level)
0095 self.logger.handlers.clear()
0096
0097
0098 console_handler = logging.StreamHandler(sys.stdout)
0099 console_handler.setLevel(self.log_level)
0100 formatter = logging.Formatter(
0101 "%(asctime)s [%(name)s] %(levelname)s: %(message)s",
0102 datefmt="%Y-%m-%d %H:%M:%S"
0103 )
0104 console_handler.setFormatter(formatter)
0105 self.logger.addHandler(console_handler)
0106
0107
0108 if enable_file_logging:
0109 log_file = self.output_dir / f"{job_name}_execution.log"
0110 file_handler = logging.FileHandler(log_file)
0111 file_handler.setLevel(self.log_level)
0112 file_handler.setFormatter(formatter)
0113 self.logger.addHandler(file_handler)
0114 self.log_file = log_file
0115 else:
0116 self.log_file = None
0117
0118
0119 self.checkpoints: list[Checkpoint] = []
0120 self.checkpoint_file = self.output_dir / f"{job_name}_checkpoints.json"
0121
0122
0123 self.checkpoint(
0124 stage="initialization",
0125 status="start",
0126 message=f"Job execution started: {job_name}",
0127 context={"job_name": job_name, "output_dir": str(self.output_dir)}
0128 )
0129
0130 def checkpoint(
0131 self,
0132 stage: str,
0133 status: str,
0134 message: str,
0135 context: Optional[Dict[str, Any]] = None,
0136 details: Optional[Dict[str, Any]] = None
0137 ) -> Checkpoint:
0138 """Record a checkpoint at current execution stage.
0139
0140 Args:
0141 stage: Stage name (e.g., "rule_resolution", "payload_validation")
0142 status: Status code ("start", "success", "warning", "error", "skipped")
0143 message: Human-readable message
0144 context: Contextual data (optional)
0145 details: Additional details (optional)
0146
0147 Returns:
0148 Checkpoint object created
0149 """
0150 checkpoint = Checkpoint(
0151 stage=stage,
0152 status=status,
0153 timestamp=datetime.now().isoformat(),
0154 message=message,
0155 context=context or {},
0156 details=details
0157 )
0158
0159 self.checkpoints.append(checkpoint)
0160
0161
0162 log_msg = f"[{stage}:{status}] {message}"
0163 if context:
0164 log_msg += f" | context: {json.dumps(context, default=str, indent=0)}"
0165
0166 if status == "error":
0167 self.logger.error(log_msg)
0168 elif status == "warning":
0169 self.logger.warning(log_msg)
0170 elif status == "skipped":
0171 self.logger.info(f"[SKIP] {log_msg}")
0172 else:
0173 self.logger.info(log_msg)
0174
0175
0176 if self.enable_checkpoint_file:
0177 self._save_checkpoints()
0178
0179 return checkpoint
0180
0181 def log_debug(self, message: str, context: Optional[Dict[str, Any]] = None):
0182 """Log debug message with optional context."""
0183 if context:
0184 message += f" | {json.dumps(context, default=str, indent=0)}"
0185 self.logger.debug(message)
0186
0187 def log_info(self, message: str, context: Optional[Dict[str, Any]] = None):
0188 """Log info message with optional context."""
0189 if context:
0190 message += f" | {json.dumps(context, default=str, indent=0)}"
0191 self.logger.info(message)
0192
0193 def log_warning(self, message: str, context: Optional[Dict[str, Any]] = None):
0194 """Log warning message with optional context."""
0195 if context:
0196 message += f" | {json.dumps(context, default=str, indent=0)}"
0197 self.logger.warning(message)
0198
0199 def log_error(self, message: str, context: Optional[Dict[str, Any]] = None):
0200 """Log error message with optional context."""
0201 if context:
0202 message += f" | {json.dumps(context, default=str, indent=0)}"
0203 self.logger.error(message)
0204
0205 def get_last_checkpoint(self) -> Optional[Checkpoint]:
0206 """Get the most recent checkpoint."""
0207 return self.checkpoints[-1] if self.checkpoints else None
0208
0209 def get_checkpoint_by_stage(self, stage: str) -> list[Checkpoint]:
0210 """Get all checkpoints for a given stage."""
0211 return [cp for cp in self.checkpoints if cp.stage == stage]
0212
0213 def get_checkpoints_by_status(self, status: str) -> list[Checkpoint]:
0214 """Get all checkpoints with a given status."""
0215 return [cp for cp in self.checkpoints if cp.status == status]
0216
0217 def _save_checkpoints(self):
0218 """Save all checkpoints to JSON file."""
0219 try:
0220 checkpoint_data = {
0221 "job_name": self.job_name,
0222 "total_checkpoints": len(self.checkpoints),
0223 "execution_start": self.checkpoints[0].timestamp if self.checkpoints else None,
0224 "execution_end": self.checkpoints[-1].timestamp if self.checkpoints else None,
0225 "checkpoints": [cp.to_dict() for cp in self.checkpoints]
0226 }
0227
0228 with open(self.checkpoint_file, "w") as f:
0229 json.dump(checkpoint_data, f, indent=2, default=str)
0230 except Exception as e:
0231 self.logger.error(f"Failed to save checkpoints: {e}")
0232
0233 def execution_summary(self) -> Dict[str, Any]:
0234 """Generate execution summary from checkpoints.
0235
0236 Returns:
0237 Dict with execution summary (total stages, errors, warnings, etc.)
0238 """
0239 summary = {
0240 "job_name": self.job_name,
0241 "total_checkpoints": len(self.checkpoints),
0242 "execution_start": self.checkpoints[0].timestamp if self.checkpoints else None,
0243 "execution_end": self.checkpoints[-1].timestamp if self.checkpoints else None,
0244 "status_breakdown": {
0245 "start": len(self.get_checkpoints_by_status("start")),
0246 "success": len(self.get_checkpoints_by_status("success")),
0247 "warning": len(self.get_checkpoints_by_status("warning")),
0248 "error": len(self.get_checkpoints_by_status("error")),
0249 "skipped": len(self.get_checkpoints_by_status("skipped"))
0250 },
0251 "stages_executed": list(set(cp.stage for cp in self.checkpoints)),
0252 "has_errors": len(self.get_checkpoints_by_status("error")) > 0
0253 }
0254 return summary
0255
0256
0257 def create_job_logger(
0258 job_name: str,
0259 output_dir: str,
0260 log_level: str = "INFO"
0261 ) -> ExecutionLogger:
0262 """Convenience function to create job logger.
0263
0264 Args:
0265 job_name: Name of the job
0266 output_dir: Output directory for logs
0267 log_level: Logging level
0268
0269 Returns:
0270 ExecutionLogger instance
0271 """
0272 return ExecutionLogger(
0273 job_name=job_name,
0274 output_dir=output_dir,
0275 log_level=log_level,
0276 enable_file_logging=True,
0277 enable_checkpoint_file=True
0278 )
