Merge pull request #150 from luozirong2025/main

Preserve error codes and surface root cause in agentic RL errors

Author: lzsweb

dashscope/finetune/agentic_rl.py

@@ -106,6 +106,8 @@ async def register_functions(
             )
             logger.info("Function components registered")
         except Exception as e:
+            if hasattr(e, "error_code"):
+                raise
             raise RegistrationError(
                 "Function registration failed",
                 error_code=3002,
@@ -249,6 +251,8 @@ def submit_job(
                 **kwargs,
             )
         except Exception as e:
+            if hasattr(e, "error_code"):
+                raise
             raise RuntimeErrorWithCode(
                 "Job submission failed",
                 error_code=3005,
@@ -312,6 +316,8 @@ async def run(
                 **kwargs,
             )
         except Exception as e:
+            if hasattr(e, "error_code"):
+                raise
             raise RuntimeErrorWithCode(
                 "RL tuning workflow failed",
                 error_code=3006,

dashscope/finetune/reinforcement/common/errors.py

@@ -7,7 +7,32 @@
 from typing import Optional, Dict
 
 
-class AgenticRLError(Exception):
+class _RootCauseMixin:
+    """Mixin that provides root-cause traversal and formatting for exceptions
+    that carry an error code."""
+
+    @property
+    def root_cause(self) -> "Exception":
+        root: "Exception" = self  # type: ignore[assignment]
+        seen = {id(root)}
+        while root.__cause__ and id(root.__cause__) not in seen:
+            root = root.__cause__
+            seen.add(id(root))
+        return root
+
+    def _format_cause(self) -> str:
+        """Format root cause information if available."""
+        if self.__cause__ is not None:
+            root = self.root_cause
+            if self is not root:
+                cause_msg = str(root).split("\n", maxsplit=1)[0][:100].strip()
+                if cause_msg:
+                    return f" (caused by: {type(root).__name__}: {cause_msg})"
+                return f" (caused by: {type(root).__name__})"
+        return ""
+
+
+class AgenticRLError(_RootCauseMixin, Exception):
     """Base class for all Agentic RL exceptions."""
 
     def __init__(self, message: str, error_code: int = 1000):
@@ -16,15 +41,9 @@ def __init__(self, message: str, error_code: int = 1000):
         self.timestamp = datetime.now().isoformat()
         self.message = message
 
-    @property
-    def root_cause(self) -> Exception:
-        root = self
-        while root.__cause__:
-            root = root.__cause__
-        return root
-
     def __str__(self):
-        return f"[{self.error_code}] {self.message} (at {self.timestamp})"
+        base = f"[{self.error_code}] {self.message} (at {self.timestamp})"
+        return f"{base}{self._format_cause()}"
 
 
 class IOErrorWithCode(AgenticRLError):
@@ -42,7 +61,7 @@ def __init__(
         self.operation = operation
 
 
-class RuntimeErrorWithCode(RuntimeError):
+class RuntimeErrorWithCode(_RootCauseMixin, RuntimeError):
     """Enhanced RuntimeError that supports error codes for better error
     categorization."""
 
@@ -51,18 +70,11 @@ def __init__(self, message: str, error_code: int = 0):
         self.error_code = error_code
         self.message = message
 
-    @property
-    def root_cause(self) -> Exception:
-        root = self
-        while root.__cause__:
-            root = root.__cause__
-        return root
-
     def __str__(self):
-        return f"[{self.error_code}] {self.message}"
+        return f"[{self.error_code}] {self.message}{self._format_cause()}"
 
 
-class ValueErrorWithCode(ValueError):
+class ValueErrorWithCode(_RootCauseMixin, ValueError):
     """Enhanced ValueError that supports error codes for better error
     categorization."""
 
@@ -71,15 +83,8 @@ def __init__(self, message: str, error_code: int = 0):
         self.error_code = error_code
         self.message = message
 
-    @property
-    def root_cause(self) -> Exception:
-        root = self
-        while root.__cause__:
-            root = root.__cause__
-        return root
-
     def __str__(self):
-        return f"[{self.error_code}] {self.message}"
+        return f"[{self.error_code}] {self.message}{self._format_cause()}"
 
 
 class InputError(AgenticRLError):

dashscope/finetune/reinforcement/common/model.py

@@ -348,6 +348,8 @@ async def get_oss(self, oss_id: Optional[str] = None) -> str:
             )
             return self.oss_signed_url
 
+        except OSSConnectionError:
+            raise
         except Exception as e:
             raise OSSConnectionError(
                 "Failed to obtain OSS URL",
@@ -1243,6 +1245,8 @@ async def register_functions(
                         )
 
         except Exception as e:
+            if hasattr(e, "error_code"):
+                raise
             raise RegistrationError(
                 "Function component registration failed",
                 error_code=2055,

dashscope/finetune/reinforcement/common/utils.py

@@ -1004,3 +1004,46 @@ def serialize_for_output(data: Any) -> Any:
 
     # Return basic types directly
     return data
+
+
+def get_fc_request_id(request) -> str:
+    """Extract Function Compute request ID from request headers.
+
+    Retrieves the 'x-fc-request-id' header from a FastAPI/Starlette Request
+    object. This is useful for correlating logs with specific FC invocations.
+
+    Args:
+        request: FastAPI/Starlette Request object (or any object with a
+            `.headers` mapping).
+
+    Returns:
+        The FC request ID string, or "unknown" if the header is not present.
+    """
+    if request is None:
+        return "unknown"
+    headers = getattr(request, "headers", None)
+    if headers is not None and hasattr(headers, "get"):
+        return headers.get("x-fc-request-id", "unknown")
+    return "unknown"
+
+
+def get_business_summary(processor_input) -> str:
+    """Extract summary business information from processor input.
+
+    Returns the string representation of request_metadata if present,
+    otherwise returns an empty string.
+
+    Args:
+        processor_input: The processor input object (BaseDataModel subclass).
+
+    Returns:
+        String representation of request_metadata, or empty string.
+    """
+    if processor_input is None:
+        return ""
+
+    request_metadata = getattr(processor_input, "request_metadata", None)
+    if request_metadata is None:
+        return ""
+
+    return str(request_metadata)

dashscope/finetune/reinforcement/component/server/server.py

@@ -37,6 +37,7 @@
   GET  /health Health check
 """
 
+import asyncio
 import logging
 import os
 import time
@@ -47,6 +48,10 @@
 from fastapi.responses import JSONResponse
 
 from dashscope.finetune.reinforcement.common.log import logger
+from dashscope.finetune.reinforcement.common.utils import (
+    get_fc_request_id,
+    get_business_summary,
+)
 from dashscope.finetune.reinforcement.common.model_types import (
     FunctionType as FuncType,
 )
@@ -72,7 +77,7 @@
     "0",
     "no",
 )
-_THREAD_POOL_WORKERS = int(os.getenv("THREAD_POOL_WORKERS", "4"))
+_THREAD_POOL_WORKERS = int(os.getenv("THREAD_POOL_WORKERS", "32"))
 _THREAD_POOL_QUEUE = int(os.getenv("THREAD_POOL_QUEUE", "100"))
 
 if not _ENABLE_LOGGING:
@@ -377,6 +382,52 @@ def _serialize_result(result: Any) -> Dict:
         return {"result": result}
 
 
+async def _cancel_task_on_disconnect(
+    process_task: "asyncio.Task",
+    disconnect_listener: "asyncio.Task",
+    disconnected: "asyncio.Event",
+) -> bool:
+    """Wait for processing or disconnect, cancel task if disconnected.
+
+    Returns True if disconnected, False if processing completed normally.
+    """
+    _done, _pending = await asyncio.wait(
+        [process_task, disconnect_listener],
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+
+    if disconnected.is_set():
+        process_task.cancel()
+        try:
+            await process_task
+        except asyncio.CancelledError:
+            pass
+        return True
+
+    return False
+
+
+async def _log_request_metrics(
+    request: "Request",
+    processor_input,
+    start_time: float,
+    success: bool,
+    cancelled: bool,
+) -> None:
+    """Log request metrics and force flush if configured."""
+    elapsed = round(time.time() - start_time, 4)
+    fc_req_id = get_fc_request_id(request)
+    biz_summary = get_business_summary(processor_input)
+    biz_part = f" | {biz_summary}" if biz_summary else ""
+    logger.info(
+        f"[Server] /api/v1 | func_type={func_type.value} | "
+        f"fc_request_id={fc_req_id} | "
+        f"success={success} | cancelled={cancelled} | "
+        f"elapsed={elapsed}s{biz_part}",
+    )
+    await maybe_force_flush_async(reason="request")
+
+
 @app.post("/api/v1")
 async def handle_endpoint(request: Request) -> JSONResponse:
     """
@@ -386,6 +437,10 @@ async def handle_endpoint(request: Request) -> JSONResponse:
     request body, executes business logic using configured processor,
     and returns serialized result.
 
+    Monitors client connection state via ASGI disconnect messages. If the
+    client disconnects before processing completes, the processing task is
+    cancelled to avoid wasting resources.
+
     Request body format: JSON, fields determined by FuncType:
     - reward: See RewardInput
     - rollout: See RolloutInput
@@ -398,8 +453,25 @@ async def handle_endpoint(request: Request) -> JSONResponse:
     """
     start_time = time.time()
     success = False
+    cancelled = False
     processor_input = None
 
+    disconnect_listener = None
+    disconnected = asyncio.Event()
+
+    async def _listen_for_disconnect():
+        """Background listener for ASGI disconnect messages."""
+        try:
+            while True:
+                message = await request.receive()
+                if message.get("type") == "http.disconnect":
+                    disconnected.set()
+                    break
+        except Exception:
+            # If receive() raises (e.g. connection already closed),
+            # treat as disconnected.
+            disconnected.set()
+
     # Extract trace context from request headers
     _otel_ctx_token, _upstream_tokens = await _extract_trace_context(request)
 
@@ -413,8 +485,39 @@ async def handle_endpoint(request: Request) -> JSONResponse:
         # Check thread pool queue capacity
         await _check_queue_capacity()
 
-        # Execute processor
-        result = await func_manager.processes(processor_input)
+        # Start listening for disconnect only AFTER the request body has been
+        # fully read.  Otherwise the background listener competes with the body
+        # reader for ASGI receive() messages, causing hung requests or parse
+        # failures.
+        disconnect_listener = asyncio.create_task(_listen_for_disconnect())
+
+        # Execute processor as a task so we can cancel it on disconnect
+        process_task = asyncio.create_task(
+            func_manager.processes(processor_input),
+        )
+
+        # Wait for either the processing to finish or client disconnect
+        disconnected_flag = await _cancel_task_on_disconnect(
+            process_task,
+            disconnect_listener,
+            disconnected,
+        )
+
+        if disconnected_flag:
+            cancelled = True
+            fc_request_id = get_fc_request_id(request)
+            logger.warning(
+                "[Server] Client disconnected during processing, "
+                "task cancelled. x-fc-request-id: %s",
+                fc_request_id,
+            )
+            return JSONResponse(
+                status_code=499,
+                content={"message": "Client disconnected, request cancelled."},
+            )
+
+        # Processing completed normally
+        result = process_task.result()
         success = True
 
         # Serialize result
@@ -425,26 +528,40 @@ async def handle_endpoint(request: Request) -> JSONResponse:
     except HTTPException:
         # Re-raise HTTP exceptions as-is
         raise
+    except asyncio.CancelledError:
+        cancelled = True
+        logger.warning("[Server] Request processing was cancelled.")
+        return JSONResponse(
+            status_code=499,
+            content={"message": "Request cancelled."},
+        )
     except Exception as ex:
         logger.error(f"[Server] Unexpected error: {ex}", exc_info=True)
         return JSONResponse(
             status_code=500,
             content={"message": str(ex)},
         )
     finally:
+        # Cancel the disconnect listener if still running
+        if disconnect_listener is not None:
+            disconnect_listener.cancel()
+            try:
+                await disconnect_listener
+            except asyncio.CancelledError:
+                pass
+
         # Clean up trace context
         await _cleanup_trace_context(_otel_ctx_token, _upstream_tokens)
 
         # Log request metrics
-        elapsed = round(time.time() - start_time, 4)
-        logger.info(
-            f"[Server] /api/v1 | func_type={func_type.value} | "
-            f"success={success} | elapsed={elapsed}s",
+        await _log_request_metrics(
+            request,
+            processor_input,
+            start_time,
+            success,
+            cancelled,
         )
 
-        # Best-effort force flush based on platform/internal env config
-        await maybe_force_flush_async(reason="request")
-
 
 @app.get("/health")
 async def health_check_endpoint() -> JSONResponse: