rten 0.25.0

//! Functions for loading tensor data stored externally to the main model file.
//!
//! This is used for ONNX models when `TensorProto`s reference external data
//! files.

use std::cell::RefCell;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::ops::Range;
use std::path::{Component, Path, PathBuf};
use std::sync::Arc;

#[cfg(feature = "mmap")]
use memmap2::Mmap;

use super::load_error::{LoadError, LoadErrorImpl};
use crate::constant_storage::ConstantStorage;

/// Specifies the location of tensor data which is stored externally from the
/// main model file.
#[derive(Clone, Debug)]
pub struct DataLocation {
    /// Name of the external data file.
    pub path: String,

    /// Offset of the start of the tensor data in bytes.
    pub offset: u64,

    /// Length of the tensor data in bytes.
    pub length: u64,
}

/// A slice of a shared buffer, where the slice contains the data for one
/// tensor.
#[derive(Debug)]
pub struct DataSlice {
    /// The shared buffer. This may contain data for one or multiple tensors
    /// and may be a `Vec<u8>`, a memory-mapped file or static slice.
    pub storage: Arc<ConstantStorage>,

    /// The range of bytes within `storage` that contain the tensor data.
    pub bytes: Range<usize>,
}

impl DataSlice {
    pub fn data(&self) -> &[u8] {
        &self.storage.data()[self.bytes.clone()]
    }
}

#[derive(Debug)]
enum ExternalDataErrorKind {
    /// An IO error occurred when accessing the external file.
    IoError(std::io::Error),

    /// The length of the external data is too large.
    InvalidLength,

    /// An invalid path was specified.
    InvalidPath,

    /// External data is not supported in the current environment.
    NotSupported,

    /// The length of the external data file is too short for the offset and
    /// length of the external data.
    TooShort {
        /// Minimum length the file would need to be in bytes.
        required_len: usize,
        /// Actual length of the file in bytes.
        actual_len: usize,
    },

    /// The external data file path is disallowed.
    DisallowedPath,
}

/// Errors reading tensor data from an external file.
#[derive(Debug)]
pub struct ExternalDataError {
    /// The path involved in the error.
    ///
    /// This should be the absolute path of the data file when available, but
    /// for some errors it will be the path of the model or the relative data
    /// file path specified in the model.
    path: PathBuf,
    /// The type of error that occurred.
    kind: ExternalDataErrorKind,
}

impl ExternalDataError {
    fn new(path: &Path, kind: ExternalDataErrorKind) -> Self {
        Self {
            path: path.to_path_buf(),
            kind,
        }
    }

    fn from_io_error(path: &Path, err: std::io::Error) -> Self {
        Self {
            path: path.to_path_buf(),
            kind: ExternalDataErrorKind::IoError(err),
        }
    }
}

impl std::fmt::Display for ExternalDataError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        type Kind = ExternalDataErrorKind;

        write!(f, "for path \"{}\": ", self.path.display())?;
        match &self.kind {
            Kind::IoError(err) => write!(f, "io error: {}", err),
            Kind::InvalidLength => write!(f, "invalid data length"),
            Kind::InvalidPath => write!(f, "invalid path"),
            Kind::NotSupported => write!(f, "external data not supported"),
            Kind::TooShort {
                required_len,
                actual_len,
            } => write!(
                f,
                "file too short. required {} actual {}",
                required_len, actual_len
            ),
            Kind::DisallowedPath => {
                write!(f, "disallowed path")
            }
        }
    }
}

impl std::error::Error for ExternalDataError {}

impl From<ExternalDataError> for LoadError {
    fn from(err: ExternalDataError) -> LoadError {
        LoadErrorImpl::ExternalDataError(Box::new(err)).into()
    }
}

/// Trait for loading data from an external file.
pub trait DataLoader {
    /// Load data from the file and offset specified by `location`.
    fn load(&self, location: &DataLocation) -> Result<DataSlice, ExternalDataError>;
}

