import asyncio import concurrent.futures import ctypes import logging import multiprocessing as mp import os import pickle import time import traceback from http import HTTPStatus from typing import Dict, List, Optional, Set, Tuple import aiohttp import numpy as np import torch import uvicorn import zmq import zmq.asyncio from fastapi import FastAPI from fastapi.responses import ORJSONResponse, Response from transformers import AutoImageProcessor from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.disaggregation.encode_receiver import EmbeddingData from sglang.srt.distributed.parallel_state import ( get_mooncake_transfer_engine, init_distributed_environment, initialize_model_parallel, ) from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_loader import get_model from sglang.srt.server_args import ( PortArgs, ServerArgs, set_global_server_args_for_scheduler, ) from sglang.srt.utils import ( config_socket, get_local_ip_auto, get_zmq_socket, load_audio, load_image, load_video, random_uuid, ) logger = logging.getLogger(__name__) rid_lock = asyncio.Lock() rid_to_receive_endpoint: Dict[str, List[str]] = dict() rid_to_receive_count: Dict[str, int] = dict() rid_to_err_msg: Dict[str, str] = dict() cond_dict_lock = asyncio.Lock() rid_to_cond: Dict[str, asyncio.Condition] = {} use_image_processor_gpu = ( int(os.getenv("SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU", "0")) == 1 ) class MMError(Exception): def __init__(self, message, code=HTTPStatus.INTERNAL_SERVER_ERROR): self.message = message self.code = code super().__init__(self.message) class BadRequestError(MMError): def __init__(self, message): super().__init__(message, code=HTTPStatus.BAD_REQUEST) class InternalError(MMError): def __init__(self, message): super().__init__(message, code=HTTPStatus.INTERNAL_SERVER_ERROR) class TensorWrapper: """Wrapper to keep tensor alive while exposing buffer for zero-copy.""" def __init__(self, tensor): # Ensure tensor is on CPU and contiguous if tensor.is_cuda: tensor = tensor.cpu() if not tensor.is_contiguous(): tensor = tensor.contiguous() # Keep tensor reference self.tensor = tensor self.shape = list(tensor.shape) self.dtype = tensor.dtype def __buffer__(self): data_ptr = self.tensor.data_ptr() total_bytes = self.tensor.numel() * self.tensor.element_size() c_obj = (ctypes.c_char * total_bytes).from_address(data_ptr) c_obj._keep_alive_ref = self return memoryview(c_obj) def _convert(data): if isinstance(data, torch.Tensor): return data elif isinstance(data, np.ndarray): return torch.tensor(data) elif isinstance(data, list) and isinstance(data[0], np.ndarray): return torch.tensor(np.array(data)) elif isinstance(data, list) and isinstance(data[0], (int, float)): return torch.tensor(data) else: return data _image_grid_attrs = ["image_grid_thw", "image_grid_hws"] def _get_image_grid_dim(images_input): for attr in _image_grid_attrs: if attr in images_input: return images_input[attr] raise ValueError( f"Image grid dim ({_image_grid_attrs}) not found in {images_input}" ) class MMEncoder: def __init__( self, server_args: ServerArgs, schedule_path=None, dist_init_method=None, rank: int = 0, ): logger.info(f"init MMEncoder {rank}/{server_args.tp_size}") self.server_args = server_args set_global_server_args_for_scheduler(server_args) self.rank = rank self.profiler = EncoderProfiler(rank) self.image_processor = AutoImageProcessor.from_pretrained( server_args.model_path, trust_remote_code=server_args.trust_remote_code, use_fast=True, ) self.model_config = ModelConfig.from_server_args( server_args, ) self.load_config = LoadConfig( load_format=server_args.load_format, download_dir=server_args.download_dir, model_loader_extra_config=server_args.model_loader_extra_config, remote_instance_weight_loader_seed_instance_ip=server_args.remote_instance_weight_loader_seed_instance_ip, remote_instance_weight_loader_seed_instance_service_port=server_args.remote_instance_weight_loader_seed_instance_service_port, remote_instance_weight_loader_send_weights_group_ports=server_args.remote_instance_weight_loader_send_weights_group_ports, ) self.device = server_args.device self.gpu_id = server_args.base_gpu_id + rank self.device_config = DeviceConfig( device=self.device, gpu_id=self.gpu_id, ) torch.get_device_module(self.device).set_device(self.gpu_id) self.use_image_processor_gpu = use_image_processor_gpu init_distributed_environment( world_size=server_args.tp_size, rank=rank, distributed_init_method=dist_init_method, local_rank=rank, ) initialize_model_parallel(tensor_model_parallel_size=server_args.tp_size) initialize_dp_attention(server_args, self.model_config) self.model = get_model( model_config=self.model_config, load_config=self.load_config, device_config=self.device_config, ) self.context = zmq.asyncio.Context(2) self.sync_context = zmq.Context() # Reuse sync context for thread pool self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=10) embedding_cache_size = int(os.environ.get("SGLANG_VLM_CACHE_SIZE_MB", "4096")) self.mm_cache = MultiModalStaticCache(embedding_cache_size * 1024 * 1024) self.mm_cache_lock = asyncio.Lock() self.io_executor = concurrent.futures.ThreadPoolExecutor( max_workers=int(os.environ.get("SGLANG_ENCODER_MM_LOAD_WORKERS", 4)) ) self.send_timeout = envs.SGLANG_ENCODER_SEND_TIMEOUT.get() if schedule_path is not None: self.schedule_socket = get_zmq_socket( self.context, zmq.PULL, schedule_path, True ) if self.rank == 0: logger.info( f"Using transfer backend: {self.server_args.encoder_transfer_backend}" ) if self.server_args.encoder_transfer_backend == "mooncake": self.local_ip = get_local_ip_auto() self.engine = get_mooncake_transfer_engine() self.embedding_to_send = dict() self.background_tasks: Set[asyncio.Task] = set() logger.info(f"rank {rank} init finish ") def _load_single_item( self, data, modality: Modality, frame_count_limit=None, audio_sample_rate: Optional[int] = None, discard_alpha_channel=True, ): """ Load a single multimodal data. If data is precomputed, returns directly. Static method that can be pickled for multiprocessing""" if isinstance(data, dict): return data try: if modality == Modality.IMAGE: img, _ = load_image(data) if discard_alpha_channel and img.mode != "RGB": img = img.convert("RGB") return img elif modality == Modality.VIDEO: return load_video(data, frame_count_limit) elif modality == Modality.AUDIO: return load_audio(data, audio_sample_rate) except Exception as e: raise RuntimeError(f"Error while loading data {data}: {e}") def submit_data_loading_tasks(self, items, modalities): futures = [] task_info = [] for data, modality in zip(items, modalities): if modality is not None: futures.append( self.io_executor.submit( self._load_single_item, data, modality, ) ) task_info.append((modality, data)) return futures, task_info async def _flatten_and_load_images(self, mm_items): """ Flatten mm_items structure, load images concurrently, and restore original structure. Returns: Same structure as load_images would return """ # Handle single image (not a list) if not isinstance(mm_items, (list, tuple)): futures, _ = self.submit_data_loading_tasks([mm_items], [Modality.IMAGE]) return await asyncio.wrap_future(futures[0]) # Handle nested list (list of lists) if len(mm_items) > 0 and isinstance(mm_items[0], (list, tuple)): # Flatten nested structure flat_data = [] flat_indices = [] # Track which group each item belongs to for group_idx, image_group in enumerate(mm_items): for item in image_group: flat_data.append(item) flat_indices.append(group_idx) # Submit all tasks concurrently futures, _ = self.submit_data_loading_tasks( flat_data, [Modality.IMAGE] * len(flat_data) ) # Wait for all tasks to complete asynchronously async_futures = [asyncio.wrap_future(f) for f in futures] results = await asyncio.gather(*async_futures) # Restore nested structure nested_results = [[] for _ in range(len(mm_items))] for idx, result in zip(flat_indices, results): nested_results[idx].append(result) return nested_results # Handle simple list else: futures, _ = self.submit_data_loading_tasks( mm_items, [Modality.IMAGE] * len(mm_items) ) # Wait for all tasks to complete asynchronously async_futures = [asyncio.wrap_future(f) for f in futures] return await asyncio.gather(*async_futures) async def _encode(self, mm_items) -> torch.Tensor: try: images = await self._flatten_and_load_images(mm_items) except Exception as e: raise BadRequestError(f"Failed to load images from input: {str(e)}") try: kwargs = {"device": self.device} if self.use_image_processor_gpu else {} images_input = self.image_processor(images=images, **kwargs) feature = images_input["pixel_values"] mm_item = MultimodalDataItem.from_dict( { "modality": Modality.IMAGE, "feature": _convert(feature), } ) for k, v in images_input.items(): if k == "pixel_values": continue mm_item.set(k, _convert(v)) # support mm_cache mm_embedding = None mm_hash = None if self.server_args.enable_prefix_mm_cache: mm_item.set_pad_value() mm_hash = MultiModalStaticCache.combine_hashes([mm_item.hash]) async with self.mm_cache_lock: mm_cache = self.mm_cache.get([mm_item.hash]) if mm_cache is not None: mm_embedding = mm_cache.embedding if mm_embedding is None: with torch.inference_mode(): mm_embedding: torch.Tensor = self.model.get_image_feature([mm_item]) mm_embedding = mm_embedding.cpu() if len(mm_embedding.shape) != 2: mm_embedding = mm_embedding.reshape(-1, mm_embedding.shape[-1]) if self.server_args.enable_prefix_mm_cache: async with self.mm_cache_lock: self.mm_cache.set(mm_hash, EmbeddingResult(embedding=mm_embedding)) if self.profiler is not None: self.profiler.step() return _get_image_grid_dim(images_input), mm_embedding except BadRequestError as e: raise BadRequestError(f"Bad request error: {str(e)}") except Exception as e: raise InternalError(f"Internal encoding error: {str(e)}") async def _send( self, embedding: torch.Tensor, mm_data: EmbeddingData, session_id=None, buffer_address=None, prefill_host=None, embedding_port=None, url=None, ): if self.server_args.encoder_transfer_backend == "mooncake": self.engine.register(embedding.data_ptr(), embedding.nbytes) self.engine.transfer_sync( session_id, embedding.data_ptr(), buffer_address, embedding.nbytes ) self.engine.deregister(embedding.data_ptr()) mm_data.embedding = None mm_data.embedding_list[mm_data.part_idx] = None # Send ack/data endpoint = ( f"tcp://{url}" if url is not None else f"tcp://{prefill_host}:{embedding_port}" ) logger.info(f"{endpoint = }") # Serialize data if self.server_args.encoder_transfer_backend == "mooncake": serialized_data = pickle.dumps(mm_data) buffer = None else: new_mm_data = mm_data.copy_without_embedding() if new_mm_data.error_msg is not None: buffer = None serialized_data = pickle.dumps(new_mm_data) else: embedding_tensor = TensorWrapper(mm_data.embedding) serialized_data = pickle.dumps(new_mm_data) buffer = embedding_tensor.__buffer__() # Use thread pool executor for parallel ZMQ send operations def send_with_socket(): sock = self.sync_context.socket(zmq.PUSH) config_socket(sock, zmq.PUSH) try: sock.connect(endpoint) if buffer is not None: sock.send_multipart([serialized_data, buffer], copy=False) else: sock.send_multipart([serialized_data], copy=False) finally: sock.close() await asyncio.get_event_loop().run_in_executor(self.executor, send_with_socket) async def encode(self, mm_items, req_id, num_parts, part_idx): try: image_grid_dim, mm_embedding = await self._encode(mm_items) if self.rank == 0: mm_data = EmbeddingData( req_id, num_parts, part_idx, image_grid_dim, mm_embedding ) self.embedding_to_send[req_id] = mm_data return ( mm_embedding.nbytes, mm_embedding.shape[0], mm_embedding.shape[1], None, None, ) except Exception as e: error_code = getattr(e, "code", HTTPStatus.INTERNAL_SERVER_ERROR) error_msg = str(e) logger.error(f"Rank {self.rank} encode failed: {error_msg} {error_code = }") if self.rank == 0: mm_data = EmbeddingData( req_id, num_parts, part_idx, None, error_msg=error_msg, error_code=error_code, ) self.embedding_to_send[req_id] = mm_data logger.debug(f"Created error EmbeddingData: {mm_data}") return 0, 0, 0, error_msg, error_code # For zmq_to_tokenizer zmq_to_scheduler and mooncake async def send( self, req_id, prefill_host, embedding_port, session_id=None, buffer_address=None ): mm_data: EmbeddingData = self.embedding_to_send[req_id] await self._send( mm_data.embedding, mm_data, session_id=session_id, buffer_address=buffer_address, prefill_host=prefill_host, embedding_port=embedding_port, ) # For zmq_to_scheduler async def send_with_url( self, req_id, ): mm_data = self.embedding_to_send.get(req_id) if not mm_data: return sent_urls: Set[str] = set() all_tasks: List[Tuple[asyncio.Task, str]] = [] start_time = asyncio.get_running_loop().time() timeout = self.send_timeout cond = await get_condition(req_id) try: while True: async with rid_lock: current_targets = rid_to_receive_endpoint.get(req_id, set()).copy() expected_count = rid_to_receive_count.get(req_id) new_targets = current_targets - sent_urls if new_targets: logger.info( f"Found {len(new_targets)} new endpoints for {req_id}. Starting tasks..." ) for url in new_targets: task = asyncio.create_task( self._send( mm_data.embedding, mm_data, url=url, ) ) all_tasks.append((task, url)) sent_urls.add(url) # Mark as handled immediately if expected_count is not None and len(sent_urls) >= expected_count: logger.info( f"All {expected_count} endpoints initiated for {req_id}. Breaking loop." ) break remaining = timeout - (asyncio.get_running_loop().time() - start_time) if remaining <= 0: logger.error( f"[{req_id}] Timeout! Sent {len(sent_urls)}/{expected_count}" ) break async with cond: try: await asyncio.wait_for(cond.wait(), timeout=remaining) except asyncio.TimeoutError: continue if all_tasks: logger.info( f"Loop finished. Awaiting completion of {len(all_tasks)} sending tasks..." ) tasks_only = [t[0] for t in all_tasks] results = await asyncio.gather(*tasks_only, return_exceptions=True) # Process results and log errors for i, result in enumerate(results): url = all_tasks[i][1] # Retrieve URL associated with the task if isinstance(result, Exception): logger.error(f"Failed to send to {url}: {result}") else: logger.debug(f"Successfully sent to {url}") logger.info(f"All tasks completed for req_id: {req_id}") finally: logger.info(f"Cleaning up resources for req_id {req_id}") async with rid_lock: rid_to_receive_endpoint.pop(req_id, None) rid_to_receive_count.pop(req_id, None) async with cond_dict_lock: rid_to_cond.pop(req_id, None) self.embedding_to_send.pop(req_id, None) async def get_embedding_port(self, prefill_url): async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=1800) ) as session: response = await session.post( f"{prefill_url}/embedding_bootstrap", json={"embedding_port": None}, ) response_json = await response.json() return response_json["embedding_port"] class EncoderProfiler: def __init__(self, rank: int): self.rank = rank self.profiler = None self.steps_left = None self.output_dir = None self.prefix = None self.profile_id = None def start(self, obj: ProfileReq): if self.profiler is not None: return False, "profiling already running" output_dir = obj.output_dir or os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") os.makedirs(output_dir, exist_ok=True) self.output_dir = output_dir self.prefix = obj.profile_prefix or "encoder" self.profile_id = str(time.time()) activities = obj.activities or ["CPU", "GPU"] torch_activities = [] if "CPU" in activities: torch_activities.append(torch.profiler.ProfilerActivity.CPU) if "GPU" in activities: torch_activities.append(torch.profiler.ProfilerActivity.CUDA) profile_memory = "MEM" in activities if not torch_activities and not profile_memory: return False, "no supported activities" self.profiler = torch.profiler.profile( activities=torch_activities, with_stack=True if obj.with_stack is None else obj.with_stack, record_shapes=False if obj.record_shapes is None else obj.record_shapes, profile_memory=profile_memory, ) self.profiler.start() self.steps_left = obj.num_steps logger.info( f"Encoder profiling started. output_dir={self.output_dir} profile_id={self.profile_id}" ) return True, None def step(self): if self.profiler is None: return self.profiler.step() if self.steps_left is not None: self.steps_left -= 1 if self.steps_left <= 0: self.stop() def stop(self): if self.profiler is None: return False, "profiling not running" self.profiler.stop() filename = f"{self.prefix}-rank{self.rank}-{self.profile_id}.trace.json" trace_path = os.path.join(self.output_dir, filename) self.profiler.export_chrome_trace(trace_path) logger.info("Encoder profiling saved to: %s", trace_path) self.profiler = None self.steps_left = None return True, None app = FastAPI() encoder: Optional[MMEncoder] = None send_sockets: List[zmq.Socket] = [] async def run_encoder( server_args: ServerArgs, schedule_path, dist_init_method, rank: int ): encoder = MMEncoder(server_args, schedule_path, dist_init_method, rank) while True: request = await encoder.schedule_socket.recv_pyobj() if isinstance(request, ProfileReq): if request.type == ProfileReqType.START_PROFILE: if encoder.profiler is None: encoder.profiler = EncoderProfiler(encoder.rank) encoder.profiler.start(request) else: encoder.profiler.stop() else: await encoder.encode( mm_items=request["mm_items"], req_id=request["req_id"], num_parts=request["num_parts"], part_idx=request["part_idx"], ) def launch_encoder(server_args, schedule_path, dist_init_method, rank): try: asyncio.run(run_encoder(server_args, schedule_path, dist_init_method, rank)) except KeyboardInterrupt: logger.info(f"Exit rank {rank}") except Exception: traceback.print_exc() def launch_server(server_args: ServerArgs): global encoder ctx = mp.get_context("spawn") zmq_ctx = zmq.Context(10) ipc_path_prefix = random_uuid() port_args = PortArgs.init_new(server_args) if server_args.dist_init_addr: dist_init_method = f"tcp://{server_args.dist_init_addr}" else: dist_init_method = f"tcp://127.0.0.1:{port_args.nccl_port}" 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) uvicorn.run(app, host=server_args.host, port=server_args.port) async def get_condition(rid): async with cond_dict_lock: if rid not in rid_to_cond: rid_to_cond[rid] = asyncio.Condition() return rid_to_cond[rid] @app.post("/encode") async def handle_encode_request(request: dict): req_id = request["req_id"] try: def start_background_send(req_id): task = asyncio.create_task(encoder.send_with_url(req_id=req_id)) encoder.background_tasks.add(task) task.add_done_callback(encoder.background_tasks.discard) # broadcast request request.update({"enter_time": time.time()}) for socket in send_sockets: socket.send_pyobj(request) nbytes, embedding_len, embedding_dim, error_msg, error_code = ( await encoder.encode( mm_items=request["mm_items"], req_id=request["req_id"], num_parts=request["num_parts"], part_idx=request["part_idx"], ) ) if error_msg: if encoder.server_args.encoder_transfer_backend == "zmq_to_scheduler": if request["embedding_port"] is None: start_background_send(req_id) else: for port in request["embedding_port"]: await encoder.send( req_id=req_id, prefill_host=request["prefill_host"], embedding_port=port, ) return ORJSONResponse( status_code=error_code, content={"status": "error", "message": error_msg, "req_id": req_id}, ) if encoder.server_args.encoder_transfer_backend == "mooncake": del request["mm_items"] request.update( { "embedding_size": nbytes, "embedding_len": embedding_len, "embedding_dim": embedding_dim, } ) return ORJSONResponse(content=request) elif encoder.server_args.encoder_transfer_backend == "zmq_to_scheduler": logger.info(f"{request['embedding_port'] = }") if request["embedding_port"] is None: await encoder.send_with_url( req_id=request["req_id"], ) else: assert type(request["embedding_port"]) == list tasks = [] for embedding_port in request["embedding_port"]: tasks.append( encoder.send( req_id=request["req_id"], prefill_host=request["prefill_host"], embedding_port=embedding_port, ) ) await asyncio.gather(*tasks) encoder.embedding_to_send.pop(request["req_id"], None) return ORJSONResponse(content=None) elif encoder.server_args.encoder_transfer_backend == "zmq_to_tokenizer": await encoder.send( req_id=request["req_id"], prefill_host=request["prefill_host"], embedding_port=request["embedding_port"], ) encoder.embedding_to_send.pop(request["req_id"], None) return ORJSONResponse(content=None) except Exception as e: error_msg = str(e) logger.error(f"Unexpected error in encoder logic for {req_id}: {error_msg}") rid_to_err_msg[req_id] = error_msg return ORJSONResponse( status_code=HTTPStatus.INTERNAL_SERVER_ERROR, content={ "status": "error", "message": error_msg, "req_id": req_id, }, ) @app.post("/send") async def handle_send_request(request: dict): # mooncake backend await encoder.send( req_id=request["req_id"], prefill_host=request["prefill_host"], embedding_port=request["embedding_port"], session_id=request["session_id"], buffer_address=request["buffer_address"], ) encoder.embedding_to_send.pop(request["req_id"], None) return ORJSONResponse(content=None) @app.post("/scheduler_receive_url") async def handle_scheduler_receive_url_request(request: dict): rid = request["req_id"] async with rid_lock: global rid_to_receive_endpoint if rid not in rid_to_receive_endpoint: rid_to_receive_endpoint[rid] = set() rid_to_receive_count[rid] = request["receive_count"] assert rid_to_receive_count[rid] == request["receive_count"] rid_to_receive_endpoint[rid].add(request["receive_url"]) cond = await get_condition(rid) async with cond: cond.notify_all() @app.get("/health") @app.get("/health_generate") async def health_generate(): """ Health check endpoint for the encoder server. Returns 200 if the encoder is initialized and ready. """ if encoder is None: return Response(status_code=503) return Response(status_code=200) @app.api_route("/start_profile", methods=["GET", "POST"]) async def start_profile_async(obj: Optional[ProfileReqInput] = None): if encoder is None: return Response(content="encoder not ready\n", status_code=503) req = None if obj is None: req = ProfileReq(ProfileReqType.START_PROFILE) else: req = ProfileReq( type=ProfileReqType.START_PROFILE, output_dir=obj.output_dir, start_step=obj.start_step, num_steps=obj.num_steps, activities=obj.activities, with_stack=obj.with_stack, record_shapes=obj.record_shapes, profile_by_stage=obj.profile_by_stage, profile_id=str(time.time()), merge_profiles=obj.merge_profiles, profile_prefix=obj.profile_prefix, profile_stages=obj.profile_stages, ) for socket in send_sockets: socket.send_pyobj(req) if encoder.profiler is None: encoder.profiler = EncoderProfiler(encoder.rank) ok, msg = encoder.profiler.start(req) if ok: detail = ( f"Start profiling. output_dir={encoder.profiler.output_dir} " f"profile_id={encoder.profiler.profile_id}\n" ) return Response(content=detail, status_code=200) return Response( content=(msg or "Start profiling failed.\n"), status_code=HTTPStatus.BAD_REQUEST ) @app.api_route("/stop_profile", methods=["GET", "POST"]) async def stop_profile_async(): if encoder is None: return Response(content="encoder not ready\n", status_code=503) if encoder.profiler is None: return Response( content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST ) req = ProfileReq(ProfileReqType.STOP_PROFILE) for socket in send_sockets: socket.send_pyobj(req) ok, msg = encoder.profiler.stop() if ok: return Response(content="Stop profiling.\n", status_code=200) return Response( content=(msg or "Stop profiling failed.\n"), status_code=HTTPStatus.BAD_REQUEST )