burn_pack/
reader.rs

1use super::base::{
2    Error, FORMAT_VERSION, HEADER_SIZE, Header, MAX_CBOR_RECURSION_DEPTH, MAX_METADATA_SIZE,
3    MAX_TENSOR_COUNT, MAX_TENSOR_SIZE, Metadata, TensorDescriptor, aligned_data_section_start,
4    validate_tensor_byte_len,
5};
6use super::tensor::Tensor;
7use alloc::format;
8use alloc::string::{String, ToString};
9use alloc::vec::Vec;
10use burn_std::{Bytes, Shape};
11
12#[cfg(feature = "std")]
13use super::base::MAX_FILE_SIZE;
14#[cfg(feature = "std")]
15use alloc::vec;
16#[cfg(feature = "std")]
17use std::fs::File;
18#[cfg(feature = "std")]
19use std::io::Read;
20#[cfg(feature = "std")]
21use std::path::Path;
22
23/// Reader for loading burnpack containers.
24pub struct Reader {
25    metadata: Metadata,
26    source: Source,
27    /// Absolute byte offset where the (256-byte aligned) tensor data section starts.
28    data_offset: usize,
29}
30
31impl Reader {
32    /// Load a pack from an in-memory [`Bytes`] buffer.
33    ///
34    /// Loading is lazy: only the header and metadata are parsed here. The buffer is kept as-is
35    /// (no copy, no share) and is only turned into zero-copy [`Bytes::view`] windows when you
36    /// consume the reader with [`into_tensors`](Self::into_tensors). For large models, prefer
37    /// [`from_file`](Self::from_file), which keeps tensor data file-backed and lazy.
38    pub fn from_bytes(bytes: Bytes) -> Result<Self, Error> {
39        let header = read_header(&bytes)?;
40        let metadata_end = HEADER_SIZE
41            .checked_add(header.metadata_size as usize)
42            .ok_or(Error::InvalidHeader)?;
43        if bytes.len() < metadata_end {
44            return Err(Error::InvalidHeader);
45        }
46        let metadata = parse_metadata(&bytes[HEADER_SIZE..metadata_end])?;
47
48        let available = bytes.len();
49        Self::assemble(&header, metadata, Source::Memory(bytes), available)
50    }
51
52    /// Load a pack from a file.
53    ///
54    /// Only the header and metadata are read up front; the whole file is wrapped in a single
55    /// lazy [`Bytes::from_file`] source, and each tensor's data is a [`Bytes::view`] window into
56    /// it, read from disk only when accessed. This integrates with the Burn ecosystem's file
57    /// allocation (pinned-memory staging) for fast file-to-GPU transfers.
58    ///
59    /// If `path` has no extension and does not exist as given, the canonical
60    /// [`crate::EXTENSION`] (`.bpk`) is appended.
61    #[cfg(feature = "std")]
62    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
63        let path = path.as_ref();
64        let path = if path.extension().is_none() && !path.exists() {
65            path.with_extension(crate::EXTENSION)
66        } else {
67            path.to_path_buf()
68        };
69
70        Self::from_resolved_file(&path)
71    }
72
73    /// Load a pack from a file using the exact path provided.
74    ///
75    /// Unlike [`from_file`](Self::from_file), this never appends [`crate::EXTENSION`] when the
76    /// requested path has no extension and does not probe for a fallback path. It is useful when
77    /// a higher-level caller has already applied its own path policy.
78    #[cfg(feature = "std")]
79    pub fn from_file_exact<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
80        Self::from_resolved_file(path.as_ref())
81    }
82
83    #[cfg(feature = "std")]
84    fn from_resolved_file(path: &Path) -> Result<Self, Error> {
85        let mut file = File::open(path)
86            .map_err(|e| Error::IoError(format!("cannot open '{}': {e}", path.display())))?;
87
88        let file_size = file.metadata().map_err(io_err)?.len();
89        if file_size > MAX_FILE_SIZE {
90            return Err(Error::ValidationError(format!(
91                "File size {file_size} bytes exceeds maximum allowed size of {MAX_FILE_SIZE} bytes"
92            )));
93        }
94
95        let mut header_bytes = [0u8; HEADER_SIZE];
96        file.read_exact(&mut header_bytes).map_err(io_err)?;
97        let header = read_header(&header_bytes)?;
98
99        let mut metadata_bytes = vec![0u8; header.metadata_size as usize];
100        file.read_exact(&mut metadata_bytes).map_err(io_err)?;
101        let metadata = parse_metadata(&metadata_bytes)?;
102
103        let source = Source::File(Bytes::from_file(path, file_size, 0));
104        Self::assemble(&header, metadata, source, file_size as usize)
105    }
106
107    /// Finish construction once the header, metadata, and data source are known.
108    ///
109    /// Centralizes the truncation check and the aligned data-section offset so both
110    /// [`from_bytes`](Self::from_bytes) and [`from_file`](Self::from_file) stay in sync.
111    /// `available` is the number of bytes the source can actually supply.
112    fn assemble(
113        header: &Header,
114        metadata: Metadata,
115        source: Source,
116        available: usize,
117    ) -> Result<Self, Error> {
118        let data_offset = aligned_data_section_start(header.metadata_size as usize);
119        validate_total_size(&metadata, data_offset, available)?;
120
121        Ok(Self {
122            metadata,
123            source,
124            data_offset,
125        })
126    }
127
128    /// Consume the reader, returning all tensors in sorted (alphabetical) name order.
129    ///
130    /// Each tensor's bytes are a zero-copy [`Bytes::view`] window into the source — no per-tensor
131    /// copy. For an in-memory source the buffer is [shared](Bytes::shared) once here (a cheap
132    /// `Arc` move, never a data copy) to make those views available; for a file source the windows
133    /// are file-backed and read lazily on access. Consuming `self` lets us hand the source's
134    /// ownership to the views directly, so loading never reads or copies tensor data eagerly.
135    pub fn into_tensors(self) -> Result<Vec<Tensor>, Error> {
136        let Reader {
137            metadata,
138            source,
139            data_offset,
140        } = self;
141
142        // Make the source view-capable: a plain in-memory buffer has no zero-copy window until
143        // it's shared behind an `Arc`, whereas a file-backed source already windows lazily (and
144        // must NOT be shared, or every view would materialize the whole file).
145        let source = match source {
146            Source::Memory(bytes) => bytes.shared(),
147            #[cfg(feature = "std")]
148            Source::File(bytes) => bytes,
149        };
150
151        let mut tensors = Vec::with_capacity(metadata.tensors.len());
152        for (name, descriptor) in &metadata.tensors {
153            let (start, end) = tensor_range(data_offset, name, descriptor)?;
154            let bytes = source.view(start, end).map_err(|_| {
155                Error::ValidationError(format!(
156                    "Tensor '{name}' data range {start}..{end} could not be viewed (source is {} bytes)",
157                    source.len()
158                ))
159            })?;
160            tensors.push(make_tensor(name, descriptor, bytes)?);
161        }
162        Ok(tensors)
163    }
164
165    /// The user-supplied key/value metadata stored alongside the tensors.
166    ///
167    /// For per-tensor info (dtype/shape/param id), use [`into_tensors`](Self::into_tensors) — for a
168    /// file-backed reader that does not read any tensor data until a tensor's bytes are accessed.
169    pub fn metadata(&self) -> &alloc::collections::BTreeMap<String, String> {
170        &self.metadata.metadata
171    }
172
173    /// The typed scalars stored alongside the tensors.
174    ///
175    /// Empty for files written before scalar support.
176    pub fn scalars(&self) -> &alloc::collections::BTreeMap<String, crate::Scalar> {
177        &self.metadata.scalars
178    }
179
180    /// The names of all tensors in the pack, in sorted (alphabetical) order.
181    pub fn tensor_names(&self) -> Vec<&str> {
182        self.metadata.tensors.keys().map(|n| n.as_str()).collect()
183    }
184
185    /// Read a single tensor's raw little-endian bytes by name (always copies).
186    ///
187    /// Returns [`Error::TensorNotFound`] if no tensor with that name exists.
188    pub fn tensor_data(&self, name: &str) -> Result<Vec<u8>, Error> {
189        let descriptor = self
190            .metadata
191            .tensors
192            .get(name)
193            .ok_or_else(|| Error::TensorNotFound(name.to_string()))?;
194        let (start, end) = tensor_range(self.data_offset, name, descriptor)?;
195
196        match &self.source {
197            #[cfg(feature = "std")]
198            Source::File(bytes) => {
199                // A file-backed view reads just this tensor's range from disk.
200                let view = bytes.view(start, end).map_err(|_| {
201                    Error::ValidationError(format!(
202                        "Tensor '{name}' data range {start}..{end} could not be viewed"
203                    ))
204                })?;
205                let slice: &[u8] = &view;
206                Ok(slice.to_vec())
207            }
208            Source::Memory(bytes) => Ok(memory_chunk(bytes, start, end)?.to_vec()),
209        }
210    }
211}
212
213/// Compute and validate the absolute `[start, end)` byte range of a tensor.
214fn tensor_range(
215    data_offset: usize,
216    name: &str,
217    descriptor: &TensorDescriptor,
218) -> Result<(usize, usize), Error> {
219    let to_usize = |offset: u64| -> Result<usize, Error> {
220        offset.try_into().map_err(|_| {
221            Error::ValidationError(format!(
222                "Tensor '{name}' has corrupted offset data: offset {offset} exceeds platform maximum"
223            ))
224        })
225    };
226    let overflow = || {
227        Error::ValidationError(format!(
228            "Tensor '{name}' has corrupted offset data: overflow"
229        ))
230    };
231
232    let start = data_offset
233        .checked_add(to_usize(descriptor.data_offsets.0)?)
234        .ok_or_else(overflow)?;
235    let end = data_offset
236        .checked_add(to_usize(descriptor.data_offsets.1)?)
237        .ok_or_else(overflow)?;
238
239    if end < start {
240        return Err(Error::ValidationError(format!(
241            "Tensor '{name}' has corrupted offset data: end {end} < start {start}"
242        )));
243    }
244    if end - start > MAX_TENSOR_SIZE {
245        return Err(Error::ValidationError(format!(
246            "Tensor '{name}' size {} exceeds maximum allowed size of {MAX_TENSOR_SIZE} bytes (potential DoS attack)",
247            end - start
248        )));
249    }
250    Ok((start, end))
251}
252
253/// Parse and validate a header from a buffer that starts with it.
254fn read_header(buf: &[u8]) -> Result<Header, Error> {
255    if buf.len() < HEADER_SIZE {
256        return Err(Error::InvalidHeader);
257    }
258    let header = Header::from_bytes(&buf[..HEADER_SIZE])?;
259    if header.version > FORMAT_VERSION {
260        return Err(Error::InvalidVersion);
261    }
262    if header.metadata_size > MAX_METADATA_SIZE {
263        return Err(Error::ValidationError(format!(
264            "Metadata size {} exceeds maximum allowed size of {MAX_METADATA_SIZE} bytes (potential DoS attack)",
265            header.metadata_size
266        )));
267    }
268    Ok(header)
269}
270
271/// Deserialize the CBOR metadata and validate the tensor count.
272fn parse_metadata(bytes: &[u8]) -> Result<Metadata, Error> {
273    let metadata: Metadata =
274        ciborium::de::from_reader_with_recursion_limit(bytes, MAX_CBOR_RECURSION_DEPTH)
275            .map_err(|e| Error::MetadataDeserializationError(e.to_string()))?;
276    if metadata.tensors.len() > MAX_TENSOR_COUNT {
277        return Err(Error::ValidationError(format!(
278            "File contains {} tensors, exceeding maximum of {MAX_TENSOR_COUNT} (potential DoS attack)",
279            metadata.tensors.len()
280        )));
281    }
282    Ok(metadata)
283}
284
285/// Ensure the available bytes can hold every tensor the metadata claims.
286fn validate_total_size(
287    metadata: &Metadata,
288    data_offset: usize,
289    available: usize,
290) -> Result<(), Error> {
291    let mut min_size = None;
292    for (name, descriptor) in &metadata.tensors {
293        let (start, end) = tensor_range(data_offset, name, descriptor)?;
294        validate_tensor_byte_len(
295            name,
296            descriptor.dtype,
297            &descriptor.shape,
298            (end - start) as u64,
299        )?;
300        min_size = Some(min_size.map_or(end, |current: usize| current.max(end)));
301    }
302
303    let Some(min_size) = min_size else {
304        return Ok(());
305    };
306    if available < min_size {
307        return Err(Error::ValidationError(format!(
308            "File truncated: expected at least {min_size} bytes, got {available} bytes"
309        )));
310    }
311    Ok(())
312}
313
314/// Borrow a tensor's `[start, end)` range out of an in-memory buffer, bounds-checked.
315fn memory_chunk(source: &Bytes, start: usize, end: usize) -> Result<&[u8], Error> {
316    let data: &[u8] = source;
317    data.get(start..end).ok_or_else(|| {
318        Error::ValidationError(format!(
319            "Tensor data range {start}..{end} is out of bounds (buffer is {} bytes)",
320            data.len()
321        ))
322    })
323}
324
325/// Build a [`Tensor`] entry from a descriptor + its data bytes.
326fn make_tensor(name: &str, descriptor: &TensorDescriptor, bytes: Bytes) -> Result<Tensor, Error> {
327    let shape = descriptor
328        .shape
329        .iter()
330        .map(|&s| {
331            s.try_into().map_err(|_| {
332                Error::ValidationError(format!(
333                    "Tensor '{name}' has corrupted shape data: dimension {s} exceeds platform maximum"
334                ))
335            })
336        })
337        .collect::<Result<Vec<usize>, Error>>()?;
338
339    Ok(Tensor::new(
340        name.to_string(),
341        descriptor.dtype,
342        Shape::from(shape),
343        descriptor.param_id,
344        bytes,
345    ))
346}
347
348/// Where a [`Reader`] gets its tensor data from.
349///
350/// Both variants hold a single [`Bytes`] spanning the whole container, and tensors are carved
351/// out of it with zero-copy [`Bytes::view`] windows.
352enum Source {
353    /// The whole pack lives in memory.
354    Memory(Bytes),
355    /// The pack lives in a file; tensor data is read lazily via file-backed [`Bytes::view`].
356    #[cfg(feature = "std")]
357    File(Bytes),
358}
359
360#[cfg(feature = "std")]
361fn io_err(e: std::io::Error) -> Error {
362    Error::IoError(e.to_string())
363}
364
365// Verifies the on-disk layout invariant (256-byte tensor alignment), which needs access to the
366// internal `TensorDescriptor` offsets. Public round-trip/error tests live in `tests/`.
367#[cfg(all(test, feature = "std"))]
368mod tests {
369    use super::*;
370    use crate::{TENSOR_ALIGNMENT, Tensor, Writer};
371    use alloc::collections::BTreeMap;
372    use burn_std::DType;
373
374    fn tensor(name: &str, elems: usize) -> Tensor {
375        Tensor::new(
376            name.to_string(),
377            DType::F32,
378            alloc::vec![elems],
379            None,
380            Bytes::from_bytes_vec(alloc::vec![0u8; elems * 4]),
381        )
382    }
383
384    fn pack_with_descriptor(shape: Vec<u64>, data_len: u64) -> Bytes {
385        let mut tensors = BTreeMap::new();
386        tensors.insert(
387            "weight".to_string(),
388            TensorDescriptor {
389                dtype: DType::F32,
390                shape,
391                data_offsets: (0, data_len),
392                param_id: None,
393            },
394        );
395        let metadata = Metadata {
396            tensors,
397            metadata: BTreeMap::new(),
398            scalars: BTreeMap::new(),
399        };
400        let mut metadata_bytes = Vec::new();
401        ciborium::ser::into_writer(&metadata, &mut metadata_bytes).unwrap();
402        let header = Header::new(metadata_bytes.len() as u32);
403        let data_offset = aligned_data_section_start(metadata_bytes.len());
404
405        let mut bytes = header.into_bytes().to_vec();
406        bytes.extend(metadata_bytes);
407        bytes.resize(data_offset + data_len as usize, 0);
408        Bytes::from_bytes_vec(bytes)
409    }
410
411    fn assert_byte_length_error(bytes: Bytes, expected: &str) {
412        match Reader::from_bytes(bytes) {
413            Err(Error::ValidationError(message)) => assert!(
414                message.contains(expected),
415                "expected error containing {expected:?}, got {message:?}"
416            ),
417            _ => panic!("expected a byte-length validation error"),
418        }
419    }
420
421    #[test]
422    fn tensor_offsets_are_256_aligned() {
423        // Odd sizes force the writer to insert padding between tensors.
424        let packed = Writer::new(vec![tensor("a", 3), tensor("b", 1), tensor("c", 2)])
425            .into_bytes()
426            .unwrap();
427        let reader = Reader::from_bytes(packed).unwrap();
428
429        for (name, descriptor) in &reader.metadata.tensors {
430            assert_eq!(
431                descriptor.data_offsets.0 % TENSOR_ALIGNMENT,
432                0,
433                "tensor '{name}' start offset is not 256-aligned"
434            );
435        }
436    }
437
438    #[test]
439    fn rejects_tensor_data_shorter_than_declared_shape() {
440        assert_byte_length_error(pack_with_descriptor(vec![2, 3], 16), "need 24");
441    }
442
443    #[test]
444    fn rejects_tensor_data_longer_than_declared_shape() {
445        assert_byte_length_error(pack_with_descriptor(vec![2, 3], 32), "need 24");
446    }
447
448    #[test]
449    fn rejects_tensor_byte_length_overflow() {
450        assert_byte_length_error(
451            pack_with_descriptor(vec![u64::MAX, 2], 0),
452            "calculation overflows u64",
453        );
454    }
455}
burn_pack/reader.rs

burn_pack/
reader.rs
