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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

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

use super::*;
use crate::block_manager::{
    OffloadFilter,
    locality::{Logical, LogicalBlockData, LogicalResources},
};

#[derive(Debug)]
pub struct LogicalBlockFactory<S: Storage, R: LogicalResources> {
    layout_config: Arc<LayoutConfig>,
    block_set_idx: usize,
    worker_id: WorkerID,
    resources: Arc<R>,
    storage_type: StorageType,
    storage: std::marker::PhantomData<S>,
    offload_filter: Option<Arc<dyn OffloadFilter>>,
}

impl<S: Storage, R: LogicalResources> LogicalBlockFactory<S, R> {
    pub fn new(
        layout_config: Arc<LayoutConfig>,
        block_set_idx: usize,
        worker_id: WorkerID,
        resources: Arc<R>,
        storage_type: StorageType,
        offload_filter: Option<Arc<dyn OffloadFilter>>,
    ) -> Self {
        Self {
            layout_config,
            block_set_idx,
            worker_id,
            resources,
            storage_type,
            storage: std::marker::PhantomData,
            offload_filter,
        }
    }
}

impl<S: Storage, R: LogicalResources> BlockFactory<S, Logical<R>> for LogicalBlockFactory<S, R> {
    fn create_block_data(&self, block_idx: BlockId) -> BlockResult<LogicalBlockData<S, R>> {
        if block_idx >= self.num_blocks() {
            return Err(BlockError::InvalidBlockID(block_idx));
        }

        let data = LogicalBlockData::new(
            block_idx,
            self.block_set_idx,
            self.worker_id,
            self.resources.clone(),
            self.storage_type,
            self.layout_config.page_size,
        );
        Ok(data)
    }

    fn num_blocks(&self) -> usize {
        self.layout_config.num_blocks
    }

    fn layout_config(&self) -> &LayoutConfig {
        &self.layout_config
    }

    fn offload_filter(&self) -> Option<Arc<dyn OffloadFilter>> {
        self.offload_filter.clone()
    }
}

impl<S: Storage, R: LogicalResources> IntoBlocks<S, Logical<R>> for LogicalBlockFactory<S, R> {}

#[cfg(test)]
mod tests {
    use crate::block_manager::block::data::logical::null::NullResources;
    use crate::block_manager::{ManagedBlockPool, PinnedStorage};

    use super::*;

    const TEST_BLOCK_SET_ID: usize = 42;
    const TEST_WORKER_ID: WorkerID = 1337;

    #[tokio::test]
    async fn test_logical_block_factory() {
        let layout_config = LayoutConfig::builder()
            .num_blocks(10)
            .page_size(16)
            .num_layers(3)
            .outer_dim(2)
            .inner_dim(8192)
            .dtype_width_bytes(2)
            .build()
            .unwrap();

        let factory = LogicalBlockFactory::<PinnedStorage, NullResources>::new(
            Arc::new(layout_config),
            TEST_BLOCK_SET_ID,
            TEST_WORKER_ID,
            Arc::new(NullResources),
            StorageType::Pinned,
            None,
        );

        let block_data = factory.create_block_data(0).unwrap();
        assert_eq!(block_data.block_id(), 0);
        assert_eq!(block_data.block_set_id(), TEST_BLOCK_SET_ID);
        assert_eq!(block_data.worker_id(), TEST_WORKER_ID);
        assert_eq!(block_data.storage_type(), &StorageType::Pinned);

        let _resources = block_data.resources();

        let blocks = factory
            .into_blocks_with_metadata(BasicMetadata::default())
            .unwrap();

        ManagedBlockPool::builder().blocks(blocks).build().unwrap();
    }
}