block_manager.rs 13.4 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
4
// SPDX-License-Identifier: Apache-2.0

use super::*;
5
use anyhow::Result;
Ryan Olson's avatar
Ryan Olson committed
6
7
8
use dynamo_llm::block_manager::block::{
    data::logical::distributed_leader_worker::DistributedLeaderWorkerResources, locality::Logical,
};
9
use dynamo_llm::block_manager::kv_consolidator::EventSource;
10
use dynamo_llm::block_manager::offload::filter::FrequencyFilter;
Ryan Olson's avatar
Ryan Olson committed
11
use dynamo_llm::block_manager::{BasicMetadata, BlockParallelismStrategy};
Richard Huo's avatar
Richard Huo committed
12
use dynamo_runtime::DistributedRuntime;
13
use dynamo_runtime::config::environment_names::kvbm as env_kvbm;
14
use pyo3::PyResult;
15
use std::time::Duration;
Ryan Olson's avatar
Ryan Olson committed
16
17
use tokio_util::sync::CancellationToken;

18
mod cache_stats;
Ryan Olson's avatar
Ryan Olson committed
19
20
mod controller;
mod distributed;
21

Ryan Olson's avatar
Ryan Olson committed
22
pub mod vllm;
23
24
25
26

/// Add bingings from this crate to the provided module
pub fn add_to_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<BlockManager>()?;
Ryan Olson's avatar
Ryan Olson committed
27
28
29
30
31
32
    m.add_class::<distributed::KvbmWorker>()?;
    m.add_class::<distributed::KvbmLeader>()?;
    m.add_class::<controller::BlockManagerClient>()?;
    m.add_class::<controller::BlockPoolStatus>()?;
    m.add_class::<controller::ResetBlocksResponse>()?;

33
34
35
    m.add_class::<distributed::PyNcclBootstrap>()?;
    m.add_class::<distributed::PyNcclCommRef>()?;

Ryan Olson's avatar
Ryan Olson committed
36
37
    vllm::add_to_module(m)?;

38
39
40
    Ok(())
}

Ryan Olson's avatar
Ryan Olson committed
41
42
43
44
45
46
47
48
49
50
51
52
type VllmBlockManager = dynamo_llm::block_manager::KvBlockManager<
    Logical<DistributedLeaderWorkerResources>,
    BasicMetadata,
>;

type VllmController = Arc<
    dynamo_llm::block_manager::controller::Controller<
        Logical<DistributedLeaderWorkerResources>,
        BasicMetadata,
    >,
>;

Richard Huo's avatar
Richard Huo committed
53
54
55
56
57
58
59
60
/// Creates a disk offload filter based on environment configuration.
/// Returns `Ok(None)` if the filter is disabled via `DYN_KVBM_DISABLE_DISK_OFFLOAD_FILTER`,
/// otherwise constructs a `FrequencyFilter` with standard parameters.
fn create_disk_offload_filter(
    cancel_token: &CancellationToken,
    runtime: &tokio::runtime::Handle,
) -> Result<Option<Arc<FrequencyFilter>>> {
    // Check if disk offload filter is disabled via environment variable
61
    let disable_filter = std::env::var(env_kvbm::DYN_KVBM_DISABLE_DISK_OFFLOAD_FILTER)
Richard Huo's avatar
Richard Huo committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
        .map(|v| v == "true" || v == "1")
        .unwrap_or(false);

    if disable_filter {
        return Ok(None);
    }

    // TODO: These values seem plausible for most use cases, but we need to figure out a better way to configure them.
    let frequency_filter = FrequencyFilter::new(
        2,
        Duration::from_secs(600),
        1_000_000,
        cancel_token.child_token(),
        runtime.clone(),
    )?;

    Ok(Some(Arc::new(frequency_filter)))
}