/// Check if `path` is an allowed path for external data files for a given
/// model path.
///
/// Any data loaded from an external data file can potentially be returned in
/// inference results (directly or indirectly). Hence some measures are taken to
/// prevent loading of data from files not intended for this. The ONNX
/// documentation states that the only restriction is that parent directory
/// components ("..") are disallowed
/// (https://onnx.ai/onnx/repo-docs/ExternalData.html#external-data-field). This
/// implementation imposes additional restrictions:
///
///  - The file must have one of the known extensions used for tensor data
///    ("data", "onnx_data", "onnx_data_N" etc.)
///  - The file must be a relative path with only a filename component (ie. it
///    must be located in the same directory as the model)
fn is_allowed_external_data_path(path: &Path) -> bool {
    // Data file path must be relative and consist only of a filename.
    let mut components = path.components();
    let Some(Component::Normal(_)) = components.next() else {
        return false;
    };
    if components.next().is_some() {
        return false;
    }

    // Check for allowed extension. The most common extensions used are "data"
    // or "onnx_data", but large files are sometimes split into pieces eg.
    // "onnx_data_N".
    match path.extension().and_then(|ext| ext.to_str()) {
        Some(ext) if ext.starts_with("data") => true,
        Some(ext) if ext.starts_with("onnx_data") => true,
        _ => false,
    }
}

/// External data loader that uses standard file IO.
pub struct FileLoader {
    /// Path to directory containing external data.
    dir_path: PathBuf,

    /// Map of external data file name to (resolved file path, open file).
    files: RefCell<HashMap<PathBuf, (PathBuf, File)>>,
}

impl FileLoader {
    /// Create an external data loader which loads data for the model file
    /// specified by `model_path`.
    ///
    /// Data file paths will be resolved relative to the directory containing
    /// `model_path`.
    pub fn new(model_path: &Path) -> Result<Self, ExternalDataError> {
        let dir_path = dir_path_from_model_path(model_path)?;

        Ok(Self {
            dir_path,
            files: HashMap::new().into(),
        })
    }

    fn read(&self, location: &DataLocation) -> Result<Vec<u8>, ExternalDataError> {
        // Errors at the top of this function refer to the relative path in `location`.
        let make_err = |kind| ExternalDataError::new(location.path.as_ref(), kind);

        // On a big-endian system we'd need to perform byte-swapping while loading.
        if cfg!(target_endian = "big") {
            return Err(make_err(ExternalDataErrorKind::NotSupported));
        }

        // On a 32-bit systems assume we can't load more than 2GB of data.
        if location.length > isize::MAX as u64 {
            return Err(make_err(ExternalDataErrorKind::InvalidLength));
        }
        let vec_len = location.length as usize;

        let mut files = self.files.borrow_mut();
        let (file_path, file) =
            get_or_open_file(&mut files, &self.dir_path, Path::new(&location.path))?;

        // Subsequent errors refer to the resolved data file path.
        let make_err = |kind| ExternalDataError::new(file_path, kind);
        let make_io_err = |err| ExternalDataError::from_io_error(file_path, err);

        file.seek(SeekFrom::Start(location.offset))
            .map_err(make_io_err)?;

        // Ideally we would fill the buffer in one call via [`Read::read_buf`].
        // Since that API is not stabilized yet, we fill in small chunks, which
        // requires extra copying.
        let mut remaining = vec_len;
        let mut buf = Vec::with_capacity(remaining);

        // Buffer size chosen to match BufReader's default.
        const TMP_SIZE: usize = 8192;
        let mut tmp = [0u8; TMP_SIZE];

        loop {
            let tmp_size = remaining.min(TMP_SIZE);
            let n_read = read_fill(file, &mut tmp[..tmp_size]).map_err(make_io_err)?;
            let chunk = &tmp[..n_read];
            remaining -= chunk.len();
            buf.extend_from_slice(chunk);

            if n_read < tmp.len() || remaining == 0 {
                break;
            }
        }

        if buf.len() != vec_len {
            return Err(make_err(ExternalDataErrorKind::TooShort {
                required_len: vec_len,
                actual_len: buf.len(),
            }));
        }

        Ok(buf)
    }
}

/// Read from `src` repeatedly until `buf` is filled or we reach the end of the
/// file.
fn read_fill<R: Read>(src: &mut R, buf: &mut [u8]) -> std::io::Result<usize> {
    let mut total_read = 0;
    loop {
        let n = src.read(&mut buf[total_read..])?;
        total_read += n;
        if n == 0 || total_read == buf.len() {
            break;
        }
    }
    Ok(total_read)
}

