state.rs 18.7 KB
Newer Older
Ryan Olson's avatar
Ryan Olson committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::*;

18
19
use super::offload::OffloadManager;
use super::{
20
    block::{Block, GlobalRegistry, ImmutableBlock},
21
    config::NixlOptions,
22
    events::{EventManager, NullEventManager},
23
    metrics::{BlockManagerMetrics, PoolMetrics},
24
};
Ryan Olson's avatar
Ryan Olson committed
25
use std::sync::Arc;
26
use tokio::runtime::Handle;
Ryan Olson's avatar
Ryan Olson committed
27
28
29
30
31
32

#[allow(dead_code)]
pub struct KvBlockManagerState<Metadata: BlockMetadata> {
    worker_id: WorkerID,
    cancellation_token: CancellationToken,

33
    nixl_agent: Arc<Option<NixlAgent>>,
Ryan Olson's avatar
Ryan Olson committed
34
35
    nixl_backends: HashMap<String, Arc<nixl_sys::Backend>>,

36
37
38
    disk_pool: Option<Arc<BlockPool<DiskStorage, Metadata>>>,
    host_pool: Option<Arc<BlockPool<PinnedStorage, Metadata>>>,
    device_pool: Option<Arc<BlockPool<DeviceStorage, Metadata>>>,
Ryan Olson's avatar
Ryan Olson committed
39
40
41

    local_block_set: NixlBlockSet,
    remote_block_sets: RwLock<HashMap<WorkerID, HashMap<usize, RemoteBlocks>>>,
42
43

    offload_manager: Arc<OffloadManager<Metadata>>,
Ryan Olson's avatar
Ryan Olson committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
}