81
#[pyclass]
Ryan Olson's avatar
Ryan Olson committed
82
#[derive(Clone)]
83
pub struct BlockManager {
Ryan Olson's avatar
Ryan Olson committed
84
    inner: VllmBlockManager,
Richard Huo's avatar
Richard Huo committed
85
    _drt: Option<Arc<DistributedRuntime>>,
Ryan Olson's avatar
Ryan Olson committed
86
    _controller: Option<VllmController>,
87
88
}

Ryan Olson's avatar
Ryan Olson committed
89
// TODO: This is in desperate need of a massive refactor. We bind and instantiate this in Python, but we never actually use it.
90
#[pymethods]
Ryan Olson's avatar
Ryan Olson committed
91
#[allow(unused_variables)]
92
93
impl BlockManager {
    #[new]
Ryan Olson's avatar
Ryan Olson committed
94
    #[pyo3(signature = (worker_id, leader = None, page_size = 32, num_device_blocks = None, disable_device_pool = false))]
95
96
    fn new(
        worker_id: u64,
Ryan Olson's avatar
Ryan Olson committed
97
        leader: Option<distributed::KvbmLeader>,
98
        page_size: usize,
Ryan Olson's avatar
Ryan Olson committed
99
100
        num_device_blocks: Option<usize>,
        disable_device_pool: bool,
101
    ) -> PyResult<Self> {
Ryan Olson's avatar
Ryan Olson committed
102
        let cancel_token = CancellationToken::new();
103
104
105
        let mut config = dynamo_llm::block_manager::KvBlockManagerConfig::builder().runtime(
            dynamo_llm::block_manager::KvManagerRuntimeConfig::builder()
                .worker_id(worker_id)
Ryan Olson's avatar
Ryan Olson committed
106
                .cancellation_token(cancel_token.clone())
107
                .build()
108
                .map_err(to_pyerr)?,
109
        );
Ryan Olson's avatar
Ryan Olson committed
110
111
112
113

        let model_config = dynamo_llm::block_manager::KvManagerModelConfig::builder()
            .num_layers(1)
            .outer_dim(1)
114
            .page_size(page_size)
Ryan Olson's avatar
Ryan Olson committed
115
116
            .inner_dim(1);

117
        config = config.model(model_config.build().map_err(to_pyerr)?);
Ryan Olson's avatar
Ryan Olson committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

        let (leader, drt) = if let Some(leader) = leader {
            let (leader, rt) = leader.dissolve();

            if !disable_device_pool {
                config = config.device_layout(
                    dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                        .num_blocks(leader.num_device_blocks())
                        .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                        .build()
                        .map_err(to_pyerr)?,
                );
            }

            if leader.num_host_blocks() > 0 {
                tracing::info!("Using {} host blocks", leader.num_host_blocks());
134
                let mut host_layout_config =
Ryan Olson's avatar
Ryan Olson committed
135
136
                    dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                        .num_blocks(leader.num_host_blocks())
137
138
                        .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded));

Richard Huo's avatar
Richard Huo committed
139
140
141
                if leader.num_disk_blocks() > 0
                    && let Some(filter) =
                        create_disk_offload_filter(&cancel_token, &get_current_tokio_handle())
142
                            .map_err(to_pyerr)?
Richard Huo's avatar
Richard Huo committed
143
144
                {
                    host_layout_config = host_layout_config.offload_filter(Some(filter));
145
146
147
                }

                config = config.host_layout(host_layout_config.build().map_err(to_pyerr)?);
Ryan Olson's avatar
Ryan Olson committed
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
            }

            if leader.num_disk_blocks() > 0 {
                tracing::info!("Using {} disk blocks", leader.num_disk_blocks());
                config = config.disk_layout(
                    dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                        .num_blocks(leader.num_disk_blocks())
                        .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                        .build()
                        .map_err(to_pyerr)?,
                );
            }
            (Some(leader), rt)
        } else {
            tracing::info!("Leader not provided. Block transfer functionality will be disabled.");

            // let num_device_blocks = num_device_blocks
            //     .expect("num_device_blocks must be provided if leader is not provided");

            // config = config.device_layout(
            //     dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
            //         .num_blocks(num_device_blocks)
            //         .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
            //         .build()
            //         .map_err(to_pyerr)?,
            // );

            unimplemented!("Leader not provided");
            // (
            //     None,
            //     Arc::new(
            //         tokio::runtime::Builder::new_multi_thread()
            //             .enable_all()
            //             .build()
            //             .map_err(to_pyerr)?,
            //     ),
            // )
        };