impl DataLoader for FileLoader {
    fn load(&self, location: &DataLocation) -> Result<DataSlice, ExternalDataError> {
        let bytes = self.read(location)?;
        Ok(DataSlice {
            bytes: 0..bytes.len(),
            storage: Arc::new(ConstantStorage::Buffer(bytes)),
        })
    }
}

fn get_or_open_file<'a>(
    files: &'a mut HashMap<PathBuf, (PathBuf, File)>,
    dir_path: &Path,
    data_path: &Path,
) -> Result<&'a mut (PathBuf, File), ExternalDataError> {
    let data_path = Path::new(data_path);
    if !is_allowed_external_data_path(data_path) {
        return Err(ExternalDataError::new(
            data_path,
            ExternalDataErrorKind::DisallowedPath,
        ));
    }

    // Check if we already opened the file.
    if files.get(data_path).is_none() {
        let mut file_path = dir_path.to_path_buf();
        file_path.push(data_path);
        let file = File::open(&file_path)
            .map_err(|err| ExternalDataError::from_io_error(&file_path, err))?;
        files.insert(data_path.into(), (file_path, file));
    }

    Ok(files.get_mut(data_path).unwrap())
}

fn dir_path_from_model_path(model_path: &Path) -> Result<PathBuf, ExternalDataError> {
    // Resolve the path now to avoid the possibility of loading data from
    // an unexpected location if `model_path` is relative and the current
    // working directory changes before a data file is loaded.
    let model_path = if !cfg!(target_arch = "wasm32") {
        model_path
            .canonicalize()
            .map_err(|err| ExternalDataError::from_io_error(model_path, err))?
    } else {
        // On WASM / WASI `Path::canonicalize` is not available.
        model_path.to_path_buf()
    };

    if !model_path.is_file() {
        return Err(ExternalDataError::new(
            &model_path,
            ExternalDataErrorKind::InvalidPath,
        ));
    }
    let dir_path = model_path
        .parent()
        // Since `model_path` is a file path, it cannot be the root ("/").
        .expect("should have parent dir")
        .to_path_buf();
    Ok(dir_path)
}

/// External data loader that uses memory mapping.
///
/// # Alignment requirements
///
/// The ONNX Protocol Buffers schema and external data documentation state that
/// data offsets for individual tensors should be aligned to the page size,
/// which is the granularity of the `offset` argument to `mmap`. RTen only
/// requires offsets to be a multiple of the tensor element type's alignment.
/// This is because RTen creates only one memory map for each external data
/// file, covering the whole file. Each tensor using data from that file will
/// then share the memory mapping.
#[cfg(feature = "mmap")]
pub struct MmapLoader {
    /// Path to directory containing external data.
    dir_path: PathBuf,

    /// Map of filename to (resolved path, open mmap-ed content).
    mmaps: RefCell<HashMap<PathBuf, (PathBuf, Arc<ConstantStorage>)>>,
}

#[cfg(feature = "mmap")]
impl MmapLoader {
    /// Create a data loader which will use memory mapping to load data from
    /// files in the same directory as `model_path`.
    ///
    /// One memory map will be created per external file. This map will remain
    /// open as long as the `MmapLoader` or any tensors using data from it are
    /// still open.
    ///
    /// # Safety
    ///
    /// This method is marked as unsafe because truncating the file on disk
    /// while the file is mapped could cause undefined behavior. Applications
    /// must decide this is an acceptable risk for their use. See the notes for
    /// [`Model::load_mmap`](crate::model::Model::load_mmap).
    pub unsafe fn new(model_path: &Path) -> Result<Self, ExternalDataError> {
        let dir_path = dir_path_from_model_path(model_path)?;

        Ok(Self {
            dir_path,
            mmaps: HashMap::new().into(),
        })
    }