impl<Metadata: BlockMetadata> KvBlockManagerState<Metadata> {
    pub fn new(config: KvBlockManagerConfig) -> Result<Arc<Self>> {
        config
            .runtime
            .validate()
            .context("Validating runtime config")?;

        config.model.validate().context("Validating model config")?;

        let worker_id = config.runtime.worker_id;
        let cancellation_token = config.runtime.cancellation_token;

        // Create a map of NIXL backends
        let mut nixl_backends: HashMap<String, Arc<nixl_sys::Backend>> = HashMap::new();

61
        let global_registry = GlobalRegistry::default();
62
63
64

        let metrics = BlockManagerMetrics::new(&config.runtime.metrics_registry)?;

65
66
67
68
        let event_manager = config
            .event_manager
            .clone()
            .unwrap_or_else(|| NullEventManager::new());
69

Ryan Olson's avatar
Ryan Olson committed
70
71
        // Create a NIXL agent if NIXL is enabled and instantiate requested backends
        // TODO: Build a map of NIXL backends to block pools/sets
72
        let nixl_agent = Arc::new(match config.runtime.nixl {
Ryan Olson's avatar
Ryan Olson committed
73
74
75
76
77
            NixlOptions::Enabled => {
                tracing::debug!("Creating NIXL agent");
                let agent = NixlAgent::new(&worker_id.to_string())?;

                tracing::debug!("Creating NIXL backends");
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

                if let Ok((_, ucx_params)) = agent.get_plugin_params("UCX") {
                    let backend = agent.create_backend("UCX", &ucx_params)?;
                    nixl_backends.insert("UCX".to_string(), Arc::new(backend));
                } else {
                    tracing::warn!("No UCX plugin found; will not create UCX backend");
                }

                if config.disk_layout.is_some() {
                    if let Ok((_, gds_params)) = agent.get_plugin_params("GDS") {
                        let backend = agent.create_backend("GDS", &gds_params)?;
                        nixl_backends.insert("GDS".to_string(), Arc::new(backend));
                    } else {
                        tracing::warn!("No GDS plugin found; will not create GDS backend");
                    }
                }
Ryan Olson's avatar
Ryan Olson committed
94
95
96
97
98

                Some(agent)
            }
            NixlOptions::EnabledWithAgent(agent) => Some(agent),
            NixlOptions::Disabled => None,
99
        });
Ryan Olson's avatar
Ryan Olson committed
100
101
102
103
104
105
106
107
108

        // Initialize model-specific layout config. The layout_builder is incomplete at this point.
        // We will clone this builder and apply the storage-specific configs to each clone in the
        // following steps.
        let model = &config.model;
        let mut layout_builder = LayoutConfig::builder();

        layout_builder
            .num_layers(model.num_layers)
109
            .outer_dim(model.outer_dim)
Ryan Olson's avatar
Ryan Olson committed
110
111
112
113
114
115
116
            .page_size(model.page_size)
            .inner_dim(model.inner_dim)
            .dtype(model.dtype);

        let mut next_block_set_idx = 0;
        let mut local_block_set = block::nixl::NixlBlockSet::new(worker_id);

117
118
119
120
121
122
123
124
        let async_rt_handle = match config.runtime.async_runtime {
            Some(rt) => rt.handle().clone(),
            None => match Handle::try_current() {
                Ok(handle) => handle,
                Err(e) => anyhow::bail!(e),
            },
        };

125
126
127
        let (disk_pool, disk_blocks) = if let Some(config) = config.disk_layout {
            if nixl_agent.is_none() {
                tracing::warn!("NIXL is disabled; will not allocate disk blocks.");
128
                (None, None)
129
130
131
132
133
134
135
136
137
138
139
            } else {
                next_block_set_idx += 1;
                tracing::debug!("Constructing disk pool.");
                let layout =
                    create_layout(layout_builder.clone(), config, nixl_agent.as_ref().as_ref())?;
                local_block_set.add_block_set(next_block_set_idx, layout.serialize()?);
                let (pool, blocks) = create_block_pool::<_, Metadata>(
                    layout,
                    next_block_set_idx,
                    cancellation_token.clone(),
                    worker_id,
140
141
                    global_registry.clone(),
                    async_rt_handle.clone(),
142
                    metrics.pool("disk"),
143
                    Some(event_manager.clone()),
144
                )?;
145
                (Some(Arc::new(pool)), Some(blocks))
146
147
148
            }
        } else {
            tracing::debug!("No disk layout provided; will not allocate disk blocks.");
149
            (None, None)
150
151
        };

Ryan Olson's avatar
Ryan Olson committed
152
153
154
155
        // Create the host block pool if a host layout is provided
        let (host_pool, host_blocks) = if let Some(config) = config.host_layout {
            next_block_set_idx += 1;
            tracing::debug!("Constructing host pool.");
156
157
            let layout =
                create_layout(layout_builder.clone(), config, nixl_agent.as_ref().as_ref())?;
Ryan Olson's avatar
Ryan Olson committed
158
159
160
161
162
163
            local_block_set.add_block_set(next_block_set_idx, layout.serialize()?);
            let (pool, blocks) = create_block_pool::<_, Metadata>(
                layout,
                next_block_set_idx,
                cancellation_token.clone(),
                worker_id,
164
165
                global_registry.clone(),
                async_rt_handle.clone(),
166
                metrics.pool("host"),
167
                Some(event_manager.clone()),
Ryan Olson's avatar
Ryan Olson committed
168
            )?;
169
            (Some(Arc::new(pool)), Some(blocks))
Ryan Olson's avatar
Ryan Olson committed
170
171
        } else {
            tracing::debug!("No host layout provided; will not allocate host blocks.");
172
            (None, None)
Ryan Olson's avatar
Ryan Olson committed
173
174
175
176
177
178
        };

        // Create the device block pool if a device layout is provided
        let (device_pool, device_blocks) = if let Some(config) = config.device_layout {
            next_block_set_idx += 1;
            tracing::debug!("Constructing device pool.");
179
180
            let layout =
                create_layout(layout_builder.clone(), config, nixl_agent.as_ref().as_ref())?;
Ryan Olson's avatar
Ryan Olson committed
181
182
183
184
185
186
            local_block_set.add_block_set(next_block_set_idx, layout.serialize()?);
            let (pool, blocks) = create_block_pool::<_, Metadata>(
                layout,
                next_block_set_idx,
                cancellation_token.clone(),
                worker_id,
187
188
                global_registry.clone(),
                async_rt_handle.clone(),
189
                metrics.pool("device"),
190
                Some(event_manager.clone()),
Ryan Olson's avatar
Ryan Olson committed
191
            )?;
192
            (Some(Arc::new(pool)), Some(blocks))
Ryan Olson's avatar
Ryan Olson committed
193
194
        } else {
            tracing::debug!("No device layout provided; will not allocate device blocks.");
195
            (None, None)
Ryan Olson's avatar
Ryan Olson committed
196
197
198
        };

        // Finalize the local block set by adding NIXL metadata
199
        if let Some(nixl_agent) = nixl_agent.as_ref() {
Ryan Olson's avatar
Ryan Olson committed
200
201
202
203
            tracing::debug!("Finalize NixlBlockSet: adding NIXL metadata.");
            local_block_set.set_nixl_metadata(nixl_agent.get_local_md()?);
        }

204
205
206
207
208
        let offload_manager = OffloadManager::new(
            disk_pool.clone(),
            host_pool.clone(),
            device_pool.clone(),
            nixl_agent.clone(),
209
            async_rt_handle,
210
            metrics.clone(),
211
            cancellation_token.clone(),
212
        )?;
213

Ryan Olson's avatar
Ryan Olson committed
214
215
216
217
218
        let state = Arc::new(Self {
            worker_id,
            cancellation_token,
            nixl_agent,
            nixl_backends,
219
            disk_pool,
Ryan Olson's avatar
Ryan Olson committed
220
221
222
223
            host_pool,
            device_pool,
            local_block_set,
            remote_block_sets: RwLock::new(HashMap::new()),
224
            offload_manager,
Ryan Olson's avatar
Ryan Olson committed
225
226
        });

227
228
229
230
231
232
233
234
235
236
237
238
239
        if let Some(mut blocks) = disk_blocks {
            blocks.iter_mut().for_each(|block| {
                block.set_manager(state.clone());
            });

            state
                .disk_pool
                .as_ref()
                .as_ref()
                .unwrap()
                .add_blocks_blocking(blocks)?;
        }

Ryan Olson's avatar
Ryan Olson committed
240
241
242
243
244
245
246
247
        if let Some(mut blocks) = host_blocks {
            blocks.iter_mut().for_each(|block| {
                block.set_manager(state.clone());
            });

            state
                .host_pool
                .as_ref()
248
                .as_ref()
Ryan Olson's avatar
Ryan Olson committed
249
250
251
252
253
254
255
256
257
258
259
260
                .unwrap()
                .add_blocks_blocking(blocks)?;
        }

        if let Some(mut blocks) = device_blocks {
            blocks.iter_mut().for_each(|block| {
                block.set_manager(state.clone());
            });

            state
                .device_pool
                .as_ref()
261
                .as_ref()
Ryan Olson's avatar
Ryan Olson committed
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
                .unwrap()
                .add_blocks_blocking(blocks)?;
        }

        Ok(state)
    }

    /// Exports the local blockset configuration as a serialized object.
    pub fn export_local_blockset(&self) -> Result<SerializedNixlBlockSet> {
        SerializedNixlBlockSet::try_from(&self.local_block_set)
            .context("Failed to serialize local blockset")
    }

    /// Imports a remote blockset configuration from a serialized object.
    // TODO: NIXL will validate the every descriptor list against the memory registration list for
    // a given agent; this is can be an expensive operation. To avoid this, NIXL offers the ability
    // to generate "partial pre-validated (PPV)" descriptor lists. However, to support per-block and per-layer
    // PPV lists we will need as many as `num_layers + 1` PPV lists per block:
    // - one for representing the entire block
    // - one for representing each layer individually
    //
    // A deeper dive into the performance impact of PPV lists is required to determine if this is
    // the best approach.
    //
    // If PPV are valuable, it might be beneficial to lazily instantiate PPV lists when they are
    // needed; alternatively, we could generate the entire PPV list for each block at import time.
    pub fn import_remote_blockset(
        &self,
        serialized_blockset: SerializedNixlBlockSet,
    ) -> Result<()> {
        let remote = NixlBlockSet::try_from(serialized_blockset)
            .context("Failed to deserialize remote blockset")?;

        let (block_sets, metadata, worker_id) = remote.dissolve();
        tracing::debug!("Importing remote blockset from worker {}", worker_id);

        assert_ne!(
            worker_id, self.worker_id,
            "Cannot import blockset from self"
        );

        let agent = self
            .nixl_agent
            .as_ref()
306
            .as_ref()
Ryan Olson's avatar
Ryan Olson committed
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
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
            .ok_or_else(|| anyhow::anyhow!("NIXL agent not initialized"))?;

        let mut remote_block_sets = self.remote_block_sets.write().unwrap();

        if remote_block_sets.contains_key(&worker_id) {
            anyhow::bail!(
                "Worker ID {} already exists; cannot update remote blockset",
                worker_id
            );
        }

        let mut inner_map = HashMap::new();

        for (block_set_idx, block_set_layout) in block_sets {
            // Deserialize the individual layout and create RemoteBlocks
            let remote_blocks =
                RemoteBlocks::from_serialized(block_set_layout.clone(), block_set_idx, worker_id)?;

            // check the storage type of the remote blocks
            let layout = remote_blocks.layout();
            let storage = layout.storage();

            let storage = storage
                .first()
                .ok_or_else(|| anyhow::anyhow!("No storage found in remote blockset"))?;

            match storage.mem_type() {
                MemType::Dram => {
                    tracing::trace!(block_set_idx, "Detected Host/DRAM remote descriptor");
                }
                MemType::Vram => {
                    tracing::trace!(block_set_idx, "Detected GPU/Device/VRAM remote descriptor");
                }
                _ => {
                    tracing::warn!(
                        block_set_idx,
                        "Detected unknown remote descriptor; skipping blockset..."
                    );
                    continue;
                }
            }

            inner_map.insert(block_set_idx, remote_blocks);
        }

        let agent_id = agent
            .load_remote_md(&metadata)
            .context("Loading remote metadata")?;

        // try to convert the agent_id (String) to a WorkerID (u64)
        let agent_id: WorkerID =
            agent_id // Assuming agent_id is String here
                .parse() // Parse the String into u64 (WorkerID)
                .context("Failed to parse agent ID string into WorkerID (u64)")?;

        assert_eq!(agent_id, worker_id, "Mismatch with remote worker ID");

        remote_block_sets.insert(worker_id, inner_map);

        Ok(())
    }

    /// Get a [`Vec<RemoteBlock<IsImmutable>>`] from a [`BlockDescriptorList`]
    pub fn get_remote_blocks_immutable(
        &self,
        bds: &BlockDescriptorList,
    ) -> Result<Vec<RemoteBlock<IsImmutable>>> {
        // no checks - we can always create an immutable remote block even if the bds is mutable
        self.get_remote_blocks::<IsImmutable>(bds)
    }

    /// Get a [`Vec<RemoteBlock<IsMutable>>`] from a [`BlockDescriptorList`]
    pub fn get_remote_blocks_mutable(
        &self,
        bds: &BlockDescriptorList,
    ) -> Result<Vec<RemoteBlock<IsMutable>>> {
        if bds.mutability() == BlockMutability::Mutable {
            self.get_remote_blocks::<IsMutable>(bds)
        } else {
            anyhow::bail!("Cannot get mutable remote blocks for immutable block descriptor set");
        }
    }

    /// Generate a [`Vec<RemoteBlock>`] from a [`BlockDescriptorList`]
    fn get_remote_blocks<M: MutabilityKind>(
        &self,
        bds: &BlockDescriptorList,
    ) -> Result<Vec<RemoteBlock<M>>> {
        // Get a read lock on the remote block sets
        let remote_block_sets = self.remote_block_sets.read().unwrap();

        // validate we have loaded a remote blockset for the worker and the specific block_set_idx
        let remote_blocks = remote_block_sets
            .get(&bds.worker_id())
            .and_then(|map| map.get(&bds.block_set_idx()))
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "No remote blockset found for worker {} and block_set_idx {}",
                    bds.worker_id(),
                    bds.block_set_idx()
                )
            })?;

        // Iterate through indices, call .block() for each, and collect results.
        // The collect::<Result<...>>() handles potential errors from .block()
        let blocks: Vec<block::nixl::RemoteBlock<M>> = bds
            .block_indices()
            .iter()
            .map(|block_idx| remote_blocks.block(*block_idx))
            .collect::<Result<Vec<_>, _>>()?;

        Ok(blocks)
    }

