pytorch_reader/
reader.rs

1//! Format detection and container loading behind [`PytorchReader`].
2//!
3//! The containers and their layouts are described in the crate docs.
4
5use crate::nested::{adapter::DefaultAdapter, data::NestedValue, de::Deserializer};
6use crate::{DType, Tensor};
7use serde::de::DeserializeOwned;
8use std::collections::HashMap;
9use std::fs::File;
10use std::io::{BufReader, Cursor, Read, Seek, SeekFrom};
11use std::path::Path;
12use std::sync::Arc;
13
14use crate::pickle_reader::{
15    Object, PersistentIds, PickleError, StorageRef, build_tensor, extract_tensors, key_string,
16    non_negative, read_pickle, storage_type_to_dtype,
17};
18use crate::storage::{LegacySource, StorageSource, TarSource, ZipSource};
19use byteorder::{LittleEndian, ReadBytesExt};
20
21/// Error type for PyTorch file operations
22#[derive(Debug)]
23#[non_exhaustive]
24pub enum PytorchError {
25    /// IO error
26    Io(std::io::Error),
27    /// Pickle parsing error
28    Pickle(PickleError),
29    /// Zip archive error
30    Zip(zip::result::ZipError),
31    /// TAR archive error
32    Tar(std::io::Error),
33    /// Invalid file format
34    InvalidFormat(String),
35    /// Key not found
36    KeyNotFound(String),
37    /// Serde deserialization error
38    Serde(crate::nested::error::Error),
39}
40
41impl From<std::io::Error> for PytorchError {
42    fn from(e: std::io::Error) -> Self {
43        PytorchError::Io(e)
44    }
45}
46
47impl From<PickleError> for PytorchError {
48    fn from(e: PickleError) -> Self {
49        PytorchError::Pickle(e)
50    }
51}
52
53impl From<zip::result::ZipError> for PytorchError {
54    fn from(e: zip::result::ZipError) -> Self {
55        PytorchError::Zip(e)
56    }
57}
58
59impl From<crate::nested::error::Error> for PytorchError {
60    fn from(e: crate::nested::error::Error) -> Self {
61        PytorchError::Serde(e)
62    }
63}
64
65impl std::fmt::Display for PytorchError {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        match self {
68            PytorchError::Io(e) => write!(f, "IO error: {}", e),
69            PytorchError::Pickle(e) => write!(
70                f,
71                "Pickle parsing error: {}. This may indicate an unsupported PyTorch file format or corrupted file.",
72                e
73            ),
74            PytorchError::Zip(e) => write!(f, "Zip archive error: {}", e),
75            PytorchError::Tar(e) => write!(f, "TAR archive error: {}", e),
76            PytorchError::InvalidFormat(msg) => write!(f, "Invalid PyTorch file format: {}", msg),
77            PytorchError::KeyNotFound(msg) => write!(f, "Key not found in PyTorch file: {}", msg),
78            PytorchError::Serde(e) => write!(f, "Serde deserialization error: {}", e),
79        }
80    }
81}
82
83impl std::error::Error for PytorchError {}
84
85type Result<T> = std::result::Result<T, PytorchError>;
86
87/// Metadata about a PyTorch file
88///
89/// Contains information about the file format, version, and other properties
90/// that can be useful for debugging or compatibility checking.
91#[derive(Debug, Clone)]
92pub struct PytorchMetadata {
93    /// Contents of the `.format_version` entry (e.g. `"1"`), if the archive has one
94    pub format_version: Option<String>,
95    /// File format type (ZIP, TAR, Legacy, or Pickle)
96    pub format_type: FileFormat,
97    /// Byte order. Always `LittleEndian`: big-endian files are rejected at load.
98    pub byte_order: ByteOrder,
99    /// Whether the file has storage alignment information
100    pub has_storage_alignment: bool,
101    /// Contents of the `version` entry (e.g. `"3"`): the serialized-file format version,
102    /// not the PyTorch release that wrote the file
103    pub pytorch_version: Option<String>,
104    /// Number of tensors in the file
105    pub tensor_count: usize,
106    /// Approximate size of the storage section in bytes (if available)
107    pub total_data_size: Option<usize>,
108}
109
110impl PytorchMetadata {
111    /// Check if this is a modern format file (ZIP-based, PyTorch 1.6+)
112    pub fn is_modern_format(&self) -> bool {
113        matches!(self.format_type, FileFormat::Zip)
114    }
115
116    /// Check if this is a legacy format file (PyTorch 0.1.10 - 1.5)
117    pub fn is_legacy_format(&self) -> bool {
118        matches!(self.format_type, FileFormat::Legacy)
119    }
120
121    fn for_format(format_type: FileFormat) -> Self {
122        Self {
123            format_version: None,
124            format_type,
125            byte_order: ByteOrder::LittleEndian,
126            has_storage_alignment: false,
127            pytorch_version: None,
128            tensor_count: 0,
129            total_data_size: None,
130        }
131    }
132}
133
134/// File format type
135#[derive(Debug, Clone, PartialEq)]
136pub enum FileFormat {
137    /// ZIP-based format (PyTorch 1.6+)
138    Zip,
139    /// TAR-based format (PyTorch before 0.1.10)
140    Tar,
141    /// Legacy format (PyTorch 0.1.10 - 1.5)
142    Legacy,
143    /// Simple pickle file
144    Pickle,
145}
146
147/// Byte order (endianness)
148#[derive(Debug, Clone, PartialEq)]
149pub enum ByteOrder {
150    LittleEndian,
151    BigEndian,
152}
153
154/// PyTorch checkpoint reader
155///
156/// This is the main interface for reading PyTorch checkpoint files (.pt/.pth).
157/// It supports multiple PyTorch formats including modern ZIP-based format (1.6+),
158/// legacy format (0.1.10-1.5), the early TAR format, and simple pickle files.
159///
160/// # Example
161/// ```rust,no_run
162/// # use pytorch_reader::PytorchReader;
163/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
164/// // Load a checkpoint file
165/// let reader = PytorchReader::new("model.pt")?;
166///
167/// // Get tensor names
168/// let keys = reader.keys();
169///
170/// // Access a specific tensor
171/// if let Some(tensor) = reader.get("conv1.weight") {
172///     let bytes = tensor.read()?; // Reads the tensor's bytes from the file
173/// }
174///
175/// // Check file metadata
176/// println!("Format: {:?}", reader.metadata().format_type);
177/// println!("Tensor count: {}", reader.metadata().tensor_count);
178/// # Ok(())
179/// # }
180/// ```
181#[derive(Debug)]
182pub struct PytorchReader {
183    tensors: HashMap<String, Tensor>,
184    metadata: PytorchMetadata,
185}
186
187impl PytorchReader {
188    /// Load a PyTorch checkpoint file
189    ///
190    /// # Arguments
191    /// * `path` - Path to the PyTorch file (.pt or .pth)
192    ///
193    /// # Returns
194    /// A `PytorchReader` with lazy-loaded tensors and metadata
195    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
196        Self::open(path.as_ref(), None)
197    }
198
199    /// Load a PyTorch checkpoint with a specific top-level key
200    ///
201    /// Many PyTorch checkpoints store the model weights under a specific key
202    /// like "state_dict", "model", or "model_state_dict".
203    ///
204    /// # Arguments
205    /// * `path` - Path to the PyTorch file
206    /// * `key` - Top-level key to extract (e.g., "state_dict")
207    ///
208    /// # Example
209    /// ```rust,no_run
210    /// # use pytorch_reader::PytorchReader;
211    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
212    /// let reader = PytorchReader::with_top_level_key("checkpoint.pt", "state_dict")?;
213    /// # Ok(())
214    /// # }
215    /// ```
216    pub fn with_top_level_key<P: AsRef<Path>>(path: P, key: &str) -> Result<Self> {
217        Self::open(path.as_ref(), Some(key))
218    }
219
220    fn open(path: &Path, top_level_key: Option<&str>) -> Result<Self> {
221        let Loaded { root, mut metadata } = load_file(path)?;
222        let tensors = extract_tensors_at(root, top_level_key)?;
223        metadata.tensor_count = tensors.len();
224        Ok(Self { tensors, metadata })
225    }
226
227    /// Load from a reader
228    ///
229    /// This method is useful when loading from non-file sources like memory buffers.
230    /// The reader must hold a plain pickle: tensor data lives outside the pickle in every
231    /// PyTorch container, so a checkpoint with tensors must be loaded from a file.
232    ///
233    /// # Arguments
234    /// * `reader` - Any type implementing `Read`
235    /// * `top_level_key` - Optional key to extract
236    pub fn from_reader<R: Read>(reader: R, top_level_key: Option<&str>) -> Result<Self> {
237        let root = read_pickle(&mut BufReader::new(reader), &PersistentIds::Unavailable)?;
238        let tensors = extract_tensors_at(root, top_level_key)?;
239        let mut metadata = PytorchMetadata::for_format(FileFormat::Pickle);
240        metadata.tensor_count = tensors.len();
241        Ok(Self { tensors, metadata })
242    }
243
244    /// Get all tensor names
245    pub fn keys(&self) -> Vec<String> {
246        self.tensors.keys().cloned().collect()
247    }
248
249    /// Get a tensor by name
250    pub fn get(&self, name: &str) -> Option<&Tensor> {
251        self.tensors.get(name)
252    }
253
254    /// Get all tensors
255    pub fn tensors(&self) -> &HashMap<String, Tensor> {
256        &self.tensors
257    }
258
259    /// Take ownership of all tensors
260    pub fn into_tensors(self) -> HashMap<String, Tensor> {
261        self.tensors
262    }
263
264    /// Get metadata about the loaded file
265    ///
266    /// Provides information about the file format, version, endianness, etc.
267    pub fn metadata(&self) -> &PytorchMetadata {
268        &self.metadata
269    }
270
271    /// Get the number of tensors in the file
272    pub fn len(&self) -> usize {
273        self.tensors.len()
274    }
275
276    /// Check if the file contains no tensors
277    pub fn is_empty(&self) -> bool {
278        self.tensors.is_empty()
279    }
280
281    /// Read raw pickle data from a PyTorch file
282    ///
283    /// This is useful for extracting configuration or metadata that isn't tensor data.
284    /// Returns a simplified JSON-like structure that can be easily converted to other formats.
285    /// Tensors and Python objects the reader does not interpret appear as
286    /// [`PickleValue::Unsupported`], naming their Python type.
287    ///
288    /// # Arguments
289    /// * `path` - Path to the PyTorch file
290    /// * `top_level_key` - Optional key to extract from the top-level dictionary
291    ///
292    /// # Returns
293    /// A `PickleValue` representing the pickle data structure
294    pub fn read_pickle_data<P: AsRef<Path>>(
295        path: P,
296        top_level_key: Option<&str>,
297    ) -> Result<PickleValue> {
298        let Loaded { root, .. } = load_file(path.as_ref())?;
299        let value = select_top_level(root, top_level_key)?;
300        Ok(to_pickle_value(value))
301    }
302
303    /// Load and deserialize configuration data from a PyTorch file
304    ///
305    /// This method reads configuration or metadata stored in PyTorch checkpoint files
306    /// and deserializes it into the specified type. It's particularly useful for
307    /// extracting model configurations that might be saved alongside model weights.
308    ///
309    /// # Arguments
310    /// * `path` - Path to the PyTorch file (.pt or .pth)
311    /// * `top_level_key` - Optional key to extract specific data within the pickle file.
312    ///   If `None`, the entire content is deserialized.
313    ///
314    /// # Type Parameters
315    /// * `D` - The target type to deserialize into. Must implement `DeserializeOwned`.
316    ///
317    /// # Returns
318    /// A `Result` containing the deserialized configuration data, or an `Error` if
319    /// reading or deserialization fails. A field whose value the reader cannot represent
320    /// (a numpy scalar, a `torch.dtype`, a tensor) is an error rather than a default.
321    /// Entries the target type does not name are skipped, so weights saved beside the
322    /// configuration do not get in the way.
323    ///
324    /// # Example
325    /// ```rust,no_run
326    /// # use pytorch_reader::PytorchReader;
327    /// # use serde::Deserialize;
328    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
329    /// #[derive(Debug, Deserialize)]
330    /// struct ModelConfig {
331    ///     hidden_size: usize,
332    ///     num_layers: usize,
333    /// }
334    ///
335    /// let config: ModelConfig = PytorchReader::load_config("model.pth", Some("config"))?;
336    /// # Ok(())
337    /// # }
338    /// ```
339    pub fn load_config<D, P>(path: P, top_level_key: Option<&str>) -> Result<D>
340    where
341        D: DeserializeOwned,
342        P: AsRef<Path>,
343    {
344        let pickle_value = Self::read_pickle_data(path, top_level_key)?;
345        let nested_value = to_nested_value(pickle_value);
346        let deserializer = Deserializer::<DefaultAdapter>::new(nested_value, false);
347        Ok(D::deserialize(deserializer)?)
348    }
349}
350
351/// Simplified representation of pickle data
352///
353/// This enum provides a JSON-like structure that's easier to work with
354/// than the internal pickle Object type.
355#[derive(Debug, Clone, PartialEq)]
356pub enum PickleValue {
357    /// None/null value
358    None,
359    /// Boolean value
360    Bool(bool),
361    /// Integer value
362    Int(i64),
363    /// Floating point value
364    Float(f64),
365    /// String value
366    String(String),
367    /// List/array of values
368    List(Vec<PickleValue>),
369    /// Dictionary/map of string keys to values
370    Dict(HashMap<String, PickleValue>),
371    /// Binary data
372    Bytes(Vec<u8>),
373    /// A value the reader cannot represent, holding the name of its Python type: a tensor
374    /// or storage, a class such as `torch.float32`, an object it does not interpret
375    /// (`numpy.core.multiarray.scalar`, `torch.device`), or an `int` too wide for an `i64`.
376    Unsupported(String),
377}
378
379// ---------------------------------------------------------------------------------------------
380// Format detection and loading
381// ---------------------------------------------------------------------------------------------
382
383/// A parsed file: the root pickle object plus what the container says about itself.
384struct Loaded {
385    root: Object,
386    /// `tensor_count` is filled in once tensors have been extracted.
387    metadata: PytorchMetadata,
388}
389
390/// The magic number 0x1950a86a20f9469cfc6c opening a legacy file, as protocols 2 and up
391/// write it: little-endian bytes inside a `LONG1`.
392const LEGACY_MAGIC_BYTES: [u8; 10] = [0x6c, 0xfc, 0x9c, 0x46, 0xf9, 0x20, 0x6a, 0xa8, 0x50, 0x19];
393
394/// The same value in decimal, which protocols 0 and 1 write as `LONG` text.
395const LEGACY_MAGIC_TEXT: &[u8] = b"119547037146038801333356";
396
397/// How far into the file the magic is looked for. `torch.save` picks the protocol, and
398/// each frames the first pickle differently: protocol 4 puts a `PROTO` and a 9-byte `FRAME`
399/// ahead of the value, protocol 0 none at all. Every encoding starts well inside this.
400const LEGACY_MAGIC_SEARCH: usize = 64;
401
402/// Whether a file opens with the legacy magic number in any protocol's encoding.
403fn starts_with_legacy_magic(header: &[u8]) -> bool {
404    let window = &header[..header.len().min(LEGACY_MAGIC_SEARCH)];
405    window
406        .windows(LEGACY_MAGIC_BYTES.len())
407        .any(|bytes| bytes == LEGACY_MAGIC_BYTES)
408        || window
409            .windows(LEGACY_MAGIC_TEXT.len())
410            .any(|bytes| bytes == LEGACY_MAGIC_TEXT)
411}
412
413/// Offset of the `ustar` magic within a TAR header.
414const TAR_MAGIC_OFFSET: usize = 257;
415
416fn detect_format(path: &Path) -> Result<FileFormat> {
417    let mut header = Vec::new();
418    File::open(path)?
419        .take((TAR_MAGIC_OFFSET + 5) as u64)
420        .read_to_end(&mut header)?;
421
422    if header.starts_with(b"PK\x03\x04") || header.starts_with(b"PK\x05\x06") {
423        Ok(FileFormat::Zip)
424    } else if header.get(TAR_MAGIC_OFFSET..TAR_MAGIC_OFFSET + 5) == Some(b"ustar") {
425        Ok(FileFormat::Tar)
426    } else if starts_with_legacy_magic(&header) {
427        Ok(FileFormat::Legacy)
428    } else {
429        Ok(FileFormat::Pickle)
430    }
431}
432
433fn load_file(path: &Path) -> Result<Loaded> {
434    match detect_format(path)? {
435        FileFormat::Zip => load_zip(path),
436        FileFormat::Tar => load_tar(path),
437        FileFormat::Legacy => load_legacy(path),
438        FileFormat::Pickle => load_plain_pickle(path),
439    }
440}
441
442fn load_zip(path: &Path) -> Result<Loaded> {
443    let source = ZipSource::open(path)?;
444
445    match source.read_text("byteorder")?.as_deref() {
446        None | Some("little") => {}
447        Some("big") => return Err(big_endian_error()),
448        Some(other) => {
449            return Err(PytorchError::InvalidFormat(format!(
450                "Unrecognized byteorder entry '{other}', expected 'little' or 'big'"
451            )));
452        }
453    }
454    let mut metadata = PytorchMetadata::for_format(FileFormat::Zip);
455    metadata.format_version = source.read_text(".format_version")?;
456    metadata.pytorch_version = source.read_text("version")?;
457    metadata.has_storage_alignment = source.has_entry(".storage_alignment");
458    metadata.total_data_size = Some(source.data_size()? as usize);
459
460    let pickle = source.pickle()?;
461    let ids = PersistentIds::Storages(Arc::new(StorageSource::Zip(source)));
462    let root = read_pickle(&mut Cursor::new(pickle), &ids)?;
463
464    Ok(Loaded { root, metadata })
465}
466
467fn load_legacy(path: &Path) -> Result<Loaded> {
468    let mut reader = BufReader::new(File::open(path)?);
469
470    let read_header = |reader: &mut BufReader<File>, what: &str| {
471        read_pickle(reader, &PersistentIds::Unavailable).map_err(|e| {
472            PytorchError::InvalidFormat(format!("Failed to read {what} from legacy format: {e}"))
473        })
474    };
475    // How long the magic pickle is depends on the protocol it was written at, so it is
476    // consumed as a pickle rather than skipped by a fixed offset. Its value is already
477    // known from the format detection.
478    read_header(&mut reader, "magic number")?;
479    // PyTorch refuses anything but 1001 here.
480    let protocol_version = read_header(&mut reader, "protocol version")?;
481    if !matches!(protocol_version, Object::Int(1001)) {
482        return Err(PytorchError::InvalidFormat(format!(
483            "Unsupported legacy protocol version {protocol_version:?}, expected 1001"
484        )));
485    }
486    let sys_info = read_header(&mut reader, "system info")?;
487    check_little_endian(&sys_info)?;
488
489    // The source reads storages at explicit offsets through a duplicate of the handle the
490    // headers stream through, so both come from the one open above and a replacement of
491    // `path` cannot slip in between two opens. No storage is read until `finish` below, by
492    // which point the stream is done with the cursor the duplicate shares.
493    let source = Arc::new(StorageSource::Legacy(LegacySource::new(
494        reader.get_ref().try_clone()?,
495    )));
496    let root = read_pickle(&mut reader, &PersistentIds::Storages(source.clone()))?;
497
498    // The storage keys, in the order their bytes follow.
499    let storage_keys = match read_header(&mut reader, "storage key list")? {
500        Object::List(keys) => keys
501            .iter()
502            .map(|key| key_string(key, "legacy storage key"))
503            .collect::<std::result::Result<Vec<_>, _>>()?,
504        other => {
505            return Err(PytorchError::InvalidFormat(format!(
506                "legacy storage key list must be a list, got {other:?}"
507            )));
508        }
509    };
510
511    let data_start = reader.stream_position()?;
512    let file_len = reader.seek(SeekFrom::End(0))?;
513    if let StorageSource::Legacy(legacy) = &*source {
514        legacy
515            .finish(&storage_keys, data_start, file_len)
516            .map_err(|e| PytorchError::InvalidFormat(e.to_string()))?;
517    }
518
519    let mut metadata = PytorchMetadata::for_format(FileFormat::Legacy);
520    metadata.total_data_size = Some((file_len - data_start) as usize);
521    Ok(Loaded { root, metadata })
522}
523
524fn load_plain_pickle(path: &Path) -> Result<Loaded> {
525    let mut reader = BufReader::new(File::open(path)?);
526    let root = read_pickle(&mut reader, &PersistentIds::Unavailable)?;
527    Ok(Loaded {
528        root,
529        metadata: PytorchMetadata::for_format(FileFormat::Pickle),
530    })
531}
532
533fn load_tar(path: &Path) -> Result<Loaded> {
534    let mut archive = tar::Archive::new(BufReader::new(File::open(path)?));
535    let mut entries: HashMap<String, Vec<u8>> = HashMap::new();
536
537    for entry in archive.entries().map_err(PytorchError::Tar)? {
538        let mut entry = entry.map_err(PytorchError::Tar)?;
539        let name = entry
540            .path()
541            .map_err(PytorchError::Tar)?
542            .to_string_lossy()
543            .trim_start_matches("./")
544            .to_string();
545        if matches!(
546            name.as_str(),
547            "sys_info" | "storages" | "tensors" | "pickle"
548        ) {
549            let mut data = Vec::new();
550            entry.read_to_end(&mut data).map_err(PytorchError::Tar)?;
551            entries.insert(name, data);
552        }
553    }
554
555    if let Some(sys_info) = entries.remove("sys_info") {
556        let sys_info = read_pickle(&mut Cursor::new(sys_info), &PersistentIds::Unavailable)?;
557        check_little_endian(&sys_info)?;
558    }
559    let mut take = |name: &str| {
560        entries
561            .remove(name)
562            .ok_or_else(|| PytorchError::InvalidFormat(format!("TAR file missing '{name}' entry")))
563    };
564    let storages = take("storages")?;
565    let tensors = take("tensors")?;
566    let pickle = take("pickle")?;
567
568    let total_data_size = storages.len();
569    let (tar_source, storage_info) = parse_tar_storages(storages)?;
570    let source = Arc::new(StorageSource::Tar(tar_source));
571    let tensors = parse_tar_tensors(&tensors, &storage_info, &source)?;
572    let root = read_pickle(&mut Cursor::new(pickle), &PersistentIds::Tensors(tensors))?;
573
574    let mut metadata = PytorchMetadata::for_format(FileFormat::Tar);
575    metadata.total_data_size = Some(total_data_size);
576    Ok(Loaded { root, metadata })
577}
578
579/// Per storage key, the element type and byte length a TAR `storages` entry declares.
580type TarStorageInfo = HashMap<String, (DType, usize)>;
581
582/// Parse the `storages` entry: a count, then per storage `(key, location, storage type)`,
583/// an `i64` element count and the bytes, then a list of storage views.
584fn parse_tar_storages(blob: Vec<u8>) -> Result<(TarSource, TarStorageInfo)> {
585    let mut cursor = Cursor::new(blob.as_slice());
586    let count = read_tar_count(&mut cursor, "storage")?;
587
588    // Every storage costs at least a pickle and a count, so the count bounds the blob.
589    let plausible = count.min(blob.len() / 8);
590    let mut layout: HashMap<String, (usize, usize)> = HashMap::with_capacity(plausible);
591    let mut info = TarStorageInfo::with_capacity(plausible);
592
593    for _ in 0..count {
594        let meta = read_pickle(&mut cursor, &PersistentIds::Unavailable)?;
595        let [key, _location, storage_type] = tar_fields(&meta, "storage metadata")?;
596        // The tables key storages by int `_cdata`; the main pickle's persistent ids are
597        // `str(_cdata)`. Both are normalized to strings.
598        let key = key_string(key, "TAR storage key")?;
599        let dtype = match storage_type {
600            Object::Class { name, .. } => storage_type_to_dtype(name.as_str())?,
601            other => {
602                return Err(PytorchError::InvalidFormat(format!(
603                    "TAR storage type must be a class, got {other:?}"
604                )));
605            }
606        };
607
608        let numel = read_tar_usize(&mut cursor, "storage element count")?;
609        let byte_len = numel.checked_mul(dtype.size()).ok_or_else(|| {
610            PytorchError::InvalidFormat(format!("TAR storage '{key}' byte length overflows usize"))
611        })?;
612        let offset = cursor.position() as usize;
613        let end = offset
614            .checked_add(byte_len)
615            .filter(|&end| end <= blob.len())
616            .ok_or_else(|| {
617                PytorchError::InvalidFormat(format!(
618                    "TAR storage '{key}' ({byte_len} bytes) extends beyond the storages entry"
619                ))
620            })?;
621
622        layout.insert(key.clone(), (offset, byte_len));
623        info.insert(key, (dtype, byte_len));
624        cursor.set_position(end as u64);
625    }
626
627    // Views of root storages: (view key, root key, element offset, element count).
628    // PyTorch reads this list unconditionally, so its absence means a truncated entry.
629    {
630        let views = read_pickle(&mut cursor, &PersistentIds::Unavailable).map_err(|e| {
631            PytorchError::InvalidFormat(format!("TAR storage views list missing: {e}"))
632        })?;
633        let Object::List(views) = views else {
634            return Err(PytorchError::InvalidFormat(format!(
635                "TAR storage views must be a list, got {views:?}"
636            )));
637        };
638        for view in views {
639            let [target, root, offset, numel] = tar_fields(&view, "storage view")?;
640            let target = key_string(target, "TAR storage view key")?;
641            let root = key_string(root, "TAR storage view root")?;
642            let offset = non_negative(offset, "TAR storage view offset")?;
643            let numel = non_negative(numel, "TAR storage view element count")?;
644            let &(root_offset, root_len) = layout.get(&root).ok_or_else(|| {
645                PytorchError::InvalidFormat(format!(
646                    "TAR storage view '{target}' refers to unknown storage '{root}'"
647                ))
648            })?;
649            let (dtype, _) = info[&root];
650            let element_size = dtype.size();
651            let (start, len) = offset
652                .checked_mul(element_size)
653                .zip(numel.checked_mul(element_size))
654                .filter(|&(start, len)| start.checked_add(len).is_some_and(|end| end <= root_len))
655                .ok_or_else(|| {
656                    PytorchError::InvalidFormat(format!(
657                        "TAR storage view '{target}' lies outside storage '{root}'"
658                    ))
659                })?;
660            layout.insert(target.clone(), (root_offset + start, len));
661            info.insert(target, (dtype, len));
662        }
663    }
664
665    Ok((TarSource::new(blob, layout), info))
666}
667
668/// Parse the `tensors` entry: a count, then per tensor `(key, storage key, tensor type)`,
669/// an `i32` rank, 4 unused bytes, `rank` sizes, `rank` strides and the storage offset.
670fn parse_tar_tensors(
671    blob: &[u8],
672    storage_info: &TarStorageInfo,
673    source: &Arc<StorageSource>,
674) -> Result<HashMap<String, Tensor>> {
675    let mut cursor = Cursor::new(blob);
676    let count = read_tar_count(&mut cursor, "tensor")?;
677    let mut tensors = HashMap::with_capacity(count.min(blob.len() / 8));
678
679    for _ in 0..count {
680        let meta = read_pickle(&mut cursor, &PersistentIds::Unavailable)?;
681        let [key, storage_key, _tensor_type] = tar_fields(&meta, "tensor metadata")?;
682        let key = key_string(key, "TAR tensor key")?;
683        let storage_key = key_string(storage_key, "TAR tensor storage key")?;
684
685        let rank = cursor.read_i32::<LittleEndian>()?;
686        let rank = usize::try_from(rank)
687            .ok()
688            .filter(|&rank| rank <= 64)
689            .ok_or_else(|| {
690                PytorchError::InvalidFormat(format!("TAR tensor '{key}' has invalid rank {rank}"))
691            })?;
692        cursor.read_i32::<LittleEndian>()?; // Padding: the rank was once written as 8 bytes.
693        let shape = (0..rank)
694            .map(|_| read_tar_usize(&mut cursor, "tensor shape"))
695            .collect::<Result<Vec<_>>>()?;
696        let stride = (0..rank)
697            .map(|_| read_tar_usize(&mut cursor, "tensor stride"))
698            .collect::<Result<Vec<_>>>()?;
699        let storage_offset = read_tar_usize(&mut cursor, "tensor storage offset")?;
700
701        let &(dtype, byte_len) = storage_info.get(&storage_key).ok_or_else(|| {
702            PytorchError::InvalidFormat(format!(
703                "TAR tensor '{key}' refers to unknown storage '{storage_key}'"
704            ))
705        })?;
706        let storage = StorageRef {
707            source: source.clone(),
708            key: storage_key,
709            dtype: Some(dtype),
710            // A TAR view is resolved into its own layout entry above, so `byte_len` is
711            // already the view's length and nothing is reachable past it.
712            reachable_bytes: byte_len,
713            view_offset: 0,
714        };
715        let tensor = build_tensor(storage, dtype, storage_offset, shape, stride)?;
716        tensors.insert(key, tensor);
717    }
718
719    Ok(tensors)
720}
721
722fn read_tar_count(cursor: &mut Cursor<&[u8]>, what: &str) -> Result<usize> {
723    let count = read_pickle(cursor, &PersistentIds::Unavailable)?;
724    Ok(non_negative(&count, &format!("TAR {what} count"))?)
725}
726
727/// Read one of the `i64` fields in the TAR binary tables as a `usize`.
728fn read_tar_usize(cursor: &mut Cursor<&[u8]>, what: &str) -> Result<usize> {
729    let value = cursor.read_i64::<LittleEndian>()?;
730    usize::try_from(value).map_err(|_| {
731        PytorchError::InvalidFormat(format!("TAR {what} must be non-negative, got {value}"))
732    })
733}
734
735/// View a TAR metadata pickle as a tuple of exactly `N` fields.
736fn tar_fields<'a, const N: usize>(obj: &'a Object, what: &str) -> Result<&'a [Object; N]> {
737    match obj {
738        Object::Tuple(fields) => <&[Object; N]>::try_from(fields.as_slice()).map_err(|_| {
739            PytorchError::InvalidFormat(format!(
740                "TAR {what} must have {N} fields, got {}",
741                fields.len()
742            ))
743        }),
744        other => Err(PytorchError::InvalidFormat(format!(
745            "TAR {what} must be a tuple, got {other:?}"
746        ))),
747    }
748}
749
750/// Refuse a legacy file that was not written on a little-endian machine.
751///
752/// PyTorch always records a `little_endian` bool here, so a header without a usable one is
753/// malformed rather than a file to guess about: reading big-endian storages as
754/// little-endian would silently transpose every value.
755fn check_little_endian(sys_info: &Object) -> Result<()> {
756    let Object::Dict(dict) = sys_info else {
757        return Err(PytorchError::InvalidFormat(format!(
758            "Legacy sys_info must be a dict, got {sys_info:?}"
759        )));
760    };
761    match dict.get("little_endian") {
762        Some(Object::Bool(true)) => Ok(()),
763        Some(Object::Bool(false)) => Err(big_endian_error()),
764        other => Err(PytorchError::InvalidFormat(format!(
765            "Legacy sys_info must carry a little_endian bool, got {other:?}"
766        ))),
767    }
768}
769
770fn big_endian_error() -> PytorchError {
771    PytorchError::InvalidFormat(
772        "Big-endian PyTorch files are not supported. The file was saved on a big-endian system and requires byte order conversion."
773            .to_string(),
774    )
775}
776
777// ---------------------------------------------------------------------------------------------
778// Root object handling
779// ---------------------------------------------------------------------------------------------
780
781/// Return `root`, or the entry under `key` when one is given.
782fn select_top_level(root: Object, key: Option<&str>) -> Result<Object> {
783    let Some(key) = key else {
784        return Ok(root);
785    };
786    let Object::Dict(mut dict) = root else {
787        return Err(PytorchError::InvalidFormat(format!(
788            "Cannot select top-level key '{key}': the file's root is not a dictionary"
789        )));
790    };
791    dict.remove(key).ok_or_else(|| {
792        let mut available: Vec<_> = dict.keys().cloned().collect();
793        available.sort();
794        PytorchError::KeyNotFound(format!(
795            "'{key}' not found. Available top-level keys: {available:?}"
796        ))
797    })
798}
799
800fn extract_tensors_at(
801    root: Object,
802    top_level_key: Option<&str>,
803) -> Result<HashMap<String, Tensor>> {
804    let Object::Dict(dict) = select_top_level(root, top_level_key)? else {
805        return Err(PytorchError::InvalidFormat(match top_level_key {
806            Some(key) => format!("Top-level key '{key}' does not hold a dictionary"),
807            None => "Expected a dictionary at the root of the PyTorch file, but found a different type. The file may be a full model save rather than a state_dict.".to_string(),
808        }));
809    };
810    Ok(extract_tensors(dict))
811}
812
813/// Convert an internal object to the public [`PickleValue`].
814///
815/// Tuples become lists, and anything without a JSON-like counterpart (classes, tensors,
816/// storages, uninterpreted objects) becomes [`PickleValue::Unsupported`] naming its type.
817fn to_pickle_value(obj: Object) -> PickleValue {
818    match obj {
819        Object::None => PickleValue::None,
820        Object::Bool(b) => PickleValue::Bool(b),
821        Object::Int(i) => PickleValue::Int(i),
822        Object::Float(f) => PickleValue::Float(f),
823        Object::String(s) => PickleValue::String(s),
824        Object::Bytes(b) => PickleValue::Bytes(b),
825        Object::List(items) | Object::Tuple(items) => {
826            PickleValue::List(items.into_iter().map(to_pickle_value).collect())
827        }
828        Object::Dict(dict) => PickleValue::Dict(
829            dict.into_iter()
830                .map(|(k, v)| (k, to_pickle_value(v)))
831                .collect(),
832        ),
833        Object::Class { .. } | Object::Storage(_) | Object::Opaque(_) | Object::Tensor(_) => {
834            PickleValue::Unsupported(obj.python_type_name())
835        }
836    }
837}
838
839/// Convert a [`PickleValue`] to a [`NestedValue`] for deserialization.
840fn to_nested_value(value: PickleValue) -> NestedValue {
841    match value {
842        PickleValue::None => NestedValue::Default(None),
843        PickleValue::Bool(b) => NestedValue::Bool(b),
844        PickleValue::Int(i) => NestedValue::I64(i),
845        PickleValue::Float(f) => NestedValue::F64(f),
846        PickleValue::String(s) => NestedValue::String(s),
847        PickleValue::List(list) => {
848            NestedValue::Vec(list.into_iter().map(to_nested_value).collect())
849        }
850        PickleValue::Dict(dict) => NestedValue::Map(
851            dict.into_iter()
852                .map(|(k, v)| (k, to_nested_value(v)))
853                .collect(),
854        ),
855        PickleValue::Bytes(data) => {
856            NestedValue::Vec(data.into_iter().map(NestedValue::U8).collect())
857        }
858        PickleValue::Unsupported(type_name) => NestedValue::Unsupported(type_name),
859    }
860}
pytorch_reader/reader.rs

pytorch_reader/
reader.rs
