Return routed experts when request canceled (#4197)

Author: RunningLeon

lmdeploy/cli/serve.py

@@ -237,6 +237,7 @@ def api_server(args):
                 dllm_denoising_steps=args.dllm_denoising_steps,
                 dllm_confidence_threshold=args.dllm_confidence_threshold,
                 enable_return_routed_experts=args.enable_return_routed_experts,
+                distributed_executor_backend=args.distributed_executor_backend,
             )
         else:
             from lmdeploy.messages import TurbomindEngineConfig

lmdeploy/pytorch/engine/engine.py

@@ -927,18 +927,14 @@ def _make_infer_outputs(
                 num_accepted_tokens = (batched_outputs.next_token_ids[idx] > -1).sum() - 1
                 spec_info = dict(num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens)
             req_metrics = RequestMetrics(new_token_timestamp, msg.engine_events, spec_info=spec_info)
-            routed_experts = msg.routed_experts if msg.return_routed_experts and finish else None
-            if routed_experts is not None and self.engine_config.enable_transfer_obj_ref:
-                # only serialize for api server
-                routed_experts = self.executor.serialize(routed_experts)
             out = InferOutput(session_id=session_id,
                               resp=msg.resp,
                               finish=finish,
                               token_ids=token_ids,
                               cache_block_ids=cache_block_ids,
                               req_metrics=req_metrics,
                               logprobs=cur_logprobs,
-                              routed_experts=routed_experts)
+                              routed_experts=msg.routed_experts)
             outputs[session_id] = out
 
             if msg.return_logits:

lmdeploy/pytorch/engine/engine_instance.py

@@ -84,11 +84,30 @@ def __init__(self, engine: Engine):
         self.req_sender = engine.req_manager.build_sender()
 
         self.max_input_len = self.engine.max_session_len
+        self._enable_transfer_obj_ref = engine.engine_config.enable_transfer_obj_ref and \
+            engine.engine_config.distributed_executor_backend == 'ray'
 
     def __del__(self):
         """Destructor."""
         self.engine.req_manager.senders.pop(self.req_sender.sender_id)
 
+    def _get_extra_outputs(self, resp: Response):
+        """Get extra outputs."""
+        outputs = dict(routed_experts=None)
+        routed_experts = resp.data.get('routed_experts', None) if resp.data else None
+        if routed_experts is not None and resp.type in [ResponseType.FINISH, ResponseType.CANCEL]:
+            if self._enable_transfer_obj_ref:
+                import base64
+
+                import ray
+
+                ref = ray.put(routed_experts)
+                data = ray.cloudpickle.dumps(ref)
+                outputs['routed_experts'] = base64.b64encode(data).decode('utf-8')
+            else:
+                outputs['routed_experts'] = routed_experts
+        return outputs
+
     async def _async_try_add_session(self, session_id: int):
         """Add new session.
 
@@ -152,27 +171,28 @@ async def async_stream_infer(self,
             cache_block_ids = resp.data.get('cache_block_ids', None) if resp.data else None
             req_metrics = resp.data.get('req_metrics', None) if resp.data else None
             logprobs = resp.data.pop('logprobs', None) if resp.data else None
-            routed_experts = resp.data.get('routed_experts', None) if resp.data else None
+            extra_outputs = self._get_extra_outputs(resp)
+            routed_experts = extra_outputs.get('routed_experts', None)
 
             if resp.type == ResponseType.SUCCESS:
-                token_ids = resp.data['token_ids'].tolist()
+                token_ids = resp.data['token_ids']
                 num_ids = len(token_ids) - output_offset
                 logger.debug(f'session[{session_id}] success: num_out_ids={num_ids}.')
                 yield EngineOutput(resp.type,
-                                   token_ids[output_offset:],
+                                   token_ids[output_offset:].tolist(),
                                    cache_block_ids=cache_block_ids,
                                    req_metrics=req_metrics,
                                    routed_experts=routed_experts,
                                    logprobs=logprobs)
                 output_offset = len(token_ids)
-            elif resp.type == ResponseType.FINISH:
+            elif resp.type in (ResponseType.FINISH, ResponseType.CANCEL):
                 resp_data = resp.data
-                token_ids = resp_data['token_ids'].tolist()
+                token_ids = resp_data['token_ids']
                 logits = resp_data['logits']
                 num_ids = len(token_ids) - output_offset
                 logger.debug(f'session[{session_id}] finish: num_out_ids={num_ids}.')
                 yield EngineOutput(resp.type,
-                                   token_ids[output_offset:],
+                                   token_ids[output_offset:].tolist(),
                                    logits=logits,
                                    cache_block_ids=cache_block_ids,
                                    req_metrics=req_metrics,

lmdeploy/pytorch/engine/executor/base.py

@@ -102,10 +102,6 @@ def release(self):
         """Release resources."""
         raise NotImplementedError('Not Implemented.')
 
-    def serialize(self, obj):
-        """Serialize obj."""
-        return obj
-
     async def forward_async(self, inputs):
         """Start forward."""
         raise NotImplementedError('Not Implemented')

lmdeploy/pytorch/engine/executor/ray_executor.py

@@ -1,6 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import asyncio
-import base64
 import contextlib
 import json
 import os
@@ -354,13 +353,6 @@ def wakeup(self, tags: Optional[List[str]] = None):
             self.update_configs()
         self.collective_rpc('wakeup', (tags, ))
 
-    def serialize(self, obj) -> str:
-        """Serialize obj."""
-        ref = ray.put(obj)
-        data = ray.cloudpickle.dumps(ref)
-        data = base64.b64encode(data).decode('utf-8')
-        return data
-
     def get_input_processor(self):
         """Build cache engine."""
         return ray.get(self.workers[0].get_input_processor.remote())