huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 11:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

add clear cache when batch is finished

Browse Source
This commit is contained in:
OlivierDehaene 2023-07-19 01:12:28 +02:00
parent 0111869ad0
commit 8793ae5890
2 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
router/src/infer.rs
View File

@ -349,6 +349,7 @@ async fn batching_task(
} }
metrics::gauge!("tgi_batch_current_size", 0.0); metrics::gauge!("tgi_batch_current_size", 0.0);
metrics::gauge!("tgi_batch_current_max_tokens", 0.0); metrics::gauge!("tgi_batch_current_max_tokens", 0.0);
let _ = client.clear_cache(None).await;
} }
} }
} }

2
server/text_generation_server/cache.py
View File

@ -29,6 +29,8 @@ class Cache:
keys = list(self.cache.keys()) keys = list(self.cache.keys())
for k in keys: for k in keys:
self.delete(k) self.delete(k)
if torch.cuda.is_available():
torch.cuda.empty_cache()
def __len__(self): def __len__(self):
return len(self.cache.keys()) return len(self.cache.keys())