[model gateway][0/N] router EPD support: add encoder grpc server backend support (#16552)
Co-authored-by: Zongyao Chen <ZongYao.Chen@linux.alibaba.com> Co-authored-by: Zongyao Chen <solar1s@163.com>
This commit is contained in:
co-authored by
Zongyao Chen
Zongyao Chen
parent
facde4c6d3
commit
d939e26585
@@ -162,6 +162,14 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
|
||||
self.scheduler_info = scheduler_info
|
||||
self.start_time = time.time()
|
||||
self.health_servicer = health_servicer
|
||||
self.mm_receiver = None
|
||||
if (
|
||||
self.server_args.language_only
|
||||
and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
|
||||
):
|
||||
from sglang.srt.disaggregation import encode_receiver as mm_receiver
|
||||
|
||||
self.mm_receiver = mm_receiver.create_mm_receiver(self.server_args)
|
||||
|
||||
# Start the request manager's event loop using auto_create_handle_loop
|
||||
self.request_manager.auto_create_handle_loop()
|
||||
@@ -179,6 +187,7 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
|
||||
try:
|
||||
# Convert gRPC request to internal format
|
||||
tokenized_req = self._convert_generate_request(request)
|
||||
self._handle_epd_disaggregation_encode_request(request, tokenized_req)
|
||||
|
||||
# Submit to request manager (automatically handles n>1)
|
||||
response_generator = self.request_manager.generate_request(
|
||||
@@ -248,19 +257,15 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
|
||||
logger.info(f"Receive embedding request: {request.request_id}")
|
||||
|
||||
try:
|
||||
# Convert request
|
||||
tokenized_req = self._convert_embed_request(request)
|
||||
|
||||
# Submit to request manager
|
||||
future = await self.request_manager.embedding_request(
|
||||
obj=tokenized_req,
|
||||
request_id=request.request_id,
|
||||
)
|
||||
|
||||
# Wait for result
|
||||
result = await future
|
||||
|
||||
# Create response
|
||||
return sglang_scheduler_pb2.EmbedResponse(
|
||||
request_id=request.request_id,
|
||||
complete=sglang_scheduler_pb2.EmbedComplete(
|
||||
@@ -536,6 +541,25 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer)
|
||||
aggregate=_compute_aggregate_protobuf(loads),
|
||||
)
|
||||
|
||||
def _handle_epd_disaggregation_encode_request(
|
||||
self,
|
||||
grpc_req: sglang_scheduler_pb2.GenerateRequest,
|
||||
tokenized_req: TokenizedGenerateReqInput,
|
||||
) -> None:
|
||||
if not self.mm_receiver:
|
||||
return
|
||||
|
||||
image_urls = list(grpc_req.mm_inputs.image_urls)
|
||||
if not image_urls:
|
||||
return
|
||||
|
||||
encode_req = self.mm_receiver.build_and_send_encode_request(
|
||||
image_urls=image_urls,
|
||||
rid=grpc_req.request_id,
|
||||
)
|
||||
tokenized_req.need_wait_for_image = bool(encode_req.need_wait_for_image)
|
||||
tokenized_req.num_items_assigned = encode_req.num_items_assigned
|
||||
|
||||
# Helper methods for request/response conversion
|
||||
|
||||
def _convert_generate_request(
|
||||
|
||||
Reference in New Issue
Block a user