diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index b2558841b..e2d92452a 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/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( - request_id=request.request_id, - error=sglang_scheduler_pb2.GenerateError( - message=batch_output["error"], - http_status_code=( - "500" if "abort" not in batch_output else "499" - ), - ), + await context.abort( + abort_code_from_output(batch_output), + batch_output["error"], ) 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( - request_id=request.request_id, - error=sglang_scheduler_pb2.GenerateError( - message=output["error"], - http_status_code=( - "500" if "abort" not in output else "499" - ), - ), + await context.abort( + abort_code_from_output(output), + output["error"], ) 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( - request_id=request.request_id, - error=sglang_scheduler_pb2.GenerateError( - message=str(e), - http_status_code="500", - details=get_exception_traceback(), - ), - ) + await context.abort(grpc.StatusCode.INTERNAL, str(e)) 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( - request_id=request.request_id, - error=sglang_scheduler_pb2.EmbedError( - message=str(e), - code="INTERNAL_ERROR", - details=get_exception_traceback(), - ), - ) + await context.abort(grpc.StatusCode.INTERNAL, str(e)) async def HealthCheck( self, diff --git a/python/sglang/srt/grpc/utils.py b/python/sglang/srt/grpc/utils.py new file mode 100644 index 000000000..2258a725f --- /dev/null +++ b/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