Skip to content

Commit 18175ef

Browse files
fix(sam3_image): prevent unwrap panic caused by LruCache eviction during mixed queries (#254)
* fix(sam3): prevent unwrap panic caused by LRU cache eviction * fix: use get instead of contains to promote LRU items and prevent eviction * fmt & clippy --------- Co-authored-by: jamjamjon <xxyydzml@outlook.com>
1 parent 9c37153 commit 18175ef

1 file changed

Lines changed: 20 additions & 2 deletions

File tree

src/models/vlm/sam3_image/impl.rs

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -348,12 +348,30 @@ impl Sam3Image {
348348
let mut uncached = Vec::new();
349349
let mut seen = HashSet::new();
350350
for p in prompts {
351-
if seen.insert(&p.text) && !self.text_cache.contains(&p.text) {
351+
if seen.insert(&p.text) && self.text_cache.get(&p.text).is_none() {
352352
uncached.push(p.text.as_str());
353353
}
354354
}
355+
356+
let unique_prompts_count = seen.len();
357+
if unique_prompts_count > self.text_cache.cap().get() {
358+
// If the number of unique prompts in the current batch exceeds the cache capacity, force an expansion.
359+
self.text_cache
360+
.resize(NonZeroUsize::new(unique_prompts_count).unwrap());
361+
}
362+
355363
if !uncached.is_empty() {
356-
for (text, feat) in uncached.iter().zip(self.encode_texts(engines, &uncached)?) {
364+
let encoded_feats = self.encode_texts(engines, &uncached)?;
365+
366+
if encoded_feats.len() != uncached.len() {
367+
bail!(
368+
"Encoded features length mismatch: expected {}, got {}",
369+
uncached.len(),
370+
encoded_feats.len()
371+
);
372+
}
373+
374+
for (text, feat) in uncached.iter().zip(encoded_feats) {
357375
self.text_cache.put(text.to_string(), feat);
358376
}
359377
}

0 commit comments

Comments
 (0)