[Feature] support async download features (#5003)

* support async download features

* add test case

* update code

Author: kevincheng2

fastdeploy/config.py

@@ -550,6 +550,8 @@ def __init__(
         self.use_internode_ll_two_stage: bool = False
         # disable sequence parallel moe
         self.disable_sequence_parallel_moe: bool = False
+        # enable async download features
+        self.enable_async_download_features: bool = False
 
         self.pod_ip: str = None
         # enable the custom all-reduce kernel and fall back to NCCL(dist.all_reduce).

fastdeploy/engine/args_utils.py

@@ -467,6 +467,11 @@ class EngineArgs:
     Url for router server, such as `0.0.0.0:30000`.
     """
 
+    enable_async_download_features: bool = False
+    """
+    Flag to enable async download features. Default is False (disabled).
+    """
+
     def __post_init__(self):
         """
         Post-initialization processing to set default tokenizer if not provided.
@@ -849,6 +854,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.enable_expert_parallel,
             help="Enable expert parallelism.",
         )
+        parallel_group.add_argument(
+            "--enable-async-download-features",
+            action="store_true",
+            default=EngineArgs.enable_async_download_features,
+            help="Enable async download features.",
+        )
 
         # Load group
         load_group = parser.add_argument_group("Load Configuration")

fastdeploy/engine/common_engine.py

@@ -51,14 +51,7 @@
 from fastdeploy.splitwise.splitwise_connector import SplitwiseConnector
 from fastdeploy.trace.constants import LoggingEventName
 from fastdeploy.trace.trace_logger import print as trace_print
-from fastdeploy.utils import (
-    EngineError,
-    check_download_links,
-    envs,
-    get_logger,
-    init_bos_client,
-    llm_logger,
-)
+from fastdeploy.utils import EngineError, envs, get_logger, llm_logger
 
 try:
     TokenProcessor = load_token_processor_plugins()
@@ -808,7 +801,7 @@ def _fetch_request():
                             else:
                                 raise
                 # 2. Schedule requests
-                tasks = self.resource_manager.schedule()
+                tasks, error_tasks = self.resource_manager.schedule()
 
                 # 3. Send to engine
                 if tasks:
@@ -833,7 +826,16 @@ def _fetch_request():
                         trace_print(LoggingEventName.REQUEST_SCHEDULE_END, task.request_id, getattr(task, "user", ""))
                         trace_print(LoggingEventName.INFERENCE_START, task.request_id, getattr(task, "user", ""))
                     self.engine_worker_queue.put_tasks((tasks, self.resource_manager.real_bsz))
-                else:
+
+                # 4. Response error tasks
+                if error_tasks:
+                    for request_id, failed in error_tasks:
+                        if failed is None:
+                            llm_logger.warning(f"Request {request_id} has no error, skip sending error response.")
+                            continue
+                        self._send_error_response(request_id, failed)
+
+                if not tasks and not error_tasks:
                     time.sleep(0.005)
 
             except RuntimeError as e:
@@ -909,24 +911,6 @@ def _insert_zmq_task_to_scheduler(self):
                             self.llm_logger.error(f"Receive request error: {err_msg}")
                             results.append((request.request_id, err_msg))
 
-                    if self._has_features_info(request) and err_msg is None:
-                        if self.bos_client is None:
-                            self.bos_client = init_bos_client()
-
-                        download_urls = []
-                        inputs = request.multimodal_inputs
-                        if inputs.get("video_feature_urls") is not None:
-                            download_urls.extend(inputs.get("video_feature_urls"))
-                        if inputs.get("image_feature_urls") is not None:
-                            download_urls.extend(inputs.get("image_feature_urls"))
-                        if inputs.get("audio_feature_urls") is not None:
-                            download_urls.extend(inputs.get("audio_feature_urls"))
-
-                        err_msg = check_download_links(self.bos_client, download_urls)
-                        if err_msg:
-                            llm_logger.error(f"Receive request {request.request_id} download error: {err_msg}")
-                            results.append((request.request_id, err_msg))
-
                     if err_msg is None:
                         insert_task.append(request)
 
@@ -948,21 +932,27 @@ def _insert_zmq_task_to_scheduler(self):
                         main_process_metrics.num_requests_waiting.inc(1)
                         continue
 
-                    error_result = RequestOutput(
-                        request_id=request_id,
-                        finished=True,
-                        error_code=500,
-                        error_msg=failed,
-                    )
-                    # Since the request is not in scheduler
-                    # Send result by zmq directly
-                    self.send_response_server.send_response(request_id, [error_result])
+                    self._send_error_response(request_id, failed)
             except Exception as e:
                 self.llm_logger.error(
                     f"Error happened while receiving new request from zmq, details={e}, "
                     f"traceback={traceback.format_exc()}"
                 )
 
+    def _send_error_response(self, request_id, error_msg, error_code: int = 500):
+        llm_logger.error(
+            f"Send error response to client, request_id: {request_id}, error_msg: {error_msg}, error_code: {error_code}"
+        )
+        error_result = RequestOutput(
+            request_id=request_id,
+            finished=True,
+            error_code=error_code,
+            error_msg=error_msg,
+        )
+        # Since the request is not in scheduler
+        # Send result by zmq directly
+        self.send_response_server.send_response(request_id, [error_result])
+
     def _decode_token(self, token_ids, req_id, is_end):
         delta_text = ""
         if envs.FD_ENABLE_RETURN_TEXT:
@@ -977,19 +967,6 @@ def _decode_token(self, token_ids, req_id, is_end):
                 del self.data_processor.decode_status[req_id]
         return delta_text, token_ids
 
-    def _has_features_info(self, task):
-        inputs = task.multimodal_inputs
-        if inputs is None or len(inputs) == 0:
-            return False
-
-        if (
-            (inputs.get("video_feature_urls") is not None and len(inputs["video_feature_urls"]) > 0)
-            or (inputs.get("image_feature_urls") is not None and len(inputs["image_feature_urls"]) > 0)
-            or (inputs.get("audio_feature_urls") is not None and len(inputs["audio_feature_urls"]) > 0)
-        ):
-            return True
-        return False
-
     def _zmq_send_generated_tokens(self):
         """
         Recieve output for zmq

fastdeploy/engine/request.py

@@ -173,6 +173,10 @@ def __init__(
         # dp
         self.dp_rank = dp_rank
 
+        self.async_process_futures = []
+        self.error_message = None
+        self.error_code = None
+
     @classmethod
     def from_dict(cls, d: dict):
         data_processor_logger.debug(f"{d}")

fastdeploy/engine/sched/resource_manager_v1.py

@@ -44,7 +44,7 @@
 from fastdeploy.metrics.metrics import main_process_metrics
 from fastdeploy.multimodal.hasher import MultimodalHasher
 from fastdeploy.platforms import current_platform
-from fastdeploy.utils import llm_logger
+from fastdeploy.utils import download_from_bos, init_bos_client, llm_logger
 
 
 @dataclass
@@ -195,6 +195,9 @@ def __init__(self, max_num_seqs, config, tensor_parallel_size, splitwise_role, l
             max_processor_cache_in_bytes = int(config.cache_config.max_processor_cache * 1024 * 1024 * 1024)
             self.processor_cache = ProcessorCacheManager(max_processor_cache_in_bytes)
 
+        self.bos_client = None
+        self.async_preprocess_pool = ThreadPoolExecutor(max_workers=4)
+
     def allocated_slots(self, request: Request):
         return len(request.block_tables) * self.config.cache_config.block_size
 
@@ -500,6 +503,7 @@ def schedule(self):
         with self.lock:
             scheduled_reqs: list[Request] = []
             preempted_reqs: list[Request] = []
+            error_reqs: list[tuple[str, str]] = []
             token_budget = self.config.scheduler_config.max_num_batched_tokens
 
             # First, schedule the RUNNING requests.
@@ -629,6 +633,7 @@ def _allocate_decode_and_extend():
                 req_index += 1
             # schedule the WAITING requests.
             if not preempted_reqs:
+                skip_requests: list[Request] = []
                 while self.waiting and token_budget > 0:
                     if len(self.running) == self.max_num_seqs:
                         break
@@ -639,6 +644,17 @@ def _allocate_decode_and_extend():
                     ):
                         break
                     if request.status == RequestStatus.WAITING:
+                        result = self._waiting_async_process(request)
+                        if result is None:
+                            error_reqs.append((request.request_id, request.error_message))
+                            self.waiting.popleft()
+                            continue
+                        elif result is True:
+                            # skip current request, try next request
+                            skip_requests.append(request)
+                            self.waiting.popleft()
+                            continue
+
                         self._update_mm_hashes(request)
                         # Enable prefix caching
                         if self.config.cache_config.enable_prefix_caching:
@@ -725,12 +741,102 @@ def _allocate_decode_and_extend():
                     else:
                         llm_logger.error("Unknown request status type")
 
+                for req in skip_requests:
+                    # move waiting request to end of the deque
+                    self.waiting.append(req)
+
             if scheduled_reqs:
                 llm_logger.debug(f"schedued_reqs: {scheduled_reqs}")
 
             self.update_metrics()
 
-            return scheduled_reqs
+            return scheduled_reqs, error_reqs
+
+    def _waiting_async_process(self, request: Request) -> None:
+        """
+        Check if async preprocessing is complete for a request.
+        Args:
+            request: The request to check
+        Returns:
+            None: If an error occurred during preprocessing
+            True: If preprocessing is still in progress (request should be skipped)
+            False: If preprocessing is complete (request can be scheduled)
+        """
+        for future in request.async_process_futures:
+            if future.done():
+                if request.get("error_message") is not None:
+                    return None
+            else:
+                return True
+        request.async_process_futures = []
+        return False
+
+    def _apply_async_preprocess(self, request: Request) -> None:
+        request.async_process_futures.append(self.async_preprocess_pool.submit(self._download_features, request))
+
+    def _has_features_info(self, task):
+        inputs = task.multimodal_inputs
+        if inputs is None or len(inputs) == 0:
+            return False
+
+        if (
+            (inputs.get("video_feature_urls") is not None and len(inputs["video_feature_urls"]) > 0)
+            or (inputs.get("image_feature_urls") is not None and len(inputs["image_feature_urls"]) > 0)
+            or (inputs.get("audio_feature_urls") is not None and len(inputs["audio_feature_urls"]) > 0)
+        ):
+            return True
+        return False
+
+    def _download_features(self, request: Request) -> None:
+        """
+        download multimodal features from bos
+        Note:
+            1. this function will be add features for request.multimodal_inputs
+            2. this function maybe update request.error_message and request.error_code
+        Args:
+            request (Request): request object
+        """
+
+        def download_bos_features(bos_client, features_urls):
+            result_list = []
+            for status, feature in download_from_bos(self.bos_client, features_urls):
+                if status:
+                    llm_logger.info(f"request {request.request_id} async download feature: {feature.shape}")
+                    result_list.append(feature)
+                else:
+                    error_msg = f"request {request.request_id} download features error: {feature}"
+                    llm_logger.error(error_msg)
+                    return error_msg
+            return result_list
+
+        if not self.config.parallel_config.enable_async_download_features or not self._has_features_info(request):
+            return None
+
+        if self.bos_client is None:
+            self.bos_client = init_bos_client()
+
+        inputs = request.multimodal_inputs
+        if inputs.get("video_feature_urls") is not None and len(inputs["video_feature_urls"]) > 0:
+            result = download_bos_features(self.bos_client, inputs["video_feature_urls"])
+            if isinstance(result, str):  # download error
+                request.error_message = result
+                request.error_code = 530
+                return None
+            inputs["video_features"] = result
+        if inputs.get("image_feature_urls") is not None and len(inputs["image_feature_urls"]) > 0:
+            result = download_bos_features(self.bos_client, inputs["image_feature_urls"])
+            if isinstance(result, str):  # download error
+                request.error_message = result
+                request.error_code = 530
+                return None
+            inputs["image_features"] = result
+        if inputs.get("audio_feature_urls") is not None and len(inputs["audio_feature_urls"]) > 0:
+            result = download_bos_features(self.bos_client, inputs["audio_feature_urls"])
+            if isinstance(result, str):  # download error
+                request.error_message = result
+                request.error_code = 530
+                return None
+            inputs["audio_features"] = result
 
     def get_available_position(self) -> int:
         position = 0
@@ -788,6 +894,7 @@ def get_prefix_cached_blocks(self, request: Request):
 
     def add_request(self, request: Request) -> None:
         with self.lock:
+            self._apply_async_preprocess(request)
             self.waiting.append(request)
             self.requests[request.request_id] = request