Richard Huo's avatar
Richard Huo committed
187
        let rt = get_current_tokio_handle();
Ryan Olson's avatar
Ryan Olson committed
188

189
        let config = config.build().map_err(to_pyerr)?;
190
        Ok(BlockManager {
Ryan Olson's avatar
Ryan Olson committed
191
192
193
194
195
196
197
198
199
200
201
202
            inner: rt
                .block_on(async {
                    let resources =
                        DistributedLeaderWorkerResources::new(leader, cancel_token.child_token())?;

                    dynamo_llm::block_manager::KvBlockManager::<
                        Logical<DistributedLeaderWorkerResources>,
                        BasicMetadata,
                    >::new(config, resources)
                    .await
                })
                .map_err(to_pyerr)?,
Richard Huo's avatar
Richard Huo committed
203
            _drt: drt,
Ryan Olson's avatar
Ryan Olson committed
204
            _controller: None,
205
206
207
        })
    }

Ryan Olson's avatar
Ryan Olson committed
208
209
    fn block_size(&self) -> usize {
        self.inner.block_size()
210
211
    }

Ryan Olson's avatar
Ryan Olson committed
212
213
214
215
216
    fn init_controller(&mut self, component: Component) -> PyResult<()> {
        if self._controller.is_some() {
            tracing::warn!("Controller already initialized. Ignoring init_controller call.");
            return Ok(());
        }
217

Ryan Olson's avatar
Ryan Olson committed
218
        let block_manager = self.inner.clone();
Richard Huo's avatar
Richard Huo committed
219
        let controller = get_current_tokio_handle()
Ryan Olson's avatar
Ryan Olson committed
220
221
222
223
            .block_on(controller::Controller::new(
                block_manager,
                component.inner.clone(),
            ))
224
            .map_err(to_pyerr)?;
Ryan Olson's avatar
Ryan Olson committed
225
226
227

        self._controller = Some(Arc::new(controller));

228
        let instance_id = component.inner.drt().connection_id();
Ryan Olson's avatar
Ryan Olson committed
229
230
231
232
233
234
235
236
237

        tracing::info!(
            "Dynamo KVBM Controller: {}.{}:{}",
            component.inner.namespace().name(),
            component.inner.name(),
            instance_id
        );

        Ok(())
238
    }
Ryan Olson's avatar
Ryan Olson committed
239
}
240

Ryan Olson's avatar
Ryan Olson committed
241
242
243
244
impl BlockManager {
    #[inline(always)]
    pub fn get_block_manager(&self) -> &VllmBlockManager {
        &self.inner
245
    }
246
}
247
248
249
250
251
252
253

#[derive(Default)]
pub struct BlockManagerBuilder {
    worker_id: u64,
    leader: Option<distributed::KvbmLeader>,
    page_size: usize,
    disable_device_pool: bool,
254
    kvbm_metrics: Option<dynamo_llm::block_manager::metrics_kvbm::KvbmMetrics>,
255
    consolidator_config: Option<(String, Option<String>, EventSource)>, // (engine_endpoint, output_endpoint (optional), engine_source)
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
}

impl BlockManagerBuilder {
    pub fn new() -> Self {
        Self {
            page_size: 32, // default consistent with BlockManager::new
            ..Default::default()
        }
    }

