Unverified Commit 9082a7d3 authored by Teng Ma's avatar Teng Ma Committed by GitHub
Browse files

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

parent 3b9d97f3
......@@ -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)
......
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