feat(mem_cache): add client-side metadata cache for HiCacheFile storage (#29716)
This commit is contained in:
@@ -413,6 +413,10 @@ class Envs:
|
|||||||
SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)
|
SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)
|
||||||
SGLANG_HICACHE_FILE_BACKEND_EVICTION_RATIO = EnvFloat(0.9)
|
SGLANG_HICACHE_FILE_BACKEND_EVICTION_RATIO = EnvFloat(0.9)
|
||||||
SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE = EnvStr("0")
|
SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE = EnvStr("0")
|
||||||
|
# Enable client-side metadata caching to optimize filesystem checks (e.g. for Lustre/NFS/FUSE)
|
||||||
|
SGLANG_HICACHE_FILE_BACKEND_ENABLE_METADATA_CACHE = EnvBool(False)
|
||||||
|
# Positive cache TTL for filesystem metadata lookups (-1 disables positive expiration)
|
||||||
|
SGLANG_HICACHE_FILE_BACKEND_METADATA_TTL = EnvFloat(5.0)
|
||||||
SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
|
SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
|
||||||
# Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache).
|
# Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache).
|
||||||
# Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the
|
# Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import threading
|
import threading
|
||||||
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
@@ -316,6 +317,38 @@ class HiCacheStorage(ABC):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class MetadataCache:
|
||||||
|
def __init__(self, ttl_seconds: float):
|
||||||
|
self.ttl_seconds = ttl_seconds
|
||||||
|
# key -> monotonic timestamp
|
||||||
|
self.cache: dict[str, float] = {}
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
|
||||||
|
def add(self, key: str):
|
||||||
|
with self.lock:
|
||||||
|
if key not in self.cache:
|
||||||
|
self.cache[key] = time.monotonic()
|
||||||
|
|
||||||
|
def remove(self, key: str):
|
||||||
|
with self.lock:
|
||||||
|
self.cache.pop(key, None)
|
||||||
|
|
||||||
|
def contains(self, key: str) -> bool:
|
||||||
|
with self.lock:
|
||||||
|
if key not in self.cache:
|
||||||
|
return False
|
||||||
|
if self.ttl_seconds == -1.0:
|
||||||
|
return True
|
||||||
|
if time.monotonic() - self.cache[key] > self.ttl_seconds:
|
||||||
|
del self.cache[key]
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def clear(self):
|
||||||
|
with self.lock:
|
||||||
|
self.cache.clear()
|
||||||
|
|
||||||
|
|
||||||
class HiCacheFile(HiCacheStorage):
|
class HiCacheFile(HiCacheStorage):
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -349,6 +382,30 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
os.makedirs(self.file_path)
|
os.makedirs(self.file_path)
|
||||||
logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
|
logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
|
||||||
|
|
||||||
|
# Metadata cache positive lookup toggle & TTL
|
||||||
|
enable_cache_raw = None
|
||||||
|
if storage_config.extra_config:
|
||||||
|
enable_cache_raw = storage_config.extra_config.get("enable_metadata_cache")
|
||||||
|
if enable_cache_raw is None:
|
||||||
|
enable_cache_raw = (
|
||||||
|
envs.SGLANG_HICACHE_FILE_BACKEND_ENABLE_METADATA_CACHE.get()
|
||||||
|
)
|
||||||
|
|
||||||
|
self.enable_metadata_cache = bool(enable_cache_raw)
|
||||||
|
|
||||||
|
if self.enable_metadata_cache:
|
||||||
|
ttl_raw = None
|
||||||
|
if storage_config.extra_config:
|
||||||
|
ttl_raw = storage_config.extra_config.get("metadata_ttl")
|
||||||
|
if ttl_raw is None:
|
||||||
|
ttl_raw = envs.SGLANG_HICACHE_FILE_BACKEND_METADATA_TTL.get()
|
||||||
|
|
||||||
|
self.metadata_ttl = float(ttl_raw) if ttl_raw is not None else 5.0
|
||||||
|
self.metadata_cache = MetadataCache(self.metadata_ttl)
|
||||||
|
self._scan_existing_files_to_metadata_cache()
|
||||||
|
else:
|
||||||
|
self.metadata_cache = None
|
||||||
|
|
||||||
# All LRU / size accounting and disk eviction lives in the evictor so
|
# All LRU / size accounting and disk eviction lives in the evictor so
|
||||||
# this backend stays a thin raw-bytes store. Imported lazily: the storage
|
# this backend stays a thin raw-bytes store. Imported lazily: the storage
|
||||||
# package __init__ pulls in the backend factory, which imports this
|
# package __init__ pulls in the backend factory, which imports this
|
||||||
@@ -361,6 +418,9 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
tp_rank=tp_rank,
|
tp_rank=tp_rank,
|
||||||
is_mla_model=is_mla_model,
|
is_mla_model=is_mla_model,
|
||||||
extra_config=storage_config.extra_config,
|
extra_config=storage_config.extra_config,
|
||||||
|
on_evict=(
|
||||||
|
self.metadata_cache.remove if self.metadata_cache is not None else None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
def _get_suffixed_key(self, key: str) -> str:
|
def _get_suffixed_key(self, key: str) -> str:
|
||||||
@@ -378,6 +438,19 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
self.file_path, f"{self._get_component_key(key, component_name)}.bin"
|
self.file_path, f"{self._get_component_key(key, component_name)}.bin"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _scan_existing_files_to_metadata_cache(self) -> None:
|
||||||
|
try:
|
||||||
|
names = os.listdir(self.file_path)
|
||||||
|
except FileNotFoundError:
|
||||||
|
return
|
||||||
|
for fn in names:
|
||||||
|
if not fn.endswith(".bin"):
|
||||||
|
continue
|
||||||
|
stem = fn[:-4]
|
||||||
|
# Only files belonging to this rank/model.
|
||||||
|
if stem.endswith(self.config_suffix):
|
||||||
|
self.metadata_cache.add(stem)
|
||||||
|
|
||||||
def get(
|
def get(
|
||||||
self,
|
self,
|
||||||
key: str,
|
key: str,
|
||||||
@@ -393,8 +466,12 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
if f.readinto(buf) != expected:
|
if f.readinto(buf) != expected:
|
||||||
raise IOError(f"Short read for {suffixed}")
|
raise IOError(f"Short read for {suffixed}")
|
||||||
self._evictor.touch(suffixed, tensor_path)
|
self._evictor.touch(suffixed, tensor_path)
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.add(suffixed)
|
||||||
return target_location
|
return target_location
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.remove(suffixed)
|
||||||
logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
|
logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -422,7 +499,7 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
|
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
|
||||||
|
|
||||||
# Fast path: same key already on disk. Refresh recency and skip rewrite.
|
# Fast path: same key already on disk. Refresh recency and skip rewrite.
|
||||||
if os.path.exists(tensor_path):
|
if self.exists(key):
|
||||||
logger.debug(f"Key {key} already exists. Skipped.")
|
logger.debug(f"Key {key} already exists. Skipped.")
|
||||||
self._evictor.touch(suffixed, tensor_path)
|
self._evictor.touch(suffixed, tensor_path)
|
||||||
return True
|
return True
|
||||||
@@ -443,6 +520,8 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
value.contiguous().view(dtype=torch.uint8).numpy().tofile(tmp_path)
|
value.contiguous().view(dtype=torch.uint8).numpy().tofile(tmp_path)
|
||||||
os.replace(tmp_path, tensor_path)
|
os.replace(tmp_path, tensor_path)
|
||||||
self._evictor.commit(suffixed)
|
self._evictor.commit(suffixed)
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.add(suffixed)
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to save tensor {key}: {e}")
|
logger.error(f"Failed to save tensor {key}: {e}")
|
||||||
@@ -454,6 +533,8 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
os.remove(tmp_path)
|
os.remove(tmp_path)
|
||||||
except OSError:
|
except OSError:
|
||||||
pass
|
pass
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.remove(suffixed)
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def batch_set(
|
def batch_set(
|
||||||
@@ -470,8 +551,14 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
|
|
||||||
def exists(self, key: str) -> bool:
|
def exists(self, key: str) -> bool:
|
||||||
key = self._get_suffixed_key(key)
|
key = self._get_suffixed_key(key)
|
||||||
|
if self.metadata_cache is not None and self.metadata_cache.contains(key):
|
||||||
|
return True
|
||||||
tensor_path = os.path.join(self.file_path, f"{key}.bin")
|
tensor_path = os.path.join(self.file_path, f"{key}.bin")
|
||||||
return os.path.exists(tensor_path)
|
if os.path.exists(tensor_path):
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.add(key)
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
def _collect_existing_component_keys(
|
def _collect_existing_component_keys(
|
||||||
self,
|
self,
|
||||||
@@ -483,11 +570,24 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
for key in keys:
|
for key in keys:
|
||||||
target_files.add(f"{self._get_component_key(key, transfer.name)}.bin")
|
target_files.add(f"{self._get_component_key(key, transfer.name)}.bin")
|
||||||
|
|
||||||
|
if self.metadata_cache is None:
|
||||||
|
existing_files = set()
|
||||||
|
with os.scandir(self.file_path) as entries:
|
||||||
|
for entry in entries:
|
||||||
|
if entry.is_file() and entry.name in target_files:
|
||||||
|
existing_files.add(entry.name)
|
||||||
|
return existing_files
|
||||||
|
|
||||||
existing_files = set()
|
existing_files = set()
|
||||||
with os.scandir(self.file_path) as entries:
|
for filename in target_files:
|
||||||
for entry in entries:
|
stem = filename[:-4]
|
||||||
if entry.is_file() and entry.name in target_files:
|
if self.metadata_cache.contains(stem):
|
||||||
existing_files.add(entry.name)
|
existing_files.add(filename)
|
||||||
|
else:
|
||||||
|
path = os.path.join(self.file_path, filename)
|
||||||
|
if os.path.exists(path):
|
||||||
|
self.metadata_cache.add(stem)
|
||||||
|
existing_files.add(filename)
|
||||||
return existing_files
|
return existing_files
|
||||||
|
|
||||||
def batch_exists_v2(
|
def batch_exists_v2(
|
||||||
@@ -607,6 +707,8 @@ class HiCacheFile(HiCacheStorage):
|
|||||||
if os.path.isfile(file_path):
|
if os.path.isfile(file_path):
|
||||||
os.remove(file_path)
|
os.remove(file_path)
|
||||||
self._evictor.clear()
|
self._evictor.clear()
|
||||||
|
if self.metadata_cache is not None:
|
||||||
|
self.metadata_cache.clear()
|
||||||
logger.info("Cleared all entries in HiCacheFile storage.")
|
logger.info("Cleared all entries in HiCacheFile storage.")
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import threading
|
import threading
|
||||||
from collections import OrderedDict
|
from collections import OrderedDict
|
||||||
from typing import Any, Optional, Set, Tuple
|
from typing import Any, Callable, Optional, Set, Tuple
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.utils.common import human_readable_int
|
from sglang.srt.utils.common import human_readable_int
|
||||||
@@ -71,10 +71,12 @@ class LRUFileEvictor:
|
|||||||
tp_rank: int,
|
tp_rank: int,
|
||||||
is_mla_model: bool,
|
is_mla_model: bool,
|
||||||
extra_config: Optional[dict] = None,
|
extra_config: Optional[dict] = None,
|
||||||
|
on_evict: Optional[Callable[[str], None]] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.file_path = file_path
|
self.file_path = file_path
|
||||||
self.config_suffix = config_suffix
|
self.config_suffix = config_suffix
|
||||||
self._tp_rank = tp_rank
|
self._tp_rank = tp_rank
|
||||||
|
self._on_evict = on_evict
|
||||||
|
|
||||||
# MLA ranks share the same physical files, so centralize LRU bookkeeping
|
# MLA ranks share the same physical files, so centralize LRU bookkeeping
|
||||||
# on rank 0; non-MLA ranks each own their own files via the suffix.
|
# on rank 0; non-MLA ranks each own their own files via the suffix.
|
||||||
@@ -342,8 +344,12 @@ class LRUFileEvictor:
|
|||||||
try:
|
try:
|
||||||
os.remove(tensor_path)
|
os.remove(tensor_path)
|
||||||
freed = evict_size
|
freed = evict_size
|
||||||
|
if self._on_evict is not None:
|
||||||
|
self._on_evict(evict_stem)
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
freed = 0 # file already gone; still drop the stale index entry
|
freed = 0 # file already gone; still drop the stale index entry
|
||||||
|
if self._on_evict is not None:
|
||||||
|
self._on_evict(evict_stem)
|
||||||
except OSError as e:
|
except OSError as e:
|
||||||
logger.warning(f"HiCacheFile eviction failed for {evict_stem}: {e}")
|
logger.warning(f"HiCacheFile eviction failed for {evict_stem}: {e}")
|
||||||
self._lru[evict_stem] = evict_size
|
self._lru[evict_stem] = evict_size
|
||||||
|
|||||||
@@ -27,7 +27,11 @@ from unittest import mock
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.mem_cache.hicache_storage import HiCacheFile, HiCacheStorageConfig
|
from sglang.srt.mem_cache.hicache_storage import (
|
||||||
|
HiCacheFile,
|
||||||
|
HiCacheStorageConfig,
|
||||||
|
MetadataCache,
|
||||||
|
)
|
||||||
from sglang.srt.mem_cache.storage.file.lru_file_evictor import _parse_size_to_bytes
|
from sglang.srt.mem_cache.storage.file.lru_file_evictor import _parse_size_to_bytes
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
@@ -83,6 +87,8 @@ class _BackendBuilder:
|
|||||||
is_mla=False,
|
is_mla=False,
|
||||||
model="testmodel",
|
model="testmodel",
|
||||||
subdir=None,
|
subdir=None,
|
||||||
|
metadata_ttl=None,
|
||||||
|
enable_metadata_cache=None,
|
||||||
) -> HiCacheFile:
|
) -> HiCacheFile:
|
||||||
# Each backend gets its own subdir so MLA / non-MLA tests don't
|
# Each backend gets its own subdir so MLA / non-MLA tests don't
|
||||||
# contaminate each other's file_path.
|
# contaminate each other's file_path.
|
||||||
@@ -101,6 +107,8 @@ class _BackendBuilder:
|
|||||||
"max_size": max_size,
|
"max_size": max_size,
|
||||||
"eviction_ratio": eviction_ratio,
|
"eviction_ratio": eviction_ratio,
|
||||||
"min_free_space": min_free,
|
"min_free_space": min_free,
|
||||||
|
"metadata_ttl": metadata_ttl,
|
||||||
|
"enable_metadata_cache": enable_metadata_cache,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
return HiCacheFile(cfg, file_path=d)
|
return HiCacheFile(cfg, file_path=d)
|
||||||
@@ -435,5 +443,125 @@ class TestPreReservationConcurrency(HiCacheFileLRUTestBase):
|
|||||||
self.assertLessEqual(b._evictor._total_bytes, 100)
|
self.assertLessEqual(b._evictor._total_bytes, 100)
|
||||||
|
|
||||||
|
|
||||||
|
class TestMetadataCache(CustomTestCase):
|
||||||
|
def test_metadata_cache_basic(self):
|
||||||
|
cache = MetadataCache(ttl_seconds=1.0)
|
||||||
|
cache.add("k1")
|
||||||
|
self.assertTrue(cache.contains("k1"))
|
||||||
|
self.assertFalse(cache.contains("k2"))
|
||||||
|
cache.remove("k1")
|
||||||
|
self.assertFalse(cache.contains("k1"))
|
||||||
|
|
||||||
|
def test_metadata_cache_ttl(self):
|
||||||
|
cache = MetadataCache(ttl_seconds=0.1)
|
||||||
|
cache.add("k1")
|
||||||
|
self.assertTrue(cache.contains("k1"))
|
||||||
|
time.sleep(0.2)
|
||||||
|
self.assertFalse(cache.contains("k1"))
|
||||||
|
|
||||||
|
def test_metadata_cache_hard_ttl(self):
|
||||||
|
cache = MetadataCache(ttl_seconds=0.3)
|
||||||
|
cache.add("k1")
|
||||||
|
time.sleep(0.15)
|
||||||
|
# Try updating k1
|
||||||
|
cache.add("k1")
|
||||||
|
# Expiry is still 0.3s from original timestamp, i.e. 0.15s from now.
|
||||||
|
time.sleep(0.2)
|
||||||
|
self.assertFalse(cache.contains("k1"))
|
||||||
|
|
||||||
|
def test_metadata_cache_infinite_ttl(self):
|
||||||
|
cache = MetadataCache(ttl_seconds=-1.0)
|
||||||
|
cache.add("k1")
|
||||||
|
time.sleep(0.3)
|
||||||
|
self.assertTrue(cache.contains("k1"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestHiCacheFileMetadataIntegration(HiCacheFileLRUTestBase):
|
||||||
|
def test_disabled_by_default(self):
|
||||||
|
b = self.make_backend()
|
||||||
|
self.assertIsNone(b.metadata_cache)
|
||||||
|
self.assertFalse(b.enable_metadata_cache)
|
||||||
|
|
||||||
|
def test_startup_scanning_populates_cache(self):
|
||||||
|
d = tempfile.mkdtemp(prefix="hicache_metadata_seed_", dir=self.tmpdir)
|
||||||
|
cfg = _make_config(
|
||||||
|
model="seedmodel",
|
||||||
|
extra_config={"metadata_ttl": 5.0, "enable_metadata_cache": True},
|
||||||
|
)
|
||||||
|
suffix = f"_seedmodel_0_1"
|
||||||
|
|
||||||
|
# Pre-create a suffixed bin file on disk
|
||||||
|
with open(os.path.join(d, f"k1{suffix}.bin"), "wb") as f:
|
||||||
|
f.write(b"data")
|
||||||
|
|
||||||
|
b = HiCacheFile(cfg, file_path=d)
|
||||||
|
# It should be found in metadata cache on startup
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
|
||||||
|
def test_write_and_read_populates_cache(self):
|
||||||
|
b = self.make_backend(metadata_ttl=5.0, enable_metadata_cache=True)
|
||||||
|
suffix = b.config_suffix
|
||||||
|
|
||||||
|
self.assertFalse(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
b.set("k1", _t(50))
|
||||||
|
# After set, it must be in the metadata cache
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
|
||||||
|
# Evict manually from metadata cache and call get
|
||||||
|
b.metadata_cache.clear()
|
||||||
|
self.assertFalse(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
b.get("k1", target_location=_t(50))
|
||||||
|
# Get should populate it back
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
|
||||||
|
def test_eviction_removes_from_metadata_cache(self):
|
||||||
|
# max_size=200, so setting three 100B tensors will evict the oldest
|
||||||
|
b = self.make_backend(
|
||||||
|
max_size="200",
|
||||||
|
eviction_ratio=1.0,
|
||||||
|
metadata_ttl=-1.0,
|
||||||
|
enable_metadata_cache=True,
|
||||||
|
)
|
||||||
|
suffix = b.config_suffix
|
||||||
|
|
||||||
|
b.set("k1", _t(100))
|
||||||
|
b.set("k2", _t(100))
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k2{suffix}"))
|
||||||
|
|
||||||
|
# Forces eviction of k1
|
||||||
|
b.set("k3", _t(100))
|
||||||
|
self.assertFalse(b.metadata_cache.contains(f"k1{suffix}"))
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k2{suffix}"))
|
||||||
|
self.assertTrue(b.metadata_cache.contains(f"k3{suffix}"))
|
||||||
|
|
||||||
|
def test_batch_exists_bypass_scandir(self):
|
||||||
|
b = self.make_backend(metadata_ttl=5.0, enable_metadata_cache=True)
|
||||||
|
suffix = b.config_suffix
|
||||||
|
|
||||||
|
b.set("k1", _t(50))
|
||||||
|
b.set("k2", _t(50))
|
||||||
|
|
||||||
|
# Now patch os.scandir and os.path.exists
|
||||||
|
with mock.patch("os.scandir") as mock_scandir, mock.patch(
|
||||||
|
"os.path.exists"
|
||||||
|
) as mock_exists:
|
||||||
|
mock_exists.return_value = True
|
||||||
|
|
||||||
|
# batch_exists_v2 for k1 and k2 should hit the metadata cache and NOT call os.scandir or os.path.exists
|
||||||
|
res = b.batch_exists_v2(["k1", "k2"])
|
||||||
|
self.assertEqual(res.kv_hit_pages, 2)
|
||||||
|
mock_scandir.assert_not_called()
|
||||||
|
mock_exists.assert_not_called()
|
||||||
|
|
||||||
|
# Querying "k3" (miss) should fall back to os.path.exists once but still NOT call os.scandir
|
||||||
|
res = b.batch_exists_v2(["k3"])
|
||||||
|
self.assertEqual(
|
||||||
|
res.kv_hit_pages, 1
|
||||||
|
) # since mock_exists returns True, k3 exists physically
|
||||||
|
mock_scandir.assert_not_called()
|
||||||
|
mock_exists.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main(verbosity=2)
|
unittest.main(verbosity=2)
|
||||||
|
|||||||
Reference in New Issue
Block a user