MCPcopy Create free account
hub / github.com/Open-Bee/DataStudio / ShardedLMDBManager

Class ShardedLMDBManager

datastudio/utils/database.py:92–461  ·  view source on GitHub ↗

Sharded LMDB manager with singleton pattern per cache directory. Distributes images across LMDB shards using hash-based partitioning. Loads shards on-demand with LRU eviction to minimize memory pressure. Example: >>> manager = ShardedLMDBManager.get_instance('./cache', num_shar

Source from the content-addressed store, hash-verified

90
91
92class ShardedLMDBManager:
93 """Sharded LMDB manager with singleton pattern per cache directory.
94
95 Distributes images across LMDB shards using hash-based partitioning.
96 Loads shards on-demand with LRU eviction to minimize memory pressure.
97
98 Example:
99 >>> manager = ShardedLMDBManager.get_instance('./cache', num_shards=32)
100 >>> manager.put('path/to/image.jpg', image_bytes)
101 >>> exists = manager.exists('path/to/image.jpg')
102 >>> data = manager.get('path/to/image.jpg')
103 """
104
105 _instances: Dict[str, "ShardedLMDBManager"] = {}
106
107 def __init__(
108 self,
109 cache_dir: str,
110 num_shards: int = 32,
111 map_size_per_shard: int = 1024 * 1024 * 1024 * 1024, # 1TB per shard
112 readonly: bool = False,
113 lock: bool = True,
114 logger=None,
115 max_open_shards: int = 0, # 0 means no limit
116 ):
117 """Initialize the sharded LMDB manager.
118
119 Args:
120 cache_dir: Cache directory path.
121 num_shards: Number of shards to use.
122 map_size_per_shard: Maximum size per shard in bytes.
123 readonly: Whether to open in read-only mode.
124 lock: Whether to use file locking.
125 logger: Logger instance.
126 max_open_shards: Maximum open shards (0 = no limit).
127 """
128 self.cache_dir = os.path.abspath(os.path.expanduser(cache_dir))
129 self.num_shards = num_shards
130 self.map_size_per_shard = map_size_per_shard
131 self.readonly = readonly
132 self.lock = lock
133 self.logger = logger
134 # 0 means no limit, default allows all shards open
135 self.max_open_shards = max_open_shards if max_open_shards > 0 else num_shards
136
137 # Use OrderedDict for LRU-style shard management
138 self._envs: OrderedDict[int, lmdb.Environment] = OrderedDict()
139 self._lock = threading.Lock()
140 os.makedirs(self.cache_dir, exist_ok=True)
141
142 if self.logger:
143 self.logger.info(
144 f"ShardedLMDBManager initialized: cache_dir={self.cache_dir}, "
145 f"num_shards={num_shards}, map_size_per_shard={map_size_per_shard/1024**3:.1f}GB"
146 )
147
148 @classmethod
149 def get_instance(

Callers 3

test_initMethod · 0.90
test_existsMethod · 0.90
test_close_allMethod · 0.90

Calls

no outgoing calls

Tested by 3

test_initMethod · 0.72
test_existsMethod · 0.72
test_close_allMethod · 0.72