    fn get_or_open_mmap<'a>(
        &self,
        mmaps: &'a mut HashMap<PathBuf, (PathBuf, Arc<ConstantStorage>)>,
        data_path: &Path,
    ) -> Result<&'a (PathBuf, Arc<ConstantStorage>), ExternalDataError> {
        let data_path = Path::new(data_path);
        if !is_allowed_external_data_path(data_path) {
            return Err(ExternalDataError::new(
                data_path,
                ExternalDataErrorKind::DisallowedPath,
            ));
        }

        // Check if we already opened the file.
        if mmaps.get(data_path).is_none() {
            let mut file_path = self.dir_path.to_path_buf();
            file_path.push(data_path);

            let make_io_err = |err| ExternalDataError::from_io_error(&file_path, err);
            let file = File::open(&file_path).map_err(make_io_err)?;

            // Safety: By constructing an instance of `Self`, the caller has
            // accepted the risks that come with mmap.
            let mmap = unsafe { Mmap::map(&file) }.map_err(make_io_err)?;

            let storage = Arc::new(ConstantStorage::Mmap(mmap));
            mmaps.insert(data_path.into(), (file_path, storage));
        }

        Ok(mmaps.get(data_path).unwrap())
    }
}

#[cfg(feature = "mmap")]
impl DataLoader for MmapLoader {
    fn load(&self, location: &DataLocation) -> Result<DataSlice, ExternalDataError> {
        let mut mmaps = self.mmaps.borrow_mut();
        let (file_path, storage) = self.get_or_open_mmap(&mut mmaps, Path::new(&location.path))?;

        let end_offset = location.offset.saturating_add(location.length);
        if end_offset > storage.data().len() as u64 {
            return Err(ExternalDataError::new(
                file_path,
                ExternalDataErrorKind::TooShort {
                    required_len: end_offset as usize,
                    actual_len: storage.data().len(),
                },
            ));
        }

        Ok(DataSlice {
            storage: storage.clone(),
            bytes: location.offset as usize..location.offset as usize + location.length as usize,
        })
    }
}

/// A [`DataLoader`] which loads data from in-memory buffers.
///
/// This supports loading models with external data in contexts where a file
/// system is not available (eg. a browser) or inconvenient to use (eg. in
/// tests).
pub struct MemLoader(HashMap<String, Arc<ConstantStorage>>);

impl MemLoader {
    pub fn new(map: HashMap<String, Arc<ConstantStorage>>) -> Self {
        Self(map)
    }

    #[cfg(test)]
    pub fn from_entries(entries: impl IntoIterator<Item = (String, Vec<u8>)>) -> Self {
        let map = entries
            .into_iter()
            .map(|(path, buf)| {
                let storage = Arc::new(ConstantStorage::Buffer(buf));
                (path, storage)
            })
            .collect();
        Self(map)
    }
}

impl DataLoader for MemLoader {
    fn load(&self, location: &DataLocation) -> Result<DataSlice, ExternalDataError> {
        let make_err = |kind| ExternalDataError::new(location.path.as_ref(), kind);

        // The restrictions on allowed data paths don't matter for the in-memory
        // data loader, but we apply them for consistency with other loaders.
        if !is_allowed_external_data_path(Path::new(&location.path)) {
            return Err(make_err(ExternalDataErrorKind::DisallowedPath));
        }
        let Some(storage) = self.0.get(&location.path) else {
            // Error message chosen to match the file loader.
            return Err(ExternalDataError::from_io_error(
                location.path.as_ref(),
                std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    "No such file or directory".to_string(),
                ),
            ));
        };
        let end_offset = location.offset + location.length;
        if end_offset > storage.data().len() as u64 {
            return Err(make_err(ExternalDataErrorKind::TooShort {
                required_len: end_offset as usize,
                actual_len: storage.data().len(),
            }));
        }

