utilities/workflows/rule_resolution.py

0001 """Rule resolution and payload validation for job execution.
0002
0003 Implements template-based command construction and payload
0004 validation for conditional job execution.
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 re
0012 from typing import Any, Dict, Optional, Tuple
0013 from aid2e.utilities.configurations.workflow_config import JobDefinition
0014 from aid2e.utilities.workflows.execution_logger import ExecutionLogger
0015
0016
0017 class RuleResolutionError(Exception):
0018     """Raised when rule resolution fails."""
0019     pass
0020
0021
0022 class PayloadValidationError(Exception):
0023     """Raised when payload validation fails."""
0024     pass
0025
0026
0027 def resolve_payload_templates(
0028     payload: Dict[str, Any],
0029     context: Dict[str, Any],
0030     logger: Optional[ExecutionLogger] = None
0031 ) -> Dict[str, Any]:
0032     """Resolve template variables in payload dict.
0033
0034     Recursively processes payload dict, replacing template variables with
0035     context values. Supports nested dicts and lists.
0036
0037     Template Variables:
0038     - {{job_id}}: Job name
0039     - {{output_dir}}: Job output directory
0040     - {{stage_outputs[stage_name]}}: Output path from previous stage
0041     - {{input_design_params}}: Design parameters file path
0042
0043     Args:
0044         payload: Payload dict potentially containing template variables
0045         context: Runtime context dict
0046         logger: Optional execution logger for logging operations
0047
0048     Returns:
0049         Dict with all templates resolved
0050
0051     Raises:
0052         RuleResolutionError: If template variable is undefined in context
0053
0054     Example:
0055         >>> payload = {
0056         ...     "input_file": "{{input_design_params}}",
0057         ...     "output_dir": "{{output_dir}}",
0058         ...     "job_id": "{{job_id}}"
0059         ... }
0060         >>> context = {
0061         ...     "input_design_params": "/path/to/design.params",
0062         ...     "output_dir": "/tmp/stage_output",
0063         ...     "job_id": 0
0064         ... }
0065         >>> resolved = resolve_payload_templates(payload, context)
0066         # resolved = {"input_file": "/path/to/design.params", ...}
0067     """
0068     if logger:
0069         logger.checkpoint(
0070             stage="payload_template_resolution",
0071             status="start",
0072             message="Resolving payload template variables",
0073             context={"payload_keys": list(payload.keys())}
0074         )
0075
0076     resolved = {}
0077
0078     for key, value in payload.items():
0079         if isinstance(value, str):
0080             # Resolve string templates
0081             resolved[key] = _resolve_template_string(value, context, logger)
0082         elif isinstance(value, dict):
0083             # Recursively resolve nested dicts
0084             resolved[key] = resolve_payload_templates(value, context, logger)
0085         elif isinstance(value, list):
0086             # Resolve list items
0087             resolved[key] = [
0088                 _resolve_template_string(item, context, logger) if isinstance(item, str) else item
0089                 for item in value
0090             ]
0091         else:
0092             # Non-string values (int, bool, float, etc.) pass through unchanged
0093             resolved[key] = value
0094
0095     if logger:
0096         logger.checkpoint(
0097             stage="payload_template_resolution",
0098             status="success",
0099             message="Payload templates resolved successfully",
0100             context={
0101                 "resolved_keys": list(resolved.keys()),
0102                 "sample_values": {k: str(v)[:50] for k, v in list(resolved.items())[:3]}
0103             }
0104         )
0105
0106     return resolved
0107
0108
0109 def _resolve_template_string(
0110     template: str,
0111     context: Dict[str, Any],
0112     logger: Optional[ExecutionLogger] = None
0113 ) -> str:
0114     """Resolve a single template string.
0115
0116     Args:
0117         template: Template string (e.g., "{{input_design_params}}")
0118         context: Runtime context
0119         logger: Optional logger
0120
0121     Returns:
0122         Resolved string
0123
0124     Raises:
0125         RuleResolutionError: If variable is undefined
0126     """
0127     # Find all template variables: {{var_name}} or {{dict[key]}}
0128     pattern = r'\{\{([^}]+)\}\}'
0129     variables = re.findall(pattern, template)
0130
0131     result = template
0132     for var in variables:
0133         # Try to get value from context
0134         target = "{{" + f"{var}" + "}}"
0135         value = _get_context_value(var, context)
0136         if value is None:
0137             raise RuleResolutionError(f"Undefined template variable: {target}")
0138
0139         result = result.replace(target, str(value))
0140
0141     return result
0142
0143
0144 def _get_context_value(var_path: str, context: Dict[str, Any]) -> Any:
0145     """Get value from context using path notation.
0146
0147     Supports:
0148     - Simple keys: "job_id"
0149     - Dict access: "stage_outputs[preparation]"
0150     - Nested access: "payload[metadata][version]"
0151
0152     Args:
0153         var_path: Variable path string
0154         context: Context dict
0155
0156     Returns:
0157         Value from context, or None if not found
0158     """
0159     # Handle dict[key] notation
0160     if "[" in var_path and "]" in var_path:
0161         # Extract base name and keys
0162         match = re.match(r'(\w+)\[([^\]]+)\]', var_path)
0163         if not match:
0164             return None
0165
0166         base_name = match.group(1)
0167         key_path = match.group(2)
0168
0169         if base_name not in context:
0170             return None
0171
0172         obj = context[base_name]
0173
0174         # Navigate through keys
0175         for key in key_path.split("]["):
0176             key = key.strip("]").strip("[")
0177             if isinstance(obj, dict):
0178                 obj = obj.get(key)
0179             else:
0180                 return None
0181
0182             if obj is None:
0183                 return None
0184
0185         return obj
0186     else:
0187         # Simple key lookup
0188         return context.get(var_path)
0189
0190
0191 def resolve_job_rule(
0192     job: JobDefinition,
0193     context: Dict[str, Any],
0194     logger: Optional[ExecutionLogger] = None
0195 ) -> str:
0196     """Resolve job rule template to final command.
0197
0198     Implements experimental_stack.py StackLayer pattern for command construction.
0199
0200     Rule Template Variables:
0201     - {{command}}: Job command
0202     - {{payload[key]}}: Value from resolved payload
0203     - {{job_id}}, {{output_dir}}, etc.: Context variables
0204
0205     Args:
0206         job: JobDefinition with command and optional rule
0207         context: Runtime context with template variables
0208         logger: Optional execution logger
0209
0210     Returns:
0211         Final command string ready for execution
0212
0213     Raises:
0214         RuleResolutionError: If rule cannot be resolved
0215
0216     Example:
0217         >>> job = JobDefinition(
0218         ...     name="test",
0219         ...     command="python compute.py",
0220         ...     rule="{{command}} {{payload[input]}} {{payload[output]}}",
0221         ...     payload={"input": "{{input_design_params}}", "output": "{{output_dir}}"}
0222         ... )
0223         >>> context = {
0224         ...     "input_design_params": "/data/design.params",
0225         ...     "output_dir": "/tmp/out",
0226         ...     "job_id": 0
0227         ... }
0228         >>> cmd = resolve_job_rule(job, context)
0229         # cmd = "python compute.py /data/design.params /tmp/out"
0230     """
0231     if logger:
0232         logger.checkpoint(
0233             stage="rule_resolution",
0234             status="start",
0235             message="Resolving job rule template",
0236             context={"job_name": job.name, "rule": job.rule}
0237         )
0238
0239     # Step 1: Resolve payload templates with context
0240     resolved_payload = resolve_payload_templates(job.payload, context, logger)
0241
0242     # Step 2: Determine rule (use default if not specified)
0243     rule = job.rule or "{{command}}"
0244
0245     if logger:
0246         logger.log_debug("Using rule template", context={"rule": rule})
0247
0248     # Step 3: Build command by substituting rule template
0249     # Create substitution dict: {"command": value, "payload": {...}}
0250     substitution_context = {
0251         "command": job.command,
0252         "payload": resolved_payload,
0253         **context  # Include all context variables
0254     }
0255
0256     try:
0257         final_command = _substitute_rule_template(rule, substitution_context)
0258
0259         if logger:
0260             logger.checkpoint(
0261                 stage="rule_resolution",
0262                 status="success",
0263                 message="Rule template resolved to final command",
0264                 context={
0265                     "rule": rule,
0266                     "final_command": final_command,
0267                     "command_length": len(final_command)
0268                 }
0269             )
0270
0271         return final_command
0272
0273     except Exception as e:
0274         if logger:
0275             logger.checkpoint(
0276                 stage="rule_resolution",
0277                 status="error",
0278                 message=f"Failed to resolve rule: {str(e)}",
0279                 context={"rule": rule},
0280                 details={"error": str(e), "job_name": job.name}
0281             )
0282         raise RuleResolutionError(f"Cannot resolve rule '{rule}': {str(e)}")
0283
0284
0285 def _substitute_rule_template(rule: str, context: Dict[str, Any]) -> str:
0286     """Substitute variables in rule template.
0287
0288     Handles {{command}}, {{payload[key]}}, and other context variables.
0289
0290     Args:
0291         rule: Rule template string
0292         context: Substitution context
0293
0294     Returns:
0295         Substituted command string
0296     """
0297     result = rule
0298
0299     # Find all template variables: {var_name} or {dict[key]}
0300     pattern = r'\{\{([^}]+)\}\}'
0301     variables = re.findall(pattern, rule)
0302
0303     for var in variables:
0304         target = "{{" + f"{var}" + "}}"
0305         value = _get_context_value(var, context)
0306         if value is None:
0307             raise ValueError(f"Undefined variable in rule: {target}")
0308
0309         result = result.replace(target, str(value))
0310
0311     # Clean up multiple spaces
0312     result = re.sub(r'\s+', ' ', result).strip()
0313
0314     return result
0315
0316
0317 def validate_job_payload(
0318     job: JobDefinition,
0319     required_keys: Optional[list[str]] = None,
0320     context: Optional[Dict[str, Any]] = None,
0321     logger: Optional[ExecutionLogger] = None
0322 ) -> Tuple[bool, Optional[str]]:
0323     """Validate job payload before execution.
0324
0325     Checks:
0326     1. Required keys are present in payload
0327     2. Template variables can be resolved (if context provided)
0328     3. Payload values are not empty
0329
0330     Args:
0331         job: JobDefinition to validate
0332         required_keys: List of required payload keys
0333         context: Optional context for template validation
0334         logger: Optional execution logger
0335
0336     Returns:
0337         Tuple of (is_valid, error_message)
0338
0339     Example:
0340         >>> job = JobDefinition(
0341         ...     name="test",
0342         ...     command="python script.py",
0343         ...     payload={"input": "/path/to/input"}
0344         ... )
0345         >>> is_valid, error = validate_job_payload(job, required_keys=["input"])
0346         # is_valid = True, error = None
0347     """
0348     if logger:
0349         logger.checkpoint(
0350             stage="payload_validation",
0351             status="start",
0352             message="Validating job payload",
0353             context={"job_name": job.name, "required_keys": required_keys}
0354         )
0355
0356     # Check required keys
0357     if required_keys:
0358         missing_keys = [key for key in required_keys if key not in job.payload]
0359         if missing_keys:
0360             error_msg = f"Missing required payload keys: {missing_keys}"
0361             if logger:
0362                 logger.checkpoint(
0363                     stage="payload_validation",
0364                     status="error",
0365                     message=error_msg,
0366                     context={"required_keys": required_keys, "missing_keys": missing_keys}
0367                 )
0368             return False, error_msg
0369
0370     # Check for empty payload values
0371     empty_values = [
0372         (key, value) for key, value in job.payload.items()
0373         if value is None or (isinstance(value, str) and value.strip() == "")
0374     ]
0375
0376     if empty_values:
0377         error_msg = f"Payload contains empty values: {empty_values}"
0378         if logger:
0379             logger.checkpoint(
0380                 stage="payload_validation",
0381                 status="warning",
0382                 message=error_msg,
0383                 context={"empty_values": {k: v for k, v in empty_values}}
0384             )
0385         # Warning, not error - proceed but log it
0386
0387     # Try to resolve templates if context provided
0388     if context:
0389         try:
0390             resolved = resolve_payload_templates(job.payload, context, logger)
0391             if logger:
0392                 logger.log_debug("Payload templates resolved during validation")
0393         except RuleResolutionError as e:
0394             error_msg = f"Cannot resolve payload templates: {str(e)}"
0395             if logger:
0396                 logger.checkpoint(
0397                     stage="payload_validation",
0398                     status="error",
0399                     message=error_msg,
0400                     context={"error": str(e)}
0401                 )
0402             return False, error_msg
0403
0404     if logger:
0405         logger.checkpoint(
0406             stage="payload_validation",
0407             status="success",
0408             message="Payload validation passed",
0409             context={"payload_keys": list(job.payload.keys())}
0410         )
0411
0412     return True, None