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:
@@ -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,
|
||||
|
||||
21
python/sglang/srt/grpc/utils.py
Normal file
21
python/sglang/srt/grpc/utils.py
Normal 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
|
||||
Reference in New Issue
Block a user