add gc interface and bugfix

Author: JimyMa

lmdeploy/pytorch/disagg/conn/proxy_conn.py

@@ -74,7 +74,7 @@ def __init__(self):
         # put migrating session to `self.migration_session_shelf` for increasing fault tolerance
         # if a session is finished, then pop it from `self.migration_session_shelf`
         # if a decode instance is disconnected, then gc all blocks of these sessions in prefill instance.
-        self.migration_session_shelf: Dict[Tuple[str, str], Set[int]] = defaultdict(set)
+        self.migration_session_shelf: Dict[str, Set[int]] = defaultdict(set)
 
         # conn_perform handler queue
         self.waiting_conn: asyncio.Queue[Tuple[PDConnectionMessage, asyncio.Event]] = (asyncio.Queue())
@@ -93,17 +93,15 @@ def __init__(self):
 
     def reg_instance(self, role: EngineRole, endpoint: str):
         if role == EngineRole.Prefill:
-            logger.error('????????????????')
             self.prefill_endpoints.add(endpoint)
         elif role == EngineRole.Decode:
-            logger.error('????????????????')
             self.decode_endpoints.add(endpoint)
         else:
             raise ValueError(f'Unsupported role: {role}')
 
     def dereg_instance(self, endpoint: str):
         if endpoint in self.prefill_endpoints:
-            self.prefill_endpoints.pop(endpoint)
+            self.prefill_endpoints.remove(endpoint)
         elif endpoint in self.decode_endpoints:
             dropped_key = []
             for conn_key in self.pool.keys():
@@ -112,7 +110,13 @@ def dereg_instance(self, endpoint: str):
             for k in dropped_key:
                 self.drop(k)
             # TODO(JimyMa): handle side-effect by kvcache migration
-            self.decode_endpoints.pop(endpoint)
+            self.decode_endpoints.remove(endpoint)
+
+    def shelf_prefill_session(self, conn_key: Tuple[str, str], session_id: int):
+        self.migration_session_shelf[conn_key].add(session_id)
+
+    def unshelf_prefill_session(self, conn_key: Tuple[str, str], session_id: int):
+        self.migration_session_shelf[conn_key].remove(session_id)
 
     async def connect(self, conn_req: PDConnectionMessage):
 

lmdeploy/serve/proxy/proxy.py

@@ -507,6 +507,12 @@ async def connection_warmup():
     return JSONResponse({'SUCCESS': True})
 
 
+@app.post('/distserve/gc')
+async def cache_block_gc_to_be_migrated():
+    # TODO (JimyMa): add garbage collection of to be migrated request
+    raise NotImplementedError
+
+
 @app.post('/v1/chat/completions', dependencies=[Depends(check_api_key)])
 async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Request = None):
     """Completion API similar to OpenAI's API.
@@ -625,17 +631,21 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque
         ).model_dump(mode='json')
 
         start = node_manager.pre_call(d_url)
+        node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id'])
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, d_url, '/v1/chat/completions')
             background_task = node_manager.create_background_tasks(d_url, start)
+            node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info['id'])
             return StreamingResponse(response, background=background_task)
         else:
             try:
                 response = await node_manager.generate(request_dict, d_url, '/v1/chat/completions')
                 node_manager.post_call(d_url, start)
                 resp = JSONResponse(json.loads(response))
             finally:
+                node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info['id'])
                 return resp
+
     else:
         raise ValueError(f'No serving strategy named {node_manager.serving_strategy}')
 
@@ -737,15 +747,18 @@ async def completions_v1(request: CompletionRequest, raw_request: Request = None
             remote_block_ids=prefill_info['cache_block_ids'],
             remote_token_id=prefill_info['remote_token_ids'][-1],
         ).model_dump(mode='json')
+        node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id'])
 
         start = node_manager.pre_call(d_url)
         if request.stream is True:
             response = node_manager.stream_generate(request_dict, d_url, '/v1/completions')
             background_task = node_manager.create_background_tasks(d_url, start)
+            node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info['id'])
             return StreamingResponse(response, background=background_task)
         else:
             response = await node_manager.generate(request_dict, d_url, '/v1/completions')
             node_manager.post_call(d_url, start)
+            node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info['id'])
             return JSONResponse(json.loads(response))
     else:
         raise ValueError(f'No serving strategy named {node_manager.serving_strategy}')