[Refactor]

Signed-off-by: LastZhabka <sakhmoldin.mukhammadarif@gmail.com>

Author: LastZhabka

examples/online_serving/separated_encode/proxy/proxy1e1pd_aiohttp.py

@@ -45,31 +45,45 @@ async def shutdown_event():
     if decode_session:
         await decode_session.close()
 
+
+def has_mm_input(request_data: dict):
+    if "messages" not in request_data:
+        return False
+    for message in request_data["messages"]:  
+        if not isinstance(message.get("content"), list):  
+            continue
+        for content_item in message["content"]:  
+            if content_item.get("type") in ["image_url", "audio_url", "input_audio"]:  
+                return True 
+    return False
+
 async def forward_streaming_request(
     request_data: dict,
     request_id: str
 ) -> AsyncIterator[str]:
     headers = {"x-request-id": request_id}
-    task1 = asyncio.create_task(
-        encode_session.post(
-            f"{ENCODE_SERVER_URL}/v1/chat/completions",
-            json=request_data,
-            headers=headers
+    # Skip request to encoder instance if we don't have mm input
+    if has_mm_input(request_data):
+        task1 = asyncio.create_task(
+            encode_session.post(
+                f"{ENCODE_SERVER_URL}/v1/chat/completions",
+                json=request_data,
+                headers=headers
+            )
         )
-    )
-    try:
-        response = await task1
-        if response.status != 200:
-            error_text = await response.text()
+        try:
+            response = await task1
+            if response.status != 200:
+                error_text = await response.text()
+                raise HTTPException(
+                    status_code=response.status,
+                    detail={"error": "Request failed", "message": error_text}
+                )
+        except Exception as e:
             raise HTTPException(
-                status_code=response.status,
-                detail={"error": "Request failed", "message": error_text}
+                status_code=500,
+                detail={"error": "Internal server error", "message": str(e)}
             )
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail={"error": "Internal server error", "message": str(e)}
-        )
 
     try:
         async with decode_session.post(
@@ -83,37 +97,37 @@ async def forward_streaming_request(
                     yield chunk.decode('utf-8', errors='ignore')
     except Exception as e:
         logger.error(f"Error in streaming: {e}")
-        task1.cancel()
         raise
 
 async def forward_non_streaming_request(
     request_data: dict,
     request_id: str
 ) -> dict:
     headers = {"x-request-id": request_id}
-    
-    # Start request to encode server
-    task1 = asyncio.create_task(
-        encode_session.post(
-            f"{ENCODE_SERVER_URL}/v1/chat/completions",
-            json=request_data,
-            headers=headers
+    # Skip request to encoder instance if we don't have mm input
+    if has_mm_input(request_data):
+        # Start request to encode server
+        task1 = asyncio.create_task(
+            encode_session.post(
+                f"{ENCODE_SERVER_URL}/v1/chat/completions",
+                json=request_data,
+                headers=headers
+            )
         )
-    )
 
-    try:
-        response = await task1
-        if response.status != 200:
-            error_text = await response.text()
+        try:
+            response = await task1
+            if response.status != 200:
+                error_text = await response.text()
+                raise HTTPException(
+                    status_code=response.status,
+                    detail={"error": "Request failed", "message": error_text}
+                )
+        except Exception as e:
             raise HTTPException(
-                status_code=response.status,
-                detail={"error": "Request failed", "message": error_text}
+                status_code=500,
+                detail={"error": "Internal server error", "message": str(e)}
             )
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail={"error": "Internal server error", "message": str(e)}
-        )
 
     try:
         # Make request to decode server
@@ -127,7 +141,6 @@ async def forward_non_streaming_request(
         return result
     except Exception as e:
         logger.error(f"Error in non-streaming: {e}")
-        task1.cancel()
         raise
 
 @app.post("/v1/chat/completions")

vllm/separated_encode/README.md

@@ -3,32 +3,33 @@
 from typing import Callable, Literal, Optional
 
 import msgpack_numpy
-import numpy as np
 import redis
-from numpy.typing import NDArray
 
 from vllm.config import VllmConfig
 from vllm.separated_encode.ec_transfer.connector.template import (
     ECConnectorTemplate)
 from vllm.logger import init_logger
+import torch
 
 logger = init_logger(__name__)
 
 class RedisECConnector(ECConnectorTemplate):
 
     def __init__(self,
                  vllm_config: "VllmConfig",
+                 device: Optional[torch.device],
                  intra_instance_type: Literal["scheduler", "model-runner"],
                  preallocate_callback: Optional[Callable[[str, int, int, str],
                                                          None]],
                  injection_callback: Optional[Callable[
-                     [str, int, NDArray[np.float32], str], None]],
+                     [str, int, torch.Tensor, str], None]],
                  redis_host: str = "localhost",
                  redis_port: int = 6379):
         self.redis_client = redis.StrictRedis(host=redis_host, port=redis_port)
         self.rank = vllm_config.epd_disagg_config.epd_rank
         super().__init__(
             vllm_config,
+            device,
             intra_instance_type,
             preallocate_callback,
             injection_callback,
@@ -71,12 +72,13 @@ def _send_encoder_cache_metas(
 
     def _send_encoder_cache(
         self, request_id: str, input_id: int,
-        encoder_cache: NDArray[np.float32], mm_hash: str) -> None:
+        encoder_cache: torch.Tensor, mm_hash: str) -> None:
         # E -> PD
+        encoder_cache_numpy = encoder_cache.to("cpu", dtype=torch.float16).numpy()
         transfer_data = msgpack_numpy.packb({
             "request_id": request_id,
             "input_id": input_id,
-            "encoder_cache": encoder_cache,
+            "encoder_cache": encoder_cache_numpy,
             "mm_hash": mm_hash
         })
         rank = self._get_request_ranks(request_id)[1]
@@ -113,7 +115,7 @@ def _recv_encoder_cache_metas(
 
     def _recv_encoder_cache(
         self, 
-        injection_callback: Callable[[str, int, NDArray[np.float32], str],None]
+        injection_callback: Callable[[str, int, torch.Tensor, str],None]
     ) -> None:
         transfered_data = self.redis_client.blpop(f"cache{self.rank}")[1]
         transfered_data = msgpack_numpy.unpackb(transfered_data, raw=False)
@@ -123,5 +125,7 @@ def _recv_encoder_cache(
             transfered_data["encoder_cache"],
             transfered_data["mm_hash"]
         )
+        encoder_cache = torch.from_numpy(encoder_cache).to(
+                device=self.device, dtype=self.dtype)   
         logger.debug(f"Received encoder cache -> {self.rank}, {request_id}")
         injection_callback(request_id, input_id, encoder_cache, mm_hash)

vllm/separated_encode/ec_transfer/connector/redis.py

@@ -6,8 +6,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Callable, Literal, Optional
 
-import numpy as np
-from numpy.typing import NDArray
+import torch
 
 from vllm.config import EPDDisaggConfig, VllmConfig
 from vllm.logger import init_logger
@@ -37,9 +36,10 @@ class ECConnectorTemplate(ABC):
     def __init__(
         self,
         vllm_config: "VllmConfig",
+        device: Optional[torch.device],
         intra_instance_type: Literal["scheduler", "model-runner"],
         preallocate_callback: Optional[Callable[[str, int, int, str], None]],
-        injection_callback: Optional[Callable[[str, int, NDArray[np.float32], str],
+        injection_callback: Optional[Callable[[str, int, torch.Tensor, str],
                                               None]],
     ):
         callback_mapping = {
@@ -55,12 +55,14 @@ def __init__(
             ("prefill+decode", "model-runner"): (self._recv_encoder_cache,
                                                  injection_callback)
         }
+        self.device = device
+        self.dtype = vllm_config.model_config.dtype
 
         self.epd_disagg_config: EPDDisaggConfig
         self.intra_instance_type: Literal["scheduler", "model-runner"]
         self.inter_instance_type: Literal["encode", "prefill",
                                           "prefill+decode"]
-        self.encoder_cache: dict[str, dict[int, NDArray[np.float32]]]
+        self.encoder_cache: dict[str, dict[int, torch.Tensor]]
         self.send_executors: ThreadPoolExecutor
         self.recv_executors: ThreadPoolExecutor
 
@@ -143,7 +145,7 @@ def _send_encoder_cache_metas(self, request_id: str, input_id: int,
     @abstractmethod
     def _send_encoder_cache(
         self, request_id: str, input_id: int,
-        encoder_cache: NDArray[np.float32], mm_hash: str
+        encoder_cache: torch.Tensor, mm_hash: str
     ) -> None:
         """Send the encoder cache.
 
@@ -204,7 +206,7 @@ def _recv_encoder_cache_metas(
     @abstractmethod
     def _recv_encoder_cache(
         self, 
-        injection_callback: Callable[[str, int, NDArray[np.float32], str],None]
+        injection_callback: Callable[[str, int, torch.Tensor, str],None]
     ) -> None:
         """Receives the encoder cache and calls injection callback
 
@@ -224,7 +226,7 @@ def _recv_encoder_cache(
         pass
 
     def add_encoder_cache(self, request_id: str, input_id: int,
-                          encoder_cache: NDArray[np.float32], mm_hash: str):
+                          encoder_cache: torch.Tensor, mm_hash: str):
         """Add an encoder cache to the EC connector.
 
         This method adds the encoder cache to the self.encoder_cache dictionary
@@ -360,7 +362,7 @@ def schedule_send_encoder_cache_metadata(self, request_id: str,
 
     def schedule_send_encoder_cache(
         self, request_id: str, input_id: int,
-        encoder_cache: NDArray[np.float32], mm_hash: str
+        encoder_cache: torch.Tensor, mm_hash: str
     ) -> None:
         """Schedule encoder cache sending
 
@@ -377,7 +379,7 @@ def schedule_send_encoder_cache(
 
     def _finish_wrapper(
         self, callback: Callable, request_id: str, input_id: int,
-        encoder_cache: NDArray[np.float32], mm_hash: str
+        encoder_cache: torch.Tensor, mm_hash: str
     ):
 
         callback(request_id, input_id, encoder_cache, mm_hash)