Преглед на файлове

Improve workflow timeout handling and runtime safety

ChenYu преди 5 месеца
родител
ревизия
327bb7b121

+ 2 - 0
.env.example

@@ -15,6 +15,8 @@ AI_FLOW_REDIS_PUBLISH_CHANNEL=WORKFLOW_TASK
 
 # 工作流线程池
 AI_FLOW_EXECUTOR_POOL_SIZE=8
+AI_FLOW_HTTP_TIMEOUT=300
+AI_FLOW_STEP_RETRY_COUNT=1
 
 # 模板缓存
 AI_FLOW_TEMPLATE_CACHE_ENABLED=true

+ 2 - 0
app/config.py

@@ -24,6 +24,8 @@ class Settings(BaseSettings):
 
     # 工作流线程池
     executor_pool_size: int = 8
+    http_timeout: int = 300
+    step_retry_count: int = 1
 
     # 模板缓存
     template_cache_enabled: bool = True

+ 67 - 60
app/service/engine/workflow_engine.py

@@ -3,6 +3,7 @@ import os
 import time
 from typing import Any
 
+from app.config import settings
 from app.service.engine.execution_context import ExecutionContext
 from app.service.engine.parameter_resolver import ParameterResolver
 from app.service.engine.workflow_result import WorkflowResult
@@ -32,6 +33,7 @@ class WorkflowEngine:
         logger.info("========== Workflow start ==========  version=%s, total_steps=%d",
                      workflow.version, total_steps)
         wf_t0 = time.time()
+        workspace: str | None = None
 
         ctx = ExecutionContext(inputs)
 
@@ -42,70 +44,75 @@ class WorkflowEngine:
             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)
+        try:
+            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()))
+
+            total_elapsed = time.time() - wf_t0
+            logger.info("========== Workflow end ==========  total_elapsed=%.2fs, success=True",
+                         total_elapsed)
+            return WorkflowResult.ok(outputs)
+        finally:
+            if workspace:
+                cleanup_workspace(workspace, task_id)
 
     def _execute_step_with_retry(
         self, task_id: str, workspace: str, step: Step, ctx: ExecutionContext
     ) -> TaskResult:
-        max_retries = step.retry_count
+        max_retries = (
+            settings.step_retry_count if step.retry_count is None else step.retry_count
+        )
         last_result: TaskResult | None = None
 
         for attempt in range(max_retries + 1):

+ 1 - 1
app/service/spec/step.py

@@ -9,4 +9,4 @@ class Step:
     resources: str = ""
     inputs: dict[str, Any] = field(default_factory=dict)
     outputs: dict[str, Any] = field(default_factory=dict)
-    retry_count: int = 0
+    retry_count: Optional[int] = None

+ 1 - 1
app/service/spec/workflow_loader.py

@@ -37,7 +37,7 @@ def _parse_workflow(root: dict) -> Workflow:
             resources=node.get("resources", ""),
             inputs=node.get("inputs") or {},
             outputs=node.get("outputs") or {},
-            retry_count=node.get("retryCount", 0),
+            retry_count=node.get("retryCount"),
         )
         steps[step_id] = step
 

+ 3 - 3
app/service/task/impl/google_veo3_task.py

@@ -48,7 +48,7 @@ class GoogleVeo3Task(BaseTask):
     def _create_request(self, inputs: dict) -> str:
         endpoint = f"{settings.gemini_base_url}/models/{settings.gemini_veo3_model}:predictLongRunning"
         logger.info("[google_veo3] Submitting generation request: POST %s", endpoint)
-        with logged_client(timeout=120) as client:
+        with logged_client(timeout=settings.http_timeout) as client:
             resp = client.post(endpoint, headers=self._headers(), json=inputs)
             logger.info("[google_veo3] Submit response: status=%d", resp.status_code)
             resp.raise_for_status()
@@ -65,7 +65,7 @@ class GoogleVeo3Task(BaseTask):
 
         for i in range(max_retries):
             try:
-                with logged_client(timeout=60) as client:
+                with logged_client(timeout=settings.http_timeout) as client:
                     resp = client.get(url, headers=self._headers())
                     if not resp.is_success:
                         logger.warning("[google_veo3] Poll #%d: HTTP %d, elapsed=%.1fs",
@@ -97,7 +97,7 @@ class GoogleVeo3Task(BaseTask):
 
     def _download_video(self, video_uri: str) -> bytes:
         t0 = time.time()
-        with logged_client(timeout=120, follow_redirects=True) as client:
+        with logged_client(timeout=settings.http_timeout, follow_redirects=True) as client:
             resp = client.get(video_uri, headers=self._headers())
             resp.raise_for_status()
             logger.info("[google_veo3] Video downloaded: %d bytes in %.2fs",

+ 9 - 8
app/service/task/impl/image_editor_task.py

@@ -35,12 +35,12 @@ class ImageEditorTask(BaseTask):
             self._execute_p2p(images, param.resources, output_path, param.workspace)
 
             out_file = os.path.join(param.workspace, output_path)
-            if os.path.isfile(out_file):
-                file_size = os.path.getsize(out_file)
-                logger.info("[image_editor] Completed: output=%s (%.1f KB) in %.2fs",
-                             output_path, file_size / 1024, time.time() - t0)
-            else:
-                logger.warning("[image_editor] Completed but output file not found: %s", output_path)
+            if not os.path.isfile(out_file):
+                raise RuntimeError(f"Output file not found: {output_path}")
+
+            file_size = os.path.getsize(out_file)
+            logger.info("[image_editor] Completed: output=%s (%.1f KB) in %.2fs",
+                        output_path, file_size / 1024, time.time() - t0)
 
             return TaskResult.success({"images": [output_path]})
         except Exception as e:
@@ -71,5 +71,6 @@ class ImageEditorTask(BaseTask):
 
         if result.returncode != 0:
             logger.error("[image_editor] p2pconvertor failed: exit_code=%d in %.2fs", result.returncode, cmd_elapsed)
-        else:
-            logger.info("[image_editor] p2pconvertor completed: exit_code=0 in %.2fs", cmd_elapsed)
+            raise RuntimeError(f"p2pconvertor failed with exit code {result.returncode}")
+
+        logger.info("[image_editor] p2pconvertor completed: exit_code=0 in %.2fs", cmd_elapsed)

+ 2 - 2
app/service/task/impl/keling_image2video_task.py

@@ -43,7 +43,7 @@ class KelingImage2VideoTask(BaseTask):
         logger.info("[%s] Submitting task: POST %s  (remaining_retries=%d)",
                      self.get_type(), self.execute_url, remaining)
         try:
-            with logged_client(timeout=60) as client:
+            with logged_client(timeout=settings.http_timeout) as client:
                 resp = client.post(self.execute_url, headers=self._headers(), json=param.inputs)
                 logger.info("[%s] Submit response: status=%d", self.get_type(), resp.status_code)
                 bean = resp.json()
@@ -107,6 +107,6 @@ class KelingImage2VideoTask(BaseTask):
 
     def _poll(self, task_id: str) -> dict:
         url = f"{self.search_url}/{task_id}"
-        with logged_client(timeout=60) as client:
+        with logged_client(timeout=settings.http_timeout) as client:
             resp = client.get(url, headers=self._headers())
             return resp.json()

+ 2 - 2
app/service/task/impl/kontext_pro_task.py

@@ -28,7 +28,7 @@ class KontextProTask(BaseTask):
         logger.info("[%s] Submitting task: POST %s", self.get_type(), self.execute_url)
         t0 = time.time()
         try:
-            with logged_client(timeout=60) as client:
+            with logged_client(timeout=settings.http_timeout) as client:
                 resp = client.post(self.execute_url, headers=self._headers(), json=param.inputs)
                 elapsed = time.time() - t0
                 logger.info("[%s] Submit response: status=%d, elapsed=%.2fs",
@@ -80,6 +80,6 @@ class KontextProTask(BaseTask):
         return TaskResult.success({"images": images})
 
     def _poll(self, polling_url: str) -> dict:
-        with logged_client(timeout=60) as client:
+        with logged_client(timeout=settings.http_timeout) as client:
             resp = client.get(polling_url, headers=self._headers())
             return resp.json()

+ 1 - 1
app/service/task/impl/nano_banana_task.py

@@ -24,7 +24,7 @@ class NanoBananaTask(BaseTask):
         logger.info("[nano_banana] Calling Gemini API: POST %s", endpoint)
         t0 = time.time()
         try:
-            with logged_client(timeout=120) as client:
+            with logged_client(timeout=settings.http_timeout) as client:
                 resp = client.post(endpoint, headers=headers, json=param.inputs)
                 elapsed = time.time() - t0
                 logger.info("[nano_banana] Response: status=%d, elapsed=%.2fs, body_size=%d",

+ 2 - 2
app/service/task/impl/seedance_video_task.py

@@ -26,7 +26,7 @@ class SeedanceVideoTask(BaseTask):
 
         logger.info("[seedance_video] Submitting task: POST %s", execute_url)
         try:
-            with logged_client(timeout=60) as client:
+            with logged_client(timeout=settings.http_timeout) as client:
                 resp = client.post(execute_url, headers=self._headers(), json=param.inputs)
                 logger.info("[seedance_video] Submit response: status=%d, elapsed=%.2fs",
                             resp.status_code, time.time() - t0)
@@ -64,7 +64,7 @@ class SeedanceVideoTask(BaseTask):
 
         for attempt in range(1, max_attempts + 1):
             try:
-                with logged_client(timeout=60) as client:
+                with logged_client(timeout=settings.http_timeout) as client:
                     resp = client.get(search_url, headers=self._headers())
                     data = resp.json()
 

+ 1 - 1
app/service/task/impl/seedream_image_task.py

@@ -26,7 +26,7 @@ class SeedreamImageTask(BaseTask):
         t0 = time.time()
 
         try:
-            with logged_client(timeout=60) as client:
+            with logged_client(timeout=settings.http_timeout) as client:
                 resp = client.post(url, headers=self._headers(), json=param.inputs)
                 elapsed = time.time() - t0
                 logger.info("[seedream_image] Response: status=%d, elapsed=%.2fs",

+ 5 - 2
app/service/task/impl/video_editor_task.py

@@ -39,6 +39,8 @@ class VideoEditorTask(BaseTask):
 
             self._execute_compile(param.workspace, param.resources)
             output_path = self._move_output(param.workspace, param.resources, param.step_id)
+            if not output_path:
+                raise RuntimeError("No output path generated by video_editor")
 
             logger.info("[video_editor] Completed: output=%s in %.2fs", output_path, time.time() - t0)
             return TaskResult.success({"videos": [output_path]})
@@ -86,8 +88,9 @@ class VideoEditorTask(BaseTask):
 
         if result.returncode != 0:
             logger.error("[video_editor] compile failed: exit_code=%d in %.2fs", result.returncode, cmd_elapsed)
-        else:
-            logger.info("[video_editor] compile completed: exit_code=0 in %.2fs", cmd_elapsed)
+            raise RuntimeError(f"compile failed with exit code {result.returncode}")
+
+        logger.info("[video_editor] compile completed: exit_code=0 in %.2fs", cmd_elapsed)
 
     @staticmethod
     def _move_output(workspace: str, resources: str, step_id: str) -> str | None:

+ 2 - 2
app/service/task/task_executor.py

@@ -13,8 +13,8 @@ class _NoopTask(BaseTask):
         return "noop"
 
     def execute(self, param: TaskParam) -> TaskResult:
-        logger.warning("[noop] No task implementation found for type=%s, returning empty success", param.type)
-        return TaskResult.success({})
+        logger.error("[noop] No task implementation found for type=%s", param.type)
+        return TaskResult.failure(f"No task implementation found for type '{param.type}'")
 
 
 class TaskExecutor:

+ 3 - 1
app/util/download_util.py

@@ -5,6 +5,8 @@ from urllib.parse import urlparse
 
 import httpx
 
+from app.config import settings
+
 logger = logging.getLogger(__name__)
 
 
@@ -25,7 +27,7 @@ def download_file(url: str, target_dir: str) -> str | None:
             return filepath
 
         logger.info("Starting download from: %s to: %s", url, filepath)
-        with httpx.Client(timeout=60, follow_redirects=True) as client:
+        with httpx.Client(timeout=settings.http_timeout, follow_redirects=True) as client:
             resp = client.get(url)
             resp.raise_for_status()
             with open(filepath, "wb") as f:

+ 13 - 3
app/util/http_util.py

@@ -6,6 +6,7 @@ import httpx
 logger = logging.getLogger("app.http")
 
 _MAX_BODY_LOG_LENGTH = 2000
+_BODY_PREVIEW_LENGTH = 500
 
 
 def _truncate(text: str, max_len: int = _MAX_BODY_LOG_LENGTH) -> str:
@@ -14,6 +15,14 @@ def _truncate(text: str, max_len: int = _MAX_BODY_LOG_LENGTH) -> str:
     return text[:max_len] + f"... (truncated, total {len(text)} chars)"
 
 
+def _body_summary(text: str) -> str:
+    if not text:
+        return ""
+    if len(text) > _MAX_BODY_LOG_LENGTH:
+        return f"<omitted, body too large: {len(text)} chars>"
+    return _truncate(text, _BODY_PREVIEW_LENGTH)
+
+
 def _on_request(request: httpx.Request) -> None:
     body = request.content.decode("utf-8", errors="replace") if request.content else ""
     logger.info(
@@ -21,18 +30,19 @@ def _on_request(request: httpx.Request) -> None:
         request.method,
         request.url,
         len(body),
-        _truncate(body),
+        _body_summary(body),
     )
 
 
 def _on_response(response: httpx.Response) -> None:
     response.read()
+    body = response.text
     logger.info(
         "[HTTP <<< ] %s %s | body(%d): %s",
         response.status_code,
         response.url,
-        len(response.text),
-        _truncate(response.text),
+        len(body),
+        _body_summary(body),
     )