[HiCache] feat: add multi tenant with prefix tag (#9256)

This commit is contained in:
Teng Ma
2025-10-11 00:23:28 +08:00
committed by GitHub
parent 3b9d97f335
commit 9082a7d323

View File

@@ -151,6 +151,12 @@ class MooncakeStore(HiCacheStorage):
)
per_tp_local_buffer_size = self.config.local_buffer_size // tp_scale_factor
# Check if extra_backend_tag should be passed to MooncakeDistributedStore
self.extra_backend_tag = None
if extra_config and "extra_backend_tag" in extra_config:
self.extra_backend_tag = extra_config["extra_backend_tag"]
logger.info(f"Using extra_backend_tag: {self.extra_backend_tag}")
ret_code = self.store.setup(
self.config.local_hostname,
self.config.metadata_server,
@@ -257,6 +263,11 @@ class MooncakeStore(HiCacheStorage):
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
get_results = self._get_batch_zero_copy_impl(
key_strs, buffer_ptrs, buffer_sizes
@@ -269,6 +280,11 @@ class MooncakeStore(HiCacheStorage):
host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
# Apply extra_backend_tag prefix if available
if self.extra_backend_tag is not None:
prefix = self.extra_backend_tag
keys = [f"{prefix}_{key}" for key in keys]
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
exist_result = self._batch_exist(key_strs)