dynamo-llm 1.0.2

Dynamo LLM Library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

use crate::block_manager::{block::factory::logical::LogicalBlockFactory, storage::StorageType};

/// The local block factories for the block manager
///
/// This struct will construct the factories in a consistent order and can be
/// used as an intermediate step before creating the block pools.
///
/// This is useful for debugging and for testing.
#[derive(Dissolve)]
pub struct LogicalBlockFactories<R: LogicalResources> {
    disk_factory: Option<LogicalBlockFactory<DiskStorage, R>>,
    host_factory: Option<LogicalBlockFactory<PinnedStorage, R>>,
    device_factory: Option<LogicalBlockFactory<DeviceStorage, R>>,
}

impl<R: LogicalResources> LogicalBlockFactories<R> {
    /// Construct the local block factories
    pub fn new(resources: &mut Resources, logical_resources: R) -> Result<Self> {
        let mut next_block_set_idx = 0;
        let layout_builder = resources.layout_builder();

        let logical_resources = Arc::new(logical_resources);

        let device_factory = if let Some(config) = resources.config.device_layout.take() {
            next_block_set_idx += 1;

            let offload_filter = config.offload_filter.clone();

            let mut builder = layout_builder.clone();
            let config = Arc::new(builder.num_blocks(config.num_blocks).build()?);

            let factory = LogicalBlockFactory::new(
                config,
                next_block_set_idx,
                resources.worker_id,
                logical_resources.clone(),
                StorageType::Device(0),
                offload_filter,
            );

            Some(factory)
        } else {
            None
        };

        let host_factory = if let Some(config) = resources.config.host_layout.take() {
            next_block_set_idx += 1;

            let offload_filter = config.offload_filter.clone();

            let mut builder = layout_builder.clone();
            let config = Arc::new(builder.num_blocks(config.num_blocks).build()?);
            let factory = LogicalBlockFactory::new(
                config,
                next_block_set_idx,
                resources.worker_id,
                logical_resources.clone(),
                StorageType::Pinned,
                offload_filter,
            );

            Some(factory)
        } else {
            None
        };

        let disk_factory = if let Some(config) = resources.config.disk_layout.take() {
            next_block_set_idx += 1;

            let offload_filter = config.offload_filter.clone();

            let mut builder = layout_builder.clone();
            let config = Arc::new(builder.num_blocks(config.num_blocks).build()?);
            let factory = LogicalBlockFactory::new(
                config,
                next_block_set_idx,
                resources.worker_id,
                logical_resources.clone(),
                StorageType::Disk(0),
                offload_filter,
            );

            Some(factory)
        } else {
            None
        };

        Ok(Self {
            disk_factory,
            host_factory,
            device_factory,
        })
    }
}