burn_p2p_dataloader 0.21.8

Dataset planning and lease-assignment boundary crate for burn_p2p.
Documentation 
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

use burn_p2p_core::{ContentId, DatasetView, MicroShard, MicroShardId};
use serde::{Deserialize, Serialize};

use crate::DataloaderError;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
/// Represents a shard fetch entry.
pub struct ShardFetchEntry {
    /// The microshard ID.
    pub microshard_id: MicroShardId,
    /// The ordinal.
    pub ordinal: u32,
    /// The locator.
    pub locator: String,
    /// The content hash.
    pub content_hash: ContentId,
    /// The bytes len.
    pub bytes_len: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
/// Describes the shard fetch.
pub struct ShardFetchManifest {
    /// The dataset view ID.
    pub dataset_view_id: burn_p2p_core::DatasetViewId,
    /// The entries.
    pub entries: Vec<ShardFetchEntry>,
}

impl ShardFetchManifest {
    /// Performs the entry for microshard operation.
    pub fn entry_for_microshard(&self, microshard_id: &MicroShardId) -> Option<&ShardFetchEntry> {
        self.entries
            .iter()
            .find(|entry| &entry.microshard_id == microshard_id)
    }

    /// Creates a value from the microshards.
    pub fn from_microshards(
        dataset_view: &DatasetView,
        microshards: &[MicroShard],
        bytes_for_ordinal: impl Fn(u32) -> Vec<u8>,
    ) -> Self {
        Self {
            dataset_view_id: dataset_view.dataset_view_id.clone(),
            entries: microshards
                .iter()
                .map(|microshard| {
                    let bytes = bytes_for_ordinal(microshard.ordinal);
                    ShardFetchEntry {
                        microshard_id: microshard.microshard_id.clone(),
                        ordinal: microshard.ordinal,
                        locator: format!("{:05}.bin", microshard.ordinal),
                        content_hash: ContentId::from_multihash(
                            burn_p2p_core::codec::multihash_sha256(&bytes),
                        ),
                        bytes_len: bytes.len() as u64,
                    }
                })
                .collect(),
        }
    }
}

pub(crate) fn verify_bytes(
    microshard_id: &MicroShardId,
    entry: &ShardFetchEntry,
    bytes: &[u8],
) -> Result<(), DataloaderError> {
    if entry.bytes_len != bytes.len() as u64 {
        return Err(DataloaderError::HashMismatch {
            microshard_id: microshard_id.clone(),
        });
    }

    let content_hash = ContentId::from_multihash(burn_p2p_core::codec::multihash_sha256(bytes));
    if content_hash != entry.content_hash {
        return Err(DataloaderError::HashMismatch {
            microshard_id: microshard_id.clone(),
        });
    }

    Ok(())
}