[HiCache] feat: add multi tenant with prefix tag (#9256)
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user