Unverified Commit 889f8bb2 authored by wz1qqx's avatar wz1qqx Committed by GitHub
Browse files

[BugFix]Reclaim resources to prevent memory leaks when use LMCacheMPConnector (#30745)


Signed-off-by: default avatarwz1qqx <ziqi.wang@novita.ai>
Co-authored-by: default avatarwz1qqx <ziqi.wang@novita.ai>
parent 058926d4
......@@ -147,6 +147,14 @@ class LMCacheMPSchedulerAdapter:
"""
return self.blocks_in_chunk
def _cleanup_lookup_result(self, request_id: str) -> None:
"""
Clean up lookup future for a finished request to prevent memory leak.
Args:
request_id: The ID of the finished request.
"""
self.lookup_futures.pop(request_id, None)
# Helper functions
def _create_key(self, block_hash: bytes) -> IPCCacheEngineKey:
"""Convert a block hash to an IPC cache engine key"""
......
......@@ -701,6 +701,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
if condition
else LMCacheMPRequestState.READY
)
# Clean up lookup future in scheduler adapter
self.scheduler_adapter._cleanup_lookup_result(request.request_id)
def build_connector_meta(
self, scheduler_output: SchedulerOutput
......@@ -754,6 +756,8 @@ class LMCacheMPConnector(KVConnectorBase_V1):
Optional KVTransferParams to be included in the request outputs
returned by the engine.
"""
# Clean up request tracker to prevent memory leak
self._cleanup_request_tracker(request.request_id)
return True, None
def take_events(self) -> Iterable["KVCacheEvent"]:
......@@ -915,3 +919,15 @@ class LMCacheMPConnector(KVConnectorBase_V1):
new_tracker = LMCacheMPRequestTracker(request)
self.request_trackers[request_id] = new_tracker
return self.request_trackers[request_id]
def _cleanup_request_tracker(self, request_id: str) -> None:
"""
Clean up request tracker and associated lookup future for a request.
This should be called when a request is finished to prevent memory leak.
"""
# Clean up request tracker
if self.request_trackers.pop(request_id, None):
logger.debug(
"[KVConnector] Cleaned up request_tracker for request %s",
request_id,
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment