"deploy/operator/vscode:/vscode.git/clone" did not exist on "d8b7d394ac10ff272ba43f27c4dae3f99a4e5798"
Unverified Commit defe5de7 authored by Michael Feil's avatar Michael Feil Committed by GitHub
Browse files

feat: release gil for publisher (#3747)


Signed-off-by: default avatarmichaelfeil <me@michaelfeil.eu>
parent 7cfac6b8
...@@ -263,7 +263,7 @@ impl KvEventPublisher { ...@@ -263,7 +263,7 @@ impl KvEventPublisher {
#[pyo3(signature = (event_id, token_ids, num_block_tokens, block_hashes, lora_id, parent_hash=None))] #[pyo3(signature = (event_id, token_ids, num_block_tokens, block_hashes, lora_id, parent_hash=None))]
fn publish_stored( fn publish_stored(
&mut self, &mut self,
_py: Python, py: Python,
event_id: u64, event_id: u64,
token_ids: Vec<u32>, token_ids: Vec<u32>,
num_block_tokens: Vec<u64>, num_block_tokens: Vec<u64>,
...@@ -271,27 +271,38 @@ impl KvEventPublisher { ...@@ -271,27 +271,38 @@ impl KvEventPublisher {
lora_id: u64, lora_id: u64,
parent_hash: Option<i64>, parent_hash: Option<i64>,
) -> PyResult<()> { ) -> PyResult<()> {
let kv_block_size = self.kv_block_size as u32;
let dp_rank = self.dp_rank;
let warning_count = self.warning_count.clone();
let inner = self.inner.clone();
py.allow_threads(|| {
let block_hashes_u64: Vec<u64> = block_hashes.iter().map(|&h| h as u64).collect(); let block_hashes_u64: Vec<u64> = block_hashes.iter().map(|&h| h as u64).collect();
let event = KvCacheEvent { let event = KvCacheEvent {
event_id, event_id,
data: KvCacheEventData::Stored(KvCacheStoreData { data: KvCacheEventData::Stored(KvCacheStoreData {
parent_hash: parent_hash.map(ExternalSequenceBlockHash::from), parent_hash: parent_hash.map(ExternalSequenceBlockHash::from),
blocks: create_stored_blocks( blocks: create_stored_blocks(
self.kv_block_size as u32, kv_block_size,
&token_ids, &token_ids,
&num_block_tokens, &num_block_tokens,
&block_hashes_u64, &block_hashes_u64,
lora_id, lora_id,
&self.warning_count, &warning_count,
), ),
}), }),
dp_rank: self.dp_rank, dp_rank,
}; };
self.inner.publish(event).map_err(to_pyerr) inner.publish(event).map_err(to_pyerr)
})
} }
fn publish_removed(&self, _py: Python, event_id: u64, block_hashes: Vec<i64>) -> PyResult<()> { fn publish_removed(&self, py: Python, event_id: u64, block_hashes: Vec<i64>) -> PyResult<()> {
let dp_rank = self.dp_rank;
let inner = self.inner.clone();
py.allow_threads(|| {
let block_hashes: Vec<ExternalSequenceBlockHash> = block_hashes let block_hashes: Vec<ExternalSequenceBlockHash> = block_hashes
.into_iter() .into_iter()
.map(ExternalSequenceBlockHash::from) .map(ExternalSequenceBlockHash::from)
...@@ -299,10 +310,11 @@ impl KvEventPublisher { ...@@ -299,10 +310,11 @@ impl KvEventPublisher {
let event = KvCacheEvent { let event = KvCacheEvent {
event_id, event_id,
data: KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }), data: KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }),
dp_rank: self.dp_rank, dp_rank,
}; };
self.inner.publish(event).map_err(to_pyerr) inner.publish(event).map_err(to_pyerr)
})
} }
} }
......
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