Skip to content

Commit f8172fe

Browse files
committed
perf: pre-warm reconstruction cache at open()
warm_reconstruction_cache() fires a background tokio task at open() that fetches the full-file reconstruction plan from CAS. A watch::Receiver signals completion; the first read() awaits it only if the warmup is not yet done. For workloads with any gap between open() and the first read (safetensors header parsing, Python overhead), the CAS round-trip overlaps with userspace work and adds zero latency to reads. FUSE open replies with FOPEN_KEEP_CACHE so the kernel page cache is reused across re-opens of the same file.
1 parent f5e7ee6 commit f8172fe

3 files changed

Lines changed: 53 additions & 16 deletions

File tree

src/fuse.rs

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -462,17 +462,8 @@ pub fn mount_fuse(
462462
config.mount_options.push(fuser::MountOption::RO);
463463
}
464464
config.acl = fuser::SessionACL::All;
465-
#[cfg(not(target_os = "macos"))]
466-
{
467-
config.clone_fd = true;
468-
config.n_threads = Some(max_threads);
469-
}
470-
#[cfg(target_os = "macos")]
471-
{
472-
// macFUSE only supports single-threaded session (fuser limitation).
473-
let _ = max_threads;
474-
config.n_threads = Some(1);
475-
}
465+
config.clone_fd = true;
466+
config.n_threads = Some(max_threads);
476467

