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>
This commit is contained in:
Chang Su
2026-03-05 14:53:18 -08:00
committed by GitHub
parent 46dced64ea
commit dd2bbe6d62
2 changed files with 41 additions and 33 deletions

View File

@@ -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,

View File

@@ -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