init

vllm_vacc/vllm/engine/multiprocessing/engine.py

@@ -0,0 +1,103 @@
+from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR,
+                                         RPCError,
+                                         RPCProcessRequest,
+                                         RPCAbortRequest)
+from vllm.config import VllmConfig
+import signal
+from vllm.logger import init_logger
+from vllm.transformers_utils.config import (
+    maybe_register_config_serialize_by_value)
+from vllm.usage.usage_lib import UsageContext
+
+logger = init_logger(__name__)
+
+
+class MQLLMEngine:
+    
+    def _handle_process_request(self, request: RPCProcessRequest):
+        """Handle RPCProcessRequest by adding it to the LLMEngine."""
+        request_id = request.request_id
+
+        if self._errored_with is not None:
+            rpc_err = RPCError(request_id=request_id,
+                               is_engine_errored=True,
+                               exception=ENGINE_DEAD_ERROR(self._errored_with))
+            self._send_outputs(rpc_err)
+
+        try:
+            self.engine.add_request(
+                request_id=request_id,
+                prompt=request.prompt,
+                params=request.params,
+                lora_request=request.lora_request,
+                trace_headers=request.trace_headers,
+                prompt_adapter_request=request.prompt_adapter_request,
+                priority=request.priority)
+
+            if self.log_requests:
+                from vllm.engine.multiprocessing.engine import logger
+                
+                if request.prompt.get('prompt_token_ids') is not None:
+                    # logger.info("Added request: %s, %s, prompt length: %s", request.request_id, request.prompt['prompt_token_ids'], len(request.prompt['prompt_token_ids']))
+                    logger.info("Added request: %s, prompt length: %s", request.request_id, len(request.prompt['prompt_token_ids']))
+                else:
+                    logger.info("Added request %s.", request.request_id)
+
+        except Exception as e:
+            # We do not set self._errored = True here, since the error
+            # is due to an issue adding this request to the engine,
+            # rather than an issue with the engine itself.
+            is_errored = self._errored_with is not None
+            rpc_err = RPCError(request_id=request_id,
+                               is_engine_errored=is_errored,
+                               exception=e)
+            self._send_outputs(rpc_err)
+
+            # Remove request from the engine.
+            self.engine.abort_request(request_id)
+            
+    def _handle_abort_request(self, request: RPCAbortRequest):
+        self.engine.abort_request(request.request_id)
+        if self.log_requests:
+            from vllm.engine.multiprocessing.engine import logger
+            import vllm_vacc.vllm.model_executor.models.vars as global_vars 
+            logger.info("Aborted request: %s, prompt length: %s", request.request_id, global_vars.DO_SEQ_LENS)
+
+def run_mp_engine(vllm_config: VllmConfig, usage_context: UsageContext,
+                  ipc_path: str, disable_log_stats: bool,
+                  disable_log_requests: bool, engine_alive):
+    
+    #patch to prevent num_speculative_tokens > 1
+    speculative_mode = hasattr(vllm_config, 'speculative_config')
+    if speculative_mode and \
+        hasattr(vllm_config.speculative_config, 'num_speculative_tokens') and \
+            vllm_config.speculative_config.num_speculative_tokens != 1:
+        raise ValueError(f'run_mp_engine: only support num_speculative_tokens == 1, but get {vllm_config.speculative_config.num_speculative_tokens}')
+    
+    default_model_infos = "default"
+    if speculative_mode:
+        if hasattr(vllm_config.speculative_config, 'method'):
+            default_model_infos = vllm_config.speculative_config.method
+            
+    from vllm_vacc.vllm.config_manager import vllm_vacc_config_manager
+    vllm_vacc_config_manager().update_model_infos(default_model_infos)
+
+    try:
+        # Ensure we can serialize transformer config before spawning
+        maybe_register_config_serialize_by_value()
+        from vllm.engine.multiprocessing.engine import MQLLMEngine,signal_handler
+        engine = MQLLMEngine.from_vllm_config(
+            vllm_config=vllm_config,
+            usage_context=usage_context,
+            disable_log_stats=disable_log_stats,
+            disable_log_requests=disable_log_requests,
+            ipc_path=ipc_path)
+
+        signal.signal(signal.SIGTERM, signal_handler)
+
+        engine.start()
+
+    except BaseException as e:
+        logger.exception(e)
+        engine_alive.value = False
+        raise e