421
    pub fn disk(&self) -> Option<&BlockPool<DiskStorage, Metadata>> {
422
        self.disk_pool.as_ref().map(|pool| pool.as_ref())
423
424
    }

Ryan Olson's avatar
Ryan Olson committed
425
    pub fn host(&self) -> Option<&BlockPool<PinnedStorage, Metadata>> {
426
        self.host_pool.as_ref().map(|pool| pool.as_ref())
Ryan Olson's avatar
Ryan Olson committed
427
428
429
    }

    pub fn device(&self) -> Option<&BlockPool<DeviceStorage, Metadata>> {
430
        self.device_pool.as_ref().map(|pool| pool.as_ref())
Ryan Olson's avatar
Ryan Olson committed
431
432
433
434
435
    }

    pub fn worker_id(&self) -> WorkerID {
        self.worker_id
    }
436
437
438
439
440
441
442
443
444
445
446

    pub(crate) async fn enqueue_offload_block<S: Storage + 'static>(
        &self,
        block: &ImmutableBlock<S, Metadata>,
        priority: u64,
    ) -> Result<()> {
        self.offload_manager.offload(block, priority).await?;

        Ok(())
    }

447
    pub async fn onboard_blocks<S: Storage>(
448
        &self,
449
450
        blocks: Vec<ImmutableBlock<S, Metadata>>,
    ) -> BlockResult<DeviceStorage, Metadata> {
451
452
        self.offload_manager.onboard(blocks).await
    }
Ryan Olson's avatar
Ryan Olson committed
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
}

impl<Metadata: BlockMetadata> std::fmt::Debug for KvBlockManagerState<Metadata> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "KvBlockManagerState")
    }
}

fn create_layout<S: Storage + NixlRegisterableStorage>(
    mut builder: LayoutConfigBuilder,
    config: KvManagerLayoutConfig<S>,
    nixl_agent: Option<&NixlAgent>,
) -> Result<Arc<dyn NixlLayout<StorageType = S>>> {
    let layout = builder.num_blocks(config.num_blocks).build()?;
    if let Some(storage) = config.storage {
        let mut layout = layout.create_layout(config.layout_type, storage)?;
        if let Some(nixl_agent) = nixl_agent {
            layout.nixl_register(nixl_agent, None)?;
        }
        return Ok(Arc::new(layout));
    }

    if let Some(allocator) = config.allocator {
        let mut layout = layout.allocate_layout(config.layout_type, allocator)?;
        if let Some(nixl_agent) = nixl_agent {
            layout.nixl_register(nixl_agent, None)?;
        }
        return Ok(Arc::new(layout));
    }

    anyhow::bail!("failed to create layout");
}

486
#[expect(clippy::type_complexity, clippy::too_many_arguments)]
Ryan Olson's avatar
Ryan Olson committed
487
488
489
490
491
fn create_block_pool<S: Storage + NixlRegisterableStorage, M: BlockMetadata>(
    layout: Arc<dyn NixlLayout<StorageType = S>>,
    block_set_idx: usize,
    cancellation_token: CancellationToken,
    worker_id: WorkerID,
492
493
    global_registry: GlobalRegistry,
    async_runtime: Handle,
494
    pool_metrics: Arc<PoolMetrics>,
495
    event_manager: Option<Arc<dyn EventManager>>,
Ryan Olson's avatar
Ryan Olson committed
496
497
) -> Result<(BlockPool<S, M>, Vec<Block<S, M>>)> {
    let blocks = block::layout_to_blocks::<_, M>(layout, block_set_idx, worker_id)?;
498
    let event_manager = event_manager.unwrap_or_else(|| NullEventManager::new());
Ryan Olson's avatar
Ryan Olson committed
499
500
    let pool = BlockPool::<S, M>::builder()
        .cancel_token(cancellation_token)
501
502
        .global_registry(global_registry)
        .async_runtime(async_runtime)
503
        .pool_metrics(pool_metrics)
504
        .event_manager(event_manager)
Ryan Olson's avatar
Ryan Olson committed
505
506
507
        .build()?;
    Ok((pool, blocks))
}