""" gRPC Encoder Server for SGLang EPD (Encode-Prefill-Decode) mode. This server provides gRPC-based encoding for multimodal inputs. Usage: python -m sglang.launch_server --model-path --encoder-only --grpc-mode """ import asyncio import logging import multiprocessing as mp import traceback from concurrent import futures from typing import List import grpc import zmq import zmq.asyncio from grpc_health.v1 import health_pb2, health_pb2_grpc from grpc_reflection.v1alpha import reflection from smg_grpc_proto import sglang_encoder_pb2, sglang_encoder_pb2_grpc from sglang.srt.disaggregation.encode_server import ( MMEncoder, handle_scheduler_receive_url_request, launch_encoder, ) from sglang.srt.managers.io_struct import async_sock_send, wrap_as_pickle from sglang.srt.managers.schedule_batch import Modality from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import random_uuid from sglang.srt.utils.network import NetworkAddress, get_zmq_socket logger = logging.getLogger(__name__) SGLangEncoderServicer = sglang_encoder_pb2_grpc.SglangEncoderServicer add_SGLangEncoderServicer_to_server = ( sglang_encoder_pb2_grpc.add_SglangEncoderServicer_to_server ) class EncoderHealthServicer(health_pb2_grpc.HealthServicer): """ Standard gRPC health check service for encoder server. Implements grpc.health.v1.Health for Kubernetes probes. """ OVERALL_SERVER = "" ENCODER_SERVICE = "sglang.grpc.encoder.SglangEncoder" def __init__(self): self._serving = False def set_serving(self): self._serving = True def set_not_serving(self): self._serving = False async def Check(self, request, context) -> health_pb2.HealthCheckResponse: if self._serving: return health_pb2.HealthCheckResponse( status=health_pb2.HealthCheckResponse.SERVING ) return health_pb2.HealthCheckResponse( status=health_pb2.HealthCheckResponse.NOT_SERVING ) async def Watch(self, request, context): yield await self.Check(request, context) class SGLangEncoderServer(SGLangEncoderServicer): """ gRPC service implementation for SGLang encoder. """ def __init__( self, encoder: MMEncoder, send_sockets: List[zmq.Socket], server_args: ServerArgs, ): self.encoder = encoder self.send_sockets = send_sockets self.server_args = server_args async def Encode( self, request: sglang_encoder_pb2.EncodeRequest, context ) -> sglang_encoder_pb2.EncodeResponse: try: request_dict = { "mm_items": list(request.mm_items), "req_id": request.req_id, "num_parts": request.num_parts, "part_idx": request.part_idx, } for socket in self.send_sockets: await async_sock_send(socket, wrap_as_pickle(request_dict)) # gRPC encode is image-only; encoder.encode() requires modality ( nbytes, embedding_len, embedding_dim, error_msg, error_code, ) = await self.encoder.encode( mm_items=list(request.mm_items), modality=Modality.IMAGE, req_id=request.req_id, num_parts=request.num_parts, part_idx=request.part_idx, ) if error_msg is not None: context.set_code(grpc.StatusCode.INTERNAL) context.set_details(error_msg) return sglang_encoder_pb2.EncodeResponse() if self.server_args.encoder_transfer_backend == "mooncake": return sglang_encoder_pb2.EncodeResponse( embedding_size=nbytes, embedding_len=embedding_len, embedding_dim=embedding_dim, ) elif self.server_args.encoder_transfer_backend == "zmq_to_scheduler": embedding_ports = list(request.embedding_port) logger.info(f"embedding_port = {embedding_ports}") if not embedding_ports: await self.encoder.send_with_url(req_id=request.req_id) else: tasks = [] for embedding_port in embedding_ports: tasks.append( self.encoder.send( req_id=request.req_id, prefill_host=request.prefill_host, embedding_port=embedding_port, ) ) await asyncio.gather(*tasks) self.encoder.embedding_to_send.pop(request.req_id, None) return sglang_encoder_pb2.EncodeResponse() elif self.server_args.encoder_transfer_backend == "zmq_to_tokenizer": embedding_port = ( request.embedding_port[0] if request.embedding_port else 0 ) await self.encoder.send( req_id=request.req_id, prefill_host=request.prefill_host, embedding_port=embedding_port, ) self.encoder.embedding_to_send.pop(request.req_id, None) return sglang_encoder_pb2.EncodeResponse() return sglang_encoder_pb2.EncodeResponse() except Exception as e: logger.error(f"Encode error: {e}") traceback.print_exc() context.set_code(grpc.StatusCode.INTERNAL) context.set_details(str(e)) return sglang_encoder_pb2.EncodeResponse() async def Send( self, request: sglang_encoder_pb2.SendRequest, context ) -> sglang_encoder_pb2.SendResponse: try: await self.encoder.send( req_id=request.req_id, prefill_host=request.prefill_host, embedding_port=request.embedding_port, session_id=request.session_id if request.session_id else None, buffer_address=( request.buffer_address if request.buffer_address else None ), ) self.encoder.embedding_to_send.pop(request.req_id, None) return sglang_encoder_pb2.SendResponse() except Exception as e: logger.error(f"Send error: {e}") traceback.print_exc() context.set_code(grpc.StatusCode.INTERNAL) context.set_details(str(e)) return sglang_encoder_pb2.SendResponse() async def SchedulerReceiveUrl( self, request: sglang_encoder_pb2.SchedulerReceiveUrlRequest, context ) -> sglang_encoder_pb2.SchedulerReceiveUrlResponse: try: await handle_scheduler_receive_url_request( { "req_id": request.req_id, "receive_count": request.receive_count, "receive_url": request.receive_url, } ) return sglang_encoder_pb2.SchedulerReceiveUrlResponse() except Exception as e: logger.error(f"SchedulerReceiveUrl error: {e}") traceback.print_exc() context.set_code(grpc.StatusCode.INTERNAL) context.set_details(str(e)) return sglang_encoder_pb2.SchedulerReceiveUrlResponse() async def serve_grpc_encoder(server_args: ServerArgs): ctx = mp.get_context("spawn") zmq_ctx = zmq.asyncio.Context(10) ipc_path_prefix = random_uuid() port_args = PortArgs.init_new(server_args) if server_args.dist_init_addr: na = NetworkAddress.parse(server_args.dist_init_addr) dist_init_method = na.to_tcp() else: dist_init_method = NetworkAddress( server_args.host or "127.0.0.1", port_args.nccl_port ).to_tcp() send_sockets: List[zmq.Socket] = [] for rank in range(1, server_args.tp_size): schedule_path = f"ipc:///tmp/{ipc_path_prefix}_schedule_{rank}" send_sockets.append( get_zmq_socket(zmq_ctx, zmq.PUSH, schedule_path, bind=False) ) ctx.Process( target=launch_encoder, args=(server_args, schedule_path, dist_init_method, rank), daemon=True, ).start() encoder = MMEncoder(server_args, dist_init_method=dist_init_method) server = grpc.aio.server( futures.ThreadPoolExecutor(max_workers=10), options=[ ("grpc.max_send_message_length", 1024 * 1024 * 256), ("grpc.max_receive_message_length", 1024 * 1024 * 256), ], ) health_servicer = EncoderHealthServicer() health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) encoder_servicer = SGLangEncoderServer( encoder=encoder, send_sockets=send_sockets, server_args=server_args, ) add_SGLangEncoderServicer_to_server(encoder_servicer, server) SERVICE_NAMES = ( sglang_encoder_pb2.DESCRIPTOR.services_by_name["SglangEncoder"].full_name, "grpc.health.v1.Health", reflection.SERVICE_NAME, ) reflection.enable_server_reflection(SERVICE_NAMES, server) listen_addr = NetworkAddress(server_args.host, server_args.port).to_host_port_str() server.add_insecure_port(listen_addr) await server.start() logger.info(f"gRPC encoder server listening on {listen_addr}") health_servicer.set_serving() try: await server.wait_for_termination() except KeyboardInterrupt: logger.info("Shutting down gRPC encoder server...") health_servicer.set_not_serving() await server.stop(grace=5)