        let bytes = (location.offset as usize)..end_offset as usize;
        Ok(DataSlice {
            storage: storage.clone(),
            bytes,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::panic::RefUnwindSafe;
    use std::path::{Path, PathBuf};
    use std::sync::Arc;

    use super::{DataLoader, DataLocation, ExternalDataError, FileLoader, MemLoader};
    use crate::constant_storage::ConstantStorage;
    use rten_testing::TestCases;

    fn temp_dir() -> PathBuf {
        if cfg!(target_arch = "wasm32") {
            // temp_dir is not available on WASM / WASI, so just use the current
            // directory.
            PathBuf::new()
        } else {
            std::env::temp_dir()
        }
    }

    struct TempFile {
        path: PathBuf,
    }

    impl TempFile {
        fn new(name: impl AsRef<Path>, content: &[u8]) -> std::io::Result<Self> {
            let mut path = temp_dir();
            path.push(name);
            std::fs::write(&path, content)?;
            Ok(Self { path })
        }

        fn path(&self) -> &Path {
            &self.path
        }
    }

    impl Drop for TempFile {
        fn drop(&mut self) {
            std::fs::remove_file(&self.path).expect("should remove file");
        }
    }

    // Run common tests for `DataLoader` impls. The `base_name` must be unique
    // for each test.
    fn test_loader<L: DataLoader>(
        base_name: &str,
        make_loader: impl Fn(&Path) -> Result<L, ExternalDataError> + RefUnwindSafe,
    ) {
        let bytes: Vec<u8> = (0..32).collect();
        let model_file = TempFile::new(format!("{base_name}.onnx"), &[]).unwrap();
        let data_file = TempFile::new(format!("{base_name}.onnx.data"), &bytes).unwrap();

        let data_filename = data_file
            .path()
            .file_name()
            .unwrap()
            .to_string_lossy()
            .to_string();

        #[derive(Debug)]
        struct Case {
            location: DataLocation,
            expected: Result<Vec<u8>, String>,
        }

        let cases = [
            // Part of file
            Case {
                location: DataLocation {
                    path: data_filename.clone(),
                    offset: 8,
                    length: 8,
                },
                expected: Ok(bytes[8..16].into()),
            },
            // Full file
            Case {
                location: DataLocation {
                    path: data_filename.clone(),
                    offset: 0,
                    length: bytes.len() as u64,
                },
                expected: Ok(bytes.clone()),
            },
            // Empty path
            Case {
                location: DataLocation {
                    path: String::new(),
                    offset: 0,
                    length: 0,
                },
                expected: Err("disallowed path".into()),
            },
            // Path containing parent directory
            Case {
                location: DataLocation {
                    path: "../foo.data".into(),
                    offset: 0,
                    length: 0,
                },
                expected: Err("disallowed path".into()),
            },
            // Path with disallowed extension
            Case {
                location: DataLocation {
                    path: "not_a_data_file.md".into(),
                    offset: 0,
                    length: 0,
                },
                expected: Err("disallowed path".into()),
            },
            // File does not exist
            Case {
                location: DataLocation {
                    path: "file_does_not_exist.data".into(),
                    offset: 0,
                    length: 0,
                },
                expected: Err("No such file or directory".into()),
            },
            // Range extends beyond end of file
            Case {
                location: DataLocation {
                    path: data_filename,
                    offset: 0,
                    length: 36,
                },
                expected: Err("file too short".into()),
            },
        ];

        cases.test_each(|case| {
            let loader = make_loader(model_file.path()).unwrap();
            let data = loader.load(&case.location).map_err(|e| e.to_string());
            match (&data, &case.expected) {
                (Ok(actual), Ok(expected)) => assert_eq!(actual.data(), expected),
                (Err(actual), Err(expected)) => assert!(
                    actual.contains(expected),
                    "{} does not contain {}",
                    actual,
                    expected
                ),
                (actual, expected) => assert_eq!(actual.is_ok(), expected.is_ok()),
            }
        });
    }

    #[test]
    fn test_file_loader() {
        test_loader("test_file_loader", FileLoader::new)
    }

    #[cfg(feature = "mmap")]
    #[test]
    fn test_mmap_loader() {
        use super::MmapLoader;
        test_loader("test_mmap_loader", |model_path| unsafe {
            MmapLoader::new(model_path)
        })
    }

    #[test]
    fn test_mem_loader() {
        test_loader("test_mem_loader", |_model_path| {
            let mut map = HashMap::new();
            let buf = (0..32).collect();
            let storage = Arc::new(ConstantStorage::Buffer(buf));
            map.insert("test_mem_loader.onnx.data".to_string(), storage);
            Ok(MemLoader::new(map))
        })
    }
}