106 lines
3.1 KiB
Rust
106 lines
3.1 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use super::*;
|
|
|
|
pub use dynamo_llm::block_manager::controller::client::ControlClient;
|
|
pub use dynamo_llm::block_manager::controller::{CacheLevel, Controller};
|
|
|
|
#[pyclass]
|
|
pub struct BlockManagerClient {
|
|
inner: ControlClient,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl BlockManagerClient {
|
|
#[new]
|
|
fn new(component: Component, instance_id: i64) -> PyResult<Self> {
|
|
let client = pyo3_async_runtimes::tokio::get_runtime()
|
|
.block_on(ControlClient::new(component.inner, instance_id))
|
|
.map_err(to_pyerr)?;
|
|
Ok(BlockManagerClient { inner: client })
|
|
}
|
|
|
|
fn reset_pool(&self, cache_level: String) -> PyResult<()> {
|
|
let cache_level = Self::cache_level_from_str(&cache_level).map_err(to_pyerr)?;
|
|
pyo3_async_runtimes::tokio::get_runtime()
|
|
.block_on(self.inner.reset_pool(cache_level))
|
|
.map_err(to_pyerr)
|
|
}
|
|
|
|
fn reset_blocks(&self, cache_level: String, blocks: Vec<u64>) -> PyResult<ResetBlocksResponse> {
|
|
let cache_level = Self::cache_level_from_str(&cache_level).map_err(to_pyerr)?;
|
|
let response = pyo3_async_runtimes::tokio::get_runtime()
|
|
.block_on(self.inner.reset_blocks(cache_level, blocks))
|
|
.map_err(to_pyerr)?;
|
|
Ok(ResetBlocksResponse { inner: response })
|
|
}
|
|
|
|
fn status(&self, cache_level: String) -> PyResult<BlockPoolStatus> {
|
|
let cache_level = Self::cache_level_from_str(&cache_level).map_err(to_pyerr)?;
|
|
let status = pyo3_async_runtimes::tokio::get_runtime()
|
|
.block_on(self.inner.status(cache_level))
|
|
.map_err(to_pyerr)?;
|
|
Ok(BlockPoolStatus { inner: status })
|
|
}
|
|
|
|
fn reset_all_pools(&self) -> PyResult<()> {
|
|
pyo3_async_runtimes::tokio::get_runtime()
|
|
.block_on(self.inner.reset_all_pools())
|
|
.map_err(to_pyerr)
|
|
}
|
|
}
|
|
|
|
impl BlockManagerClient {
|
|
// convert string to cache level
|
|
fn cache_level_from_str(cache_level: &str) -> anyhow::Result<CacheLevel> {
|
|
match cache_level.to_uppercase().as_str() {
|
|
"G1" => Ok(CacheLevel::G1),
|
|
"G2" => Ok(CacheLevel::G2),
|
|
"G3" => Ok(CacheLevel::G3),
|
|
_ => anyhow::bail!("Invalid cache level: allowed values are G1, G2, G3"),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[pyclass]
|
|
#[derive(Clone)]
|
|
pub struct BlockPoolStatus {
|
|
inner: dynamo_llm::block_manager::pool::BlockPoolStatus,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl BlockPoolStatus {
|
|
fn active_blocks(&self) -> usize {
|
|
self.inner.active_blocks
|
|
}
|
|
|
|
fn inactive_blocks(&self) -> usize {
|
|
self.inner.inactive_blocks
|
|
}
|
|
|
|
fn empty_blocks(&self) -> usize {
|
|
self.inner.empty_blocks
|
|
}
|
|
}
|
|
|
|
#[pyclass]
|
|
pub struct ResetBlocksResponse {
|
|
inner: dynamo_llm::block_manager::pool::ResetBlocksResponse,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl ResetBlocksResponse {
|
|
fn reset_blocks(&self) -> Vec<u64> {
|
|
self.inner.reset_blocks.clone()
|
|
}
|
|
|
|
fn not_found_blocks(&self) -> Vec<u64> {
|
|
self.inner.not_found.clone()
|
|
}
|
|
|
|
fn not_reset_blocks(&self) -> Vec<u64> {
|
|
self.inner.not_reset.clone()
|
|
}
|
|
}
|