spark_runtime/weights/gguf/
expert_reads.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3//! The file reads behind the expert paths: positional reads that any
4//! number of threads may issue on one descriptor, the page-cache drop after
5//! a buffered miss, and the `O_DIRECT` side (the descriptor, the segment
6//! description, the read that tolerates a window past the end of the file).
7//! Split from `expert_stream.rs` (500-LoC cap).
8
9use std::fs::File;
10use std::path::Path;
11
12use anyhow::{Result, bail};
13
14/// Positional read of exactly `dst.len()` bytes at `offset`. No file position
15/// is shared, so any number of threads may read one `File` at once.
16pub fn pread(file: &File, offset: u64, dst: &mut [u8]) -> Result<()> {
17    #[cfg(unix)]
18    {
19        use anyhow::Context;
20        use std::os::unix::fs::FileExt;
21        file.read_exact_at(dst, offset)
22            .with_context(|| format!("pread {} bytes at {offset}", dst.len()))
23    }
24    #[cfg(not(unix))]
25    {
26        let _ = (file, offset, dst);
27        bail!("expert streaming needs positional reads (unix)")
28    }
29}
30
31/// `pread`, then drop the range from the page cache: the bytes now live in
32/// the arena and the cache copy is dead weight. With a 100 GiB page-locked
33/// arena on a 121 GiB Spark, leaving 12 MiB of cache behind every miss made
34/// the kernel reclaim into swap in the middle of a step (single steps of 2 to
35/// 18 s on the 09-19 standard). `ATLAS_DS41_KEEP_PAGE_CACHE=1` keeps the old
36/// behaviour.
37pub fn pread_uncached(file: &File, offset: u64, dst: &mut [u8]) -> Result<()> {
38    pread(file, offset, dst)?;
39    #[cfg(target_os = "linux")]
40    if !keep_page_cache() {
41        use std::os::unix::io::AsRawFd;
42        // SAFETY: an advisory call on an open descriptor; the kernel checks
43        // the range.
44        unsafe {
45            libc::posix_fadvise(
46                file.as_raw_fd(),
47                offset as libc::off_t,
48                dst.len() as libc::off_t,
49                libc::POSIX_FADV_DONTNEED,
50            );
51        }
52    }
53    Ok(())
54}
55
56#[cfg(target_os = "linux")]
57fn keep_page_cache() -> bool {
58    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
59    *V.get_or_init(|| std::env::var("ATLAS_DS41_KEEP_PAGE_CACHE").is_ok_and(|v| v == "1"))
60}
61
62/// Open `path` for direct reads (offsets, lengths and buffers aligned to
63/// the 512 B logical block by the caller). `ATLAS_DS41_DIRECT_READS=0`
64/// keeps every read on the page cache.
65pub(super) fn open_direct(path: &Path) -> Option<File> {
66    if std::env::var("ATLAS_DS41_DIRECT_READS").is_ok_and(|v| v == "0") {
67        return None;
68    }
69    #[cfg(target_os = "linux")]
70    {
71        use std::os::unix::fs::OpenOptionsExt;
72        std::fs::OpenOptions::new()
73            .read(true)
74            .custom_flags(libc::O_DIRECT)
75            .open(path)
76            .ok()
77    }
78    #[cfg(not(target_os = "linux"))]
79    {
80        let _ = path;
81        None
82    }
83}
84
85/// One file range of an expert's slot image, for the direct miss path:
86/// bytes `slot_off..slot_off + len` of the slot are `file_off..` of `shard`.
87#[derive(Clone, Copy, Debug, PartialEq, Eq)]
88pub struct DirectSeg {
89    pub shard: usize,
90    pub file_off: u64,
91    pub slot_off: usize,
92    pub len: usize,
93}
94
95/// `pread` into `dst` until at least `need` bytes have landed (`need <=
96/// dst.len()`): the direct miss path reads 512 B windows whose tail may run
97/// past the end of the file, where only the tensor's own bytes are owed.
98pub fn pread_at_least(file: &File, offset: u64, dst: &mut [u8], need: usize) -> Result<()> {
99    #[cfg(unix)]
100    {
101        use anyhow::Context;
102        use std::os::unix::fs::FileExt;
103        let mut done = 0usize;
104        while done < need {
105            let n = file
106                .read_at(&mut dst[done..], offset + done as u64)
107                .with_context(|| format!("pread {} bytes at {offset}", dst.len()))?;
108            if n == 0 {
109                bail!("pread at {offset}: end of file after {done} of {need} bytes",);
110            }
111            done += n;
112        }
113        Ok(())
114    }
115    #[cfg(not(unix))]
116    {
117        let _ = (file, offset, dst, need);
118        bail!("expert streaming needs positional reads (unix)")
119    }
120}
spark_runtime/weights/gguf/expert_reads.rs

spark_runtime/weights/gguf/
expert_reads.rs
