diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index bde59f96b..9d347025d 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -1,4 +1,5 @@ import asyncio +import concurrent.futures import ctypes import logging import multiprocessing as mp @@ -17,7 +18,6 @@ import zmq.asyncio from fastapi import FastAPI from fastapi.responses import ORJSONResponse, Response from transformers import AutoImageProcessor -from transformers.image_utils import load_images from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig @@ -38,7 +38,14 @@ from sglang.srt.server_args import ( ServerArgs, set_global_server_args_for_scheduler, ) -from sglang.srt.utils import get_local_ip_auto, get_zmq_socket, random_uuid +from sglang.srt.utils import ( + get_local_ip_auto, + get_zmq_socket, + load_audio, + load_image, + load_video, + random_uuid, +) logger = logging.getLogger(__name__) @@ -46,6 +53,10 @@ rid_lock = asyncio.Lock() rid_to_receive_endpoint: Dict[str, List[str]] = dict() rid_to_receive_count: Dict[str, int] = dict() +use_image_processor_gpu = ( + int(os.getenv("SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU", "0")) == 1 +) + class TensorWrapper: """Wrapper to keep tensor alive while exposing buffer for zero-copy.""" @@ -138,6 +149,8 @@ class MMEncoder: 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, @@ -159,6 +172,10 @@ class MMEncoder: 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)) + ) + if schedule_path is not None: self.schedule_socket = get_zmq_socket( self.context, zmq.PULL, schedule_path, True @@ -182,10 +199,102 @@ class MMEncoder: logger.info(f"rank {rank} init finish ") - async def _encode(self, mm_items) -> torch.Tensor: - images = load_images(mm_items) + 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) - images_input = self.image_processor(images=images) + 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: + images = await self._flatten_and_load_images(mm_items) + + 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( {