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

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

use super::*;

#[derive(Debug, Clone, Dissolve)]
pub struct LocalBlockDataFactory<S: Storage> {
    layout: Arc<dyn BlockLayout<StorageType = S>>,
    block_set_idx: usize,
    worker_id: WorkerID,
    offload_filter: Option<Arc<dyn OffloadFilter>>,
}

impl<S: Storage> LocalBlockDataFactory<S> {
    pub fn new(
        layout: Arc<dyn BlockLayout<StorageType = S>>,
        block_set_idx: usize,
        worker_id: WorkerID,
        offload_filter: Option<Arc<dyn OffloadFilter>>,
    ) -> Self {
        Self {
            layout,
            block_set_idx,
            worker_id,
            offload_filter,
        }
    }
}

impl<S: Storage> BlockFactory<S, locality::Local> for LocalBlockDataFactory<S> {
    fn create_block_data(&self, block_idx: BlockId) -> BlockResult<BlockData<S>> {
        if block_idx >= self.layout.num_blocks() {
            return Err(BlockError::InvalidBlockID(block_idx));
        }

        let data = BlockData::new(
            self.layout.clone(),
            block_idx,
            self.block_set_idx,
            self.worker_id,
        );
        Ok(data)
    }

    fn num_blocks(&self) -> usize {
        self.layout.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> IntoBlocks<S, locality::Local> for LocalBlockDataFactory<S> {}