477468
let session = match fuser::Session::new(adapter, mount_point, &config) {
478469
Ok(s) => s,

src/virtual_fs/mod.rs

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1244,13 +1244,28 @@ impl VirtualFs {
12441244
}
12451245

12461246
/// Allocate a lazy file handle backed by a prefetch buffer.
1247+
/// Starts reconstruction cache warmup in the background — open() returns immediately.
1248+
/// The first read() waits for the warmup only if it has not completed yet.
12471249
fn open_lazy(&self, xet_hash: String, size: u64) -> VirtualFsResult<u64> {
1250+
let warm = if !xet_hash.is_empty() {
1251+
let (tx, rx) = tokio::sync::watch::channel(false);
1252+
let sessions = self.xet_sessions.clone();
1253+
let hash = xet_hash.clone();
1254+
tokio::spawn(async move {
1255+
sessions.warm_reconstruction_cache(&hash).await;
1256+
// Errors silently dropped — worst case the first read pays the CAS round-trip.
1257+
let _ = tx.send(true);
1258+
});
1259+
Some(rx)
1260+
} else {
1261+
None
1262+
};
12481263
let prefetch = Arc::new(tokio::sync::Mutex::new(PrefetchState::new(xet_hash, size)));
12491264
let file_handle = self.alloc_file_handle();
12501265
self.open_files
12511266
.write()
12521267
.expect("open_files poisoned")
1253-
.insert(file_handle, OpenFile::Lazy { prefetch });
1268+
.insert(file_handle, OpenFile::Lazy { prefetch, warm });
12541269
Ok(file_handle)
12551270
}
12561271

@@ -1380,8 +1395,9 @@ impl VirtualFs {
13801395
let files = self.open_files.read().expect("open_files poisoned");
13811396
match files.get(&file_handle) {
13821397
Some(OpenFile::Local { file, .. }) => ReadTarget::LocalFd(file.clone()),
1383-
Some(OpenFile::Lazy { prefetch, .. }) => ReadTarget::Remote {
1398+
Some(OpenFile::Lazy { prefetch, warm }) => ReadTarget::Remote {
13841399
prefetch: prefetch.clone(),
1400+
warm: warm.clone(),
13851401
},
13861402
// write-only, not readable
13871403
Some(OpenFile::Streaming { .. }) => return Err(libc::EBADF), // write-only, not readable
@@ -1411,7 +1427,12 @@ impl VirtualFs {
14111427
Ok((buf.freeze(), eof))
14121428
}
14131429
}
1414-
ReadTarget::Remote { prefetch } => {
1430+
ReadTarget::Remote { prefetch, mut warm } => {
1431+
// Wait for background reconstruction cache warmup, if not already done.
1432+
// Returns immediately on subsequent reads (watch value stays true).
1433+
if let Some(ref mut rx) = warm {
1434+
let _ = rx.wait_for(|v| *v).await;
1435+
}
14151436
let mut prefetch_state = prefetch.lock().await;
14161437

14171438
let file_size = prefetch_state.file_size;
@@ -2811,6 +2832,9 @@ enum OpenFile {
28112832
/// Lazy remote — data fetched on-demand with adaptive prefetch buffer.
28122833
Lazy {
28132834
prefetch: Arc<tokio::sync::Mutex<PrefetchState>>,
2835+
/// Background reconstruction cache warmup. Becomes `true` when the full-file plan
2836+
/// is cached. `None` for empty files (no warmup needed).
2837+
warm: Option<tokio::sync::watch::Receiver<bool>>,
28142838
},
28152839
/// Streaming append-only writer (default write mode).
28162840
Streaming { ino: u64, channel: Arc<StreamingChannel> },
@@ -2886,6 +2910,7 @@ enum ReadTarget {
28862910
LocalFd(Arc<File>),
28872911
Remote {
28882912
prefetch: Arc<tokio::sync::Mutex<PrefetchState>>,
2913+
warm: Option<tokio::sync::watch::Receiver<bool>>,
28892914
},
28902915
}
28912916

src/xet.rs

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@ use std::path::{Path, PathBuf};
22
use std::sync::Arc;
33

44
use bytes::Bytes;
5+
use cas_client::Client;
56
use data::configurations::TranslatorConfig;
67
use data::{FileDownloadSession, FileUploadSession, SingleFileCleaner, XetFileInfo};
8+
use merklehash::MerkleHash;
79
use ulid::Ulid;
810

911
use crate::error::{Error, Result};
@@ -17,6 +19,9 @@ pub trait XetOps: Send + Sync {
1719
async fn download_to_file(&self, xet_hash: &str, file_size: u64, dest: &Path) -> Result<()>;
1820
async fn upload_files(&self, paths: &[&Path]) -> Result<Vec<XetFileInfo>>;
1921
fn download_stream_boxed(&self, file_info: &XetFileInfo, offset: u64) -> Result<Box<dyn DownloadStreamOps>>;
22+
/// Pre-warm the reconstruction cache for a file by fetching its full plan.
23+
/// Errors are silently ignored — this is best-effort.
24+
async fn warm_reconstruction_cache(&self, xet_hash: &str);
2025
}
2126

2227
/// Append-only streaming writer trait (abstracts StreamingWriter for testing).
@@ -41,11 +46,21 @@ pub trait DownloadStreamOps: Send {
4146
pub struct XetSessions {
4247
session: Arc<FileDownloadSession>,
4348
upload_config: Option<Arc<TranslatorConfig>>,
49+
/// Kept separately from `session` so we can call `get_reconstruction` for cache warming.
50+
cas_client: Arc<dyn Client>,
4451
}
4552

4653
impl XetSessions {
47-
pub fn new(session: Arc<FileDownloadSession>, upload_config: Option<Arc<TranslatorConfig>>) -> Arc<Self> {
48-
Arc::new(Self { session, upload_config })
54+
pub fn new(
55+
session: Arc<FileDownloadSession>,
56+
upload_config: Option<Arc<TranslatorConfig>>,
57+
cas_client: Arc<dyn Client>,
58+
) -> Arc<Self> {
59+
Arc::new(Self {
60+
session,
61+
upload_config,
62+
cas_client,
63+
})
4964
}
5065

5166
/// Start a streaming download from a byte offset (sync, returns an iterator-like stream).
@@ -114,6 +129,12 @@ impl XetOps for XetSessions {
114129
let stream = self.download_stream(file_info, offset)?;
115130
Ok(Box::new(DownloadStreamWrapper(stream)))
116131
}
132+
133+
async fn warm_reconstruction_cache(&self, xet_hash: &str) {
134+
if let Ok(hash) = MerkleHash::from_hex(xet_hash) {
135+
let _ = self.cas_client.get_reconstruction(&hash, None).await;
136+
}
137+
}
117138
}
118139

119140
// ── DownloadStreamWrapper ─────────────────────────────────────────────

0 commit comments

Comments
 (0)