import logging import os import time from typing import Any from app.service.engine.execution_context import ExecutionContext from app.service.engine.parameter_resolver import ParameterResolver from app.service.engine.workflow_result import WorkflowResult from app.service.spec.step import Step from app.service.spec.workflow import Workflow from app.service.task.task_executor import TaskExecutor from app.service.task.task_param import TaskParam from app.service.task.task_result import TaskResult from app.util.task_util import create_task_workspace, cleanup_workspace logger = logging.getLogger(__name__) class WorkflowEngine: _BASE64_DECODE_TYPE = "base64_decode" _BASE64_MISSING_FILES_ERROR = "No files provided for base64 decoding" _MAX_PREVIOUS_STEP_RERUNS = 3 def __init__(self, task_executor: TaskExecutor): self._task_executor = task_executor self._resolver = ParameterResolver() def execute( self, task_id: str, workflow: Workflow, inputs: dict[str, Any] ) -> WorkflowResult: total_steps = len(workflow.order) logger.info("========== Workflow start ========== version=%s, total_steps=%d", workflow.version, total_steps) wf_t0 = time.time() ctx = ExecutionContext(inputs) try: workspace = create_task_workspace(task_id, workflow.template_path) logger.info("Workspace created: %s", workspace) except Exception as e: logger.error("Failed to create workspace: %s", e) return WorkflowResult.failure(f"failed: {e}") for idx, step_id in enumerate(workflow.order): step = workflow.steps.get(step_id) if step is None: logger.warning("Step '%s' not found in steps definition, skipping", step_id) continue step_no = idx + 1 logger.info("------ Step %d/%d [%s] type=%s start ------", step_no, total_steps, step.id, step.type) logger.info("[%s] raw_inputs=%s", step.id, step.inputs) step_t0 = time.time() try: step_dir = os.path.join(workspace, step.id) os.makedirs(step_dir, exist_ok=True) result = self._execute_step_with_retry(task_id, workspace, step, ctx) step_elapsed = time.time() - step_t0 if result.is_failure(): recovered_result = self._recover_base64_by_rerunning_previous_step( task_id=task_id, workspace=workspace, workflow=workflow, current_step_index=idx, current_step=step, ctx=ctx, failed_result=result, ) if recovered_result is not None: result = recovered_result step_elapsed = time.time() - step_t0 if result.is_failure(): logger.error("[%s] FAILED in %.2fs after retries. Error: %s", step_id, step_elapsed, result.error_message) return WorkflowResult.failure( f"Step {step_id} failed: {result.error_message}" ) except Exception as e: step_elapsed = time.time() - step_t0 logger.error("[%s] EXCEPTION in %.2fs: %s", step_id, step_elapsed, e, exc_info=True) return WorkflowResult.failure(f"Step {step_id} failed: {e}") logger.info("[%s] SUCCESS in %.2fs outputs=%s", step.id, step_elapsed, _shorten(result.outputs)) ctx.put_step_result(step_id, result.outputs) logger.info("All steps completed, resolving final outputs ...") outputs: dict[str, Any] = {} if workflow.outputs: for key, value in workflow.outputs.items(): outputs[key] = self._resolver.resolve(value, ctx) logger.info("Final output_keys=%s", list(outputs.keys())) cleanup_workspace(workspace, task_id) total_elapsed = time.time() - wf_t0 logger.info("========== Workflow end ========== total_elapsed=%.2fs, success=True", total_elapsed) return WorkflowResult.ok(outputs) def _execute_step_with_retry( self, task_id: str, workspace: str, step: Step, ctx: ExecutionContext ) -> TaskResult: max_retries = step.retry_count last_result: TaskResult | None = None for attempt in range(max_retries + 1): step_inputs = self._resolve_inputs(step.inputs, ctx) task = self._task_executor.resolve(step.type) param = TaskParam( workspace=workspace, step_id=step.id, type=step.type, resources=step.resources, inputs=step_inputs, ) if attempt > 0: logger.info("[%s] Retry attempt %d/%d", step.id, attempt, max_retries) logger.info("[%s] Executing task type=%s, resolved_input_keys=%s", step.id, step.type, list(step_inputs.keys())) attempt_t0 = time.time() last_result = task.execute(param) attempt_elapsed = time.time() - attempt_t0 if last_result.is_success(): logger.info("[%s] Task returned SUCCESS in %.2fs", step.id, attempt_elapsed) return last_result if attempt < max_retries: logger.warning( "[%s] Task returned FAILURE in %.2fs (attempt %d/%d), will retry. Error: %s", step.id, attempt_elapsed, attempt + 1, max_retries + 1, last_result.error_message, ) time.sleep(1) else: logger.error( "[%s] Task returned FAILURE in %.2fs (final attempt %d/%d). Error: %s", step.id, attempt_elapsed, attempt + 1, max_retries + 1, last_result.error_message, ) return last_result # type: ignore[return-value] def _resolve_inputs( self, inputs: dict[str, Any], ctx: ExecutionContext ) -> dict[str, Any]: return {k: self._resolver.resolve(v, ctx) for k, v in inputs.items()} if inputs else {} def _recover_base64_by_rerunning_previous_step( self, task_id: str, workspace: str, workflow: Workflow, current_step_index: int, current_step: Step, ctx: ExecutionContext, failed_result: TaskResult, ) -> TaskResult | None: if current_step.type != self._BASE64_DECODE_TYPE: return None if self._BASE64_MISSING_FILES_ERROR not in failed_result.error_message: return None if current_step_index <= 0: logger.error( "[%s] Missing files in base64_decode but no previous step exists", current_step.id, ) return None previous_step_id = workflow.order[current_step_index - 1] previous_step = workflow.steps.get(previous_step_id) if previous_step is None: logger.error( "[%s] Missing files in base64_decode but previous step '%s' not found", current_step.id, previous_step_id, ) return None current_result = failed_result for recover_attempt in range(1, self._MAX_PREVIOUS_STEP_RERUNS + 1): logger.warning( "[%s] Trigger fallback: rerun previous step [%s] (attempt %d/%d) " "because base64_decode has no files", current_step.id, previous_step.id, recover_attempt, self._MAX_PREVIOUS_STEP_RERUNS, ) previous_result = self._execute_step_with_retry(task_id, workspace, previous_step, ctx) if previous_result.is_failure(): logger.warning( "[%s] Previous step rerun failed on attempt %d/%d: %s", previous_step.id, recover_attempt, self._MAX_PREVIOUS_STEP_RERUNS, previous_result.error_message, ) continue ctx.put_step_result(previous_step.id, previous_result.outputs) logger.info( "[%s] Previous step rerun succeeded on attempt %d/%d, rerunning current step", previous_step.id, recover_attempt, self._MAX_PREVIOUS_STEP_RERUNS, ) current_result = self._execute_step_with_retry(task_id, workspace, current_step, ctx) if current_result.is_success(): logger.info( "[%s] Recovery succeeded after rerunning previous step [%s]", current_step.id, previous_step.id, ) return current_result if self._BASE64_MISSING_FILES_ERROR not in current_result.error_message: logger.error( "[%s] Current step rerun failed with non-recoverable error: %s", current_step.id, current_result.error_message, ) return current_result logger.error( "[%s] Recovery failed after %d attempts of rerunning previous step [%s]", current_step.id, self._MAX_PREVIOUS_STEP_RERUNS, previous_step.id, ) return current_result def _shorten(d: dict | None, max_len: int = 120) -> str: if d is None: return "null" s = str(d) return s if len(s) <= max_len else s[:max_len] + "..."