fix(grpc): use context.abort() with proper status codes instead of in-band errors (#19972)

Signed-off-by: Chang Su <chang.s.su@oracle.com>

python/sglang/srt/entrypoints/grpc_server.py

@@ -31,6 +31,7 @@ from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationM
 from sglang.srt.grpc.grpc_request_manager import GrpcRequestManager
 from sglang.srt.grpc.health_servicer import SGLangHealthServicer
 from sglang.srt.grpc.scheduler_launcher import launch_scheduler_process_only
+from sglang.srt.grpc.utils import abort_code_from_output
 from sglang.srt.managers.disagg_service import start_disagg_service
 from sglang.srt.managers.io_struct import (
     GetLoadsReqOutput,
@@ -201,14 +202,9 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
                 if isinstance(output, list):
                     for batch_output in output:
                         if "error" in batch_output:
-                            yield sglang_scheduler_pb2.GenerateResponse(
+                            await context.abort(
-                                request_id=request.request_id,
+                                abort_code_from_output(batch_output),
-                                error=sglang_scheduler_pb2.GenerateError(
+                                batch_output["error"],
-                                    message=batch_output["error"],
-                                    http_status_code=(
-                                        "500" if "abort" not in batch_output else "499"
-                                    ),
-                                ),
                             )
                         else:
                             # All non-error batch outputs are final responses
@@ -218,14 +214,9 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
                 else:
                     # Handle single response (for streaming or n=1 non-streaming)
                     if "error" in output:
-                        yield sglang_scheduler_pb2.GenerateResponse(
+                        await context.abort(
-                            request_id=request.request_id,
+                            abort_code_from_output(output),
-                            error=sglang_scheduler_pb2.GenerateError(
+                            output["error"],
-                                message=output["error"],
-                                http_status_code=(
-                                    "500" if "abort" not in output else "499"
-                                ),
-                            ),
                         )
                     elif request.stream:
                         yield self._create_chunk_response(request.request_id, output)
@@ -239,24 +230,22 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
                             request.request_id, output
                         )
 
+        except grpc.aio.AbortError:
+            raise
+        except ValueError as e:
+            logger.warning(f"Generate invalid request {request.request_id}: {e}")
+            await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
         except Exception as e:
             logger.error(
                 f"Generate failed for request {request.request_id}: {e}\n"
                 f"{get_exception_traceback()}"
             )
-            yield sglang_scheduler_pb2.GenerateResponse(
+            await context.abort(grpc.StatusCode.INTERNAL, str(e))
-                request_id=request.request_id,
-                error=sglang_scheduler_pb2.GenerateError(
-                    message=str(e),
-                    http_status_code="500",
-                    details=get_exception_traceback(),
-                ),
-            )
 
     async def Embed(
         self,
         request: sglang_scheduler_pb2.EmbedRequest,
-        _context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext,
     ) -> sglang_scheduler_pb2.EmbedResponse:
         """Handle embedding requests."""
         logger.info(f"Receive embedding request: {request.request_id}")
@@ -281,19 +270,17 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
                 ),
             )
 
+        except grpc.aio.AbortError:
+            raise
+        except ValueError as e:
+            logger.warning(f"Embed invalid request {request.request_id}: {e}")
+            await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
         except Exception as e:
             logger.error(
                 f"Embed failed for request {request.request_id}: {e}\n"
                 f"{get_exception_traceback()}"
             )
-            return sglang_scheduler_pb2.EmbedResponse(
+            await context.abort(grpc.StatusCode.INTERNAL, str(e))
-                request_id=request.request_id,
-                error=sglang_scheduler_pb2.EmbedError(
-                    message=str(e),
-                    code="INTERNAL_ERROR",
-                    details=get_exception_traceback(),
-                ),
-            )
 
     async def HealthCheck(
         self,

python/sglang/srt/grpc/utils.py

@@ -0,0 +1,21 @@
+"""gRPC utility functions."""
+
+from http import HTTPStatus
+
+import grpc
+
+_HTTP_TO_GRPC_CODE = {
+    HTTPStatus.BAD_REQUEST: grpc.StatusCode.INVALID_ARGUMENT,
+    HTTPStatus.SERVICE_UNAVAILABLE: grpc.StatusCode.UNAVAILABLE,
+    HTTPStatus.INTERNAL_SERVER_ERROR: grpc.StatusCode.INTERNAL,
+}
+
+
+def abort_code_from_output(output: dict) -> grpc.StatusCode:
+    """Map a scheduler error output to the appropriate gRPC status code."""
+    finish_reason = output.get("meta_info", {}).get("finish_reason")
+    if isinstance(finish_reason, dict):
+        status_code = finish_reason.get("status_code")
+        if status_code is not None:
+            return _HTTP_TO_GRPC_CODE.get(status_code, grpc.StatusCode.INTERNAL)
+    return grpc.StatusCode.INTERNAL