    pub fn worker_id(mut self, id: u64) -> Self {
        self.worker_id = id;
        self
    }
    pub fn page_size(mut self, ps: usize) -> Self {
        self.page_size = ps;
        self
    }
    pub fn leader(mut self, l: distributed::KvbmLeader) -> Self {
        self.leader = Some(l);
        self
    }
    pub fn disable_device_pool(mut self, yes: bool) -> Self {
        self.disable_device_pool = yes;
        self
    }
Richard Huo's avatar
Richard Huo committed
282

283
284
285
286
287
288
289
    pub fn kvbm_metrics(
        mut self,
        metrics: dynamo_llm::block_manager::metrics_kvbm::KvbmMetrics,
    ) -> Self {
        self.kvbm_metrics = Some(metrics);
        self
    }
290

291
292
293
    pub fn consolidator_config(
        mut self,
        engine_endpoint: String,
294
        output_endpoint: Option<String>,
295
296
297
        engine_source: EventSource,
    ) -> Self {
        self.consolidator_config = Some((engine_endpoint, output_endpoint, engine_source));
298
299
300
        self
    }

301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
    /// Async build (call from an async context).
    pub async fn build(self) -> Result<BlockManager> {
        let worker_id = self.worker_id;
        let leader = self.leader.ok_or_else(|| {
            anyhow::anyhow!("leader is required (runtime is always taken from leader)")
        })?;

        // Get (inner leader handle, runtime) from the provided leader.
        let (leader_inner, drt) = leader.dissolve();

        let cancel_token = CancellationToken::new();

        // Runtime & model config
        let runtime_config = dynamo_llm::block_manager::KvManagerRuntimeConfig::builder()
            .worker_id(worker_id)
            .cancellation_token(cancel_token.clone())
            .build()?;

        let mut config =
            dynamo_llm::block_manager::KvBlockManagerConfig::builder().runtime(runtime_config);

        let model_config = dynamo_llm::block_manager::KvManagerModelConfig::builder()
            .num_layers(1)
            .outer_dim(1)
            .page_size(self.page_size)
            .inner_dim(1)
            .build()?;

        config = config.model(model_config);

        // Layouts derived from leader’s counts
        if !self.disable_device_pool {
            config = config.device_layout(
                dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                    .num_blocks(leader_inner.num_device_blocks())
                    .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                    .build()?,
            );
        }

        if leader_inner.num_host_blocks() > 0 {
342
            let mut host_layout_config =
343
344
                dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                    .num_blocks(leader_inner.num_host_blocks())
345
346
                    .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded));

Richard Huo's avatar
Richard Huo committed
347
348
349
350
351
            if leader_inner.num_disk_blocks() > 0
                && let Some(filter) =
                    create_disk_offload_filter(&cancel_token, &get_current_tokio_handle())?
            {
                host_layout_config = host_layout_config.offload_filter(Some(filter));
352
353
354
            }

            config = config.host_layout(host_layout_config.build()?);
355
356
357
358
359
360
361
362
363
364
365
        }

        if leader_inner.num_disk_blocks() > 0 {
            config = config.disk_layout(
                dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
                    .num_blocks(leader_inner.num_disk_blocks())
                    .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                    .build()?,
            );
        }

366
367
368
369
        let mut config_builder = config;
        if let Some(kvbm_metrics) = self.kvbm_metrics {
            config_builder = config_builder.kvbm_metrics(Some(kvbm_metrics));
        }
370

371
        if let Some((engine_ep, output_ep, engine_source)) = self.consolidator_config {
372
373
            config_builder =
                config_builder.consolidator_config(engine_ep, output_ep, engine_source);
374
375
        }

376
        let config = config_builder.build()?;
377
378
379
380
381
382
383
384
385
386
387
388

        let resources =
            DistributedLeaderWorkerResources::new(Some(leader_inner), cancel_token.child_token())?;

        let inner = dynamo_llm::block_manager::KvBlockManager::<
            Logical<DistributedLeaderWorkerResources>,
            BasicMetadata,
        >::new(config, resources)
        .await?;

        Ok(BlockManager {
            inner,
Richard Huo's avatar
Richard Huo committed
389
            _drt: drt,
390
391
392
393
            _controller: None,
        })
    }
}