From 8c97421a723be58c2574ce2d096c7711320480b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Javier=20L=C3=B3pez?= <59578422+javieratomic@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:44:42 +0100 Subject: [PATCH] Link received payload objects into the core cache instead of copying MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rank export receives each cache payload object as a complete companion wrapper in a synced temporary file in the core cache. `store_slice` then hashes the received payload, copies it into a new wrapper in cache staging, syncs the copy, and hashes the installed copy again: three full reads and one write and sync per object, on the receiving stream's thread. When the received file already is the wrapper `store_slice` would write, the copy, its write and its sync repeat work whose result is already on disk. In a DeepSeek V4.1 Flash TP3 export on three DGX Spark GB10 hosts with an empty core cache, rank zero receives 203.5 GB from the other two ranks this way. Before this change the receiver always called `store_slice`. Now it first calls the new `Root::adopt_companion`, which links the received file itself into cache staging when it is byte-for-byte that wrapper; both hashes remain. The link step is shared with directory publication, which already hard-linked source files that are canonical wrappers from its own code in `export.rs`; that code moves into `container::link_companion`. It checks the length and leading bytes against `companion_prefix`, links the opened file through its descriptor so that a replaced path cannot be linked instead, checks that the link names the same inode, and hashes the linked file against the declared identity; a mismatched identity is an error. For a file that is not the canonical wrapper, or a link the filesystem refuses, `link_companion` returns false, `adopt_companion` returns `None`, and the receiver stores the slice as before. Adoption checks and links the descriptor the receiver opened, the same one `store_slice` copies from. An already cached identity returns the verified cached object, as `store_slice` does. Publication is unchanged: the object is made immutable, linked under its identity and verified once more, and the receiver closes its writable descriptor before the file is published. Core format, cache layout and object identities are unchanged, and directory publication keeps its behavior, error kinds and messages. Adoption needs Linux (`linkat` through `/proc/self/fd`); on other platforms received objects are copied as before. Expected result: with an empty core cache, "Receive rank export payloads" no longer copies, writes and syncs each received object. On an earlier tree, with an earlier version of this adoption, that phase took 280.6 s instead of 457.8 s; this commit itself was not measured on its current base. Evidence: helios `cargo test --release --features server,runtime-core --lib`: 1335 passed, 0 failed, 26 ignored, with FLARE_TEST_TEMPLATE_DIR prepared by tests/frontend/prepare_templates.py. heliopause `cargo test --release`: 190 passed, 0 failed, 3 ignored. A helios test covers adoption by link (same inode); a wrapper with an extra tensor and a wrapper of the right length with different header bytes, both declined; a mismatched identity, which is rejected; a second arrival of a cached identity, which returns the cached object; and that cache staging is empty after the declined, rejected and cached cases. A heliopause test checks that the receiver's install step adopts a canonical wrapper (same inode) and copies another framing into a new file. The existing directory publication tests cover the shared link step. The 203.5 GB is the byte count reported by "Receive rank export payloads" in an export on `d2b5c5f3` without this change, with an empty core cache; there that phase took 462.3 s. The measurement above used full TP3 exports with the same model and export command, an empty core cache (then named the package cache) and an empty forward compilation cache, built from `ca1f960e` plus an earlier version of #124 that published cache payloads concurrently and without the directory publication changes merged as `e13ecb74`; its adoption step linked the received file by path instead of through `link_companion`. There, with the default 8 transfer streams, "Receive rank export payloads" took 280.6 s with adoption and 457.8 s without it (477.2 s without it and `--transfer-streams 16`), and the whole export phase took 1473.2 s against 1677.4 s; the package exported with adoption served with `flare run` and answered a 262k-token prompt that asks for 16 values placed through the context with all 16 correct, with identical output with an empty and with a warm prefix cache. One export per configuration was measured. The payload is still read twice after it arrives (identity check and post-installation verification) instead of being hashed while it streams. Design-Signed-Off-By: Claude Opus 5.5; design-review=cfcf Co-Authored-By: Claude Opus 5.5 Signed-off-by: Javier López <59578422+javieratomic@users.noreply.github.com> --- heliopause/src/export/transfer.rs | 69 ++++++++++++++++- helios/src/core/cache.rs | 120 ++++++++++++++++++++++++++++-- helios/src/core/container.rs | 99 ++++++++++++++++++++++++ helios/src/core/export.rs | 104 +------------------------- 4 files changed, 279 insertions(+), 113 deletions(-) diff --git a/heliopause/src/export/transfer.rs b/heliopause/src/export/transfer.rs index 0096778f..19eeb59f 100644 --- a/heliopause/src/export/transfer.rs +++ b/heliopause/src/export/transfer.rs @@ -741,9 +741,11 @@ fn verify_cached_payload( /// Adopt a received cache payload object under its signed payload identity. /// -/// The staged file is a complete companion container; its payload bytes are -/// verified against the declared identity by the cache insertion path, so a -/// corrupt or mislabelled object never enters the cache. +/// The staged file is a complete, synced companion container; its payload bytes +/// are verified against the declared identity by the cache insertion path, so a +/// corrupt or mislabelled object never enters the cache. A canonical wrapper is +/// linked into the cache as is; otherwise, or when the filesystem refuses the +/// link, the slice is copied. fn install_payload_object( cache: &Path, id: helios_core::core::PayloadId, @@ -758,7 +760,10 @@ fn install_payload_object( "received payload object length differs from its declared group", )); } - let object = root.store_slice(Some(id), &slice)?; + let object = match root.adopt_companion(id, &slice, &partial.display().to_string())? { + Some(object) => object, + None => root.store_slice(Some(id), &slice)?, + }; if object.length() != length { return Err(invalid("received payload object has the wrong length")); } @@ -1038,6 +1043,9 @@ fn receive_rank( fs::Permissions::from_mode(0o444) })?; } + // The synced file may become a cache object itself: close the + // writable descriptor before it is published. + drop(file); if info.payload.is_some() { let id = helios_core::core::PayloadId::parse( info.payload.as_deref().unwrap_or_default(), @@ -1444,6 +1452,59 @@ mod tests { (output, id) } + #[test] + #[cfg(target_os = "linux")] + fn received_canonical_wrapper_is_adopted_and_others_are_copied() { + use std::os::unix::fs::MetadataExt; + let directory = Directory::new(); + let payload: Vec = (0..4096u32).map(|i| (i % 97) as u8).collect(); + let (_, id) = cache_backed_output(&directory, &payload); + let sent = directory.0.join("sender-cache").join("blobs").join(id.file_name()); + let cache = directory.0.join("cache"); + helios_core::core::cache::Root::open(&cache).unwrap(); + + // The receiver writes the sender's object file under the cache's blobs directory. + let partial = cache.join("blobs").join(".rank-blob-canonical"); + fs::copy(&sent, &partial).unwrap(); + let object = install_payload_object(&cache, id, payload.len() as u64, &partial).unwrap(); + assert_eq!( + fs::metadata(object.path()).unwrap().ino(), + fs::metadata(&partial).unwrap().ino(), + "a canonical received wrapper was copied instead of adopted" + ); + drop(object); + + // Another framing of the same payload is stored through its slice in a new file. + let other = directory.0.join("other-cache"); + helios_core::core::cache::Root::open(&other).unwrap(); + let framed = other.join("blobs").join(".rank-blob-framed"); + let slice = helios_core::core::container::Container::open(&sent) + .unwrap() + .tensor(helios_core::core::container::COMPANION_TENSOR) + .unwrap(); + let mut writer = helios_core::core::container::Writer::new( + fs::File::create(&framed).unwrap(), + helios_core::core::container::PAYLOAD_FORMAT, + ) + .unwrap(); + writer + .add( + helios_core::core::container::COMPANION_TENSOR, + slice.len(), + helios_core::core::container::Source::Slice(slice), + ) + .unwrap(); + writer + .add("extra.tail", 3, helios_core::core::container::Source::Bytes(vec![1, 2, 3])) + .unwrap(); + writer.finish().unwrap(); + let copied = install_payload_object(&other, id, payload.len() as u64, &framed).unwrap(); + assert_ne!( + fs::metadata(copied.path()).unwrap().ino(), + fs::metadata(&framed).unwrap().ino() + ); + } + #[test] fn repeated_identical_payload_groups_share_one_object_and_round_trip() { let directory = Directory::new(); diff --git a/helios/src/core/cache.rs b/helios/src/core/cache.rs index 62d0604a..f0070a44 100644 --- a/helios/src/core/cache.rs +++ b/helios/src/core/cache.rs @@ -45,8 +45,8 @@ thread_local! { } use crate::Result; use crate::core::container::{ - COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, open_relative, - validate_relative_path, write_companion, + COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, link_companion, + open_relative, validate_relative_path, write_companion, }; const MAX_NAME: usize = 128; @@ -1503,11 +1503,7 @@ impl Root { ) .into()); } - if let Some(object) = self.payload(id)? { - if object.length() != slice.len() { - return Err(format!("cached object {id} has a different length").into()); - } - verify_slice(object.slice(), id, None)?; + if let Some(object) = self.cached_slice(id, slice.len())? { return Ok(object); } let transaction = Transaction::begin(self)?; @@ -1529,6 +1525,48 @@ impl Root { result } + /// Adopt the complete, durable companion wrapper file that holds `slice` as + /// a payload object without copying it. + /// + /// The opened file is linked into staging only when it is byte-for-byte + /// the wrapper `store_slice` would write; the linked file must then hash to + /// `declared`, and publication verifies the installed object again. `None` + /// means the file is not that canonical wrapper or the filesystem refuses + /// the link, and the caller stores its slice instead. A mismatched identity + /// is an error. An already cached identity returns the cached object. The + /// caller must have synced the file under `slice` and must not modify it + /// afterwards: it becomes the object. `label` names the payload in errors. + pub fn adopt_companion( + &self, + declared: PayloadId, + slice: &Slice, + label: &str, + ) -> Result> { + let transaction = Transaction::begin(self)?; + let staged = transaction.staging_file("wrapped")?; + if !link_companion(slice, declared, &staged, label)? { + return Ok(None); + } + let cached = self.cached_slice(declared, slice.len()); + if !matches!(cached, Ok(None)) { + let _ = fs::remove_file(&staged); + return cached; + } + self.publish_payload(&transaction, &staged, declared, slice.len()).map(Some) + } + + /// A cached object for `id`, verified, when one is already installed. + fn cached_slice(&self, id: PayloadId, length: u64) -> Result> { + let Some(object) = self.payload(id)? else { + return Ok(None); + }; + if object.length() != length { + return Err(format!("cached object {id} has a different length").into()); + } + verify_slice(object.slice(), id, None)?; + Ok(Some(object)) + } + /// Insert an optional retention source while distinguishing a corrupt /// source from publication failure. A missing/corrupt source yields /// `Ok(None)`; object-store or staging I/O errors are returned. @@ -4855,6 +4893,74 @@ mod tests { assert_eq!(root.collect().unwrap().objects, 1); } + #[test] + #[cfg(target_os = "linux")] + fn adopt_companion_links_canonical_wrappers_and_declines_others() { + use std::os::unix::fs::MetadataExt; + let temporary = Temporary::new(); + let root = temporary.root(); + let bytes: Vec = (0..2 * 1024 * 1024 + 77).map(|i| (i % 241) as u8).collect(); + let raw = temporary.0.join("slice.bin"); + fs::write(&raw, &bytes).unwrap(); + let slice = Slice::new(Arc::new(File::open(&raw).unwrap()), 0, bytes.len() as u64).unwrap(); + let id = PayloadId::of(&bytes); + let adopt = |declared: PayloadId, path: &Path| { + let slice = Container::open(path).unwrap().tensor(COMPANION_TENSOR).unwrap(); + root.adopt_companion(declared, &slice, "test") + }; + let staging_is_empty = || { + fs::read_dir(root.root.join("staging").join("wrapped")) + .unwrap() + .count() + == 0 + }; + + // A framing that is not the canonical wrapper is left to store_slice. + let framed = temporary.0.join("framed.safetensors"); + let mut writer = Writer::new(File::create(&framed).unwrap(), PAYLOAD_FORMAT).unwrap(); + writer.add(COMPANION_TENSOR, bytes.len() as u64, Source::Slice(slice.clone())).unwrap(); + writer.add("extra.tail", 3, Source::Bytes(vec![1, 2, 3])).unwrap(); + writer.finish().unwrap(); + assert!(adopt(id, &framed).unwrap().is_none()); + + // Same length as the canonical wrapper, but its header padding differs. + let padded = temporary.0.join("padded.safetensors"); + write_companion(&mut File::create(&padded).unwrap(), &slice).unwrap(); + let mut wrapper = fs::read(&padded).unwrap(); + let header_end = 8 + u64::from_le_bytes(wrapper[..8].try_into().unwrap()) as usize; + assert_eq!(wrapper[header_end - 1], b' '); + wrapper[header_end - 1] = b'\n'; + fs::write(&padded, &wrapper).unwrap(); + assert!(adopt(id, &padded).unwrap().is_none()); + assert!(staging_is_empty()); + + let received = temporary.0.join("received.safetensors"); + write_companion(&mut File::create(&received).unwrap(), &slice).unwrap(); + assert!(adopt(PayloadId::of(b"other"), &received).is_err()); + assert!(root.payload(id).unwrap().is_none()); + assert!(staging_is_empty()); + + // The canonical wrapper becomes the object itself: no copy is written. + let object = adopt(id, &received).unwrap().unwrap(); + assert_eq!(object.id(), id); + assert_eq!(object.length(), bytes.len() as u64); + assert_eq!( + fs::metadata(object.path()).unwrap().ino(), + fs::metadata(&received).unwrap().ino() + ); + fs::remove_file(&received).unwrap(); + let mut copy = vec![0; bytes.len()]; + object.slice().read_at(0, &mut copy).unwrap(); + assert_eq!(copy, bytes); + + // A second arrival of a cached identity returns the cached object. + let again_path = temporary.0.join("again.safetensors"); + write_companion(&mut File::create(&again_path).unwrap(), &slice).unwrap(); + let again = adopt(id, &again_path).unwrap().unwrap(); + assert_eq!(again.path(), object.path()); + assert!(staging_is_empty()); + } + #[test] #[cfg(unix)] fn legacy_import_streams_large_companions_and_publishes_once() { diff --git a/helios/src/core/container.rs b/helios/src/core/container.rs index c13b1e40..796e1630 100644 --- a/helios/src/core/container.rs +++ b/helios/src/core/container.rs @@ -491,6 +491,105 @@ pub(crate) fn companion_prefix(nbytes: u64) -> Result> { Ok(prefix) } +/// Hard-link the file under `payload` to `to` when that file is exactly the +/// deterministic companion wrapper of the payload, then verify the linked file +/// against `id`. Returns false, creating no file, when the file is not that +/// wrapper or the filesystem refuses the link; the caller then copies. An +/// existing `to` is an error. If the linked file fails verification, the link +/// is removed and the error returned; `label` names the payload in errors. +#[cfg(target_os = "linux")] +pub(crate) fn link_companion(payload: &Slice, id: PayloadId, to: &Path, label: &str) -> Result { + use std::os::unix::ffi::OsStrExt; + use std::os::unix::fs::MetadataExt; + use std::os::unix::io::AsRawFd; + + // I/O failures are state errors; a linked file that is not the payload's wrapper + // or identity, and an unrepresentable path, are manifest errors. + let failed = |context: &str, error: io::Error| -> crate::Error { + format!("{context}: {error}").into() + }; + + let nbytes = payload.len(); + let prefix = companion_prefix(nbytes)?; + let source = payload.file(); + let metadata = source + .metadata() + .map_err(|e| failed("inspect payload source", e))?; + let wrapper_len = prefix.len() as u64 + nbytes; + if payload.offset() != prefix.len() as u64 + || !metadata.is_file() + || metadata.len() != wrapper_len + { + return Ok(false); + } + let mut header = vec![0u8; prefix.len()]; + // std `FileExt` keeps the short-read message that directory publication reports. + std::os::unix::fs::FileExt::read_exact_at(&**source, &mut header, 0) + .map_err(|e| failed("read payload wrapper", e))?; + if header != prefix { + return Ok(false); + } + if let Some(parent) = to.parent() { + std::fs::create_dir_all(parent).map_err(|e| failed("create payload directory", e))?; + } + // The descriptor link names the pinned file itself, not a path that could + // have been replaced since the source was opened. + let from = std::ffi::CString::new(format!("/proc/self/fd/{}", source.as_raw_fd())) + .map_err(|_| invalid("payload source descriptor path"))?; + let destination = std::ffi::CString::new(to.as_os_str().as_bytes()) + .map_err(|_| invalid(format!("payload path {label:?} contains NUL")))?; + // SAFETY: Both arguments are NUL-terminated paths that outlive the call. + let status = unsafe { + libc::linkat( + libc::AT_FDCWD, + from.as_ptr(), + libc::AT_FDCWD, + destination.as_ptr(), + libc::AT_SYMLINK_FOLLOW, + ) + }; + if status != 0 { + let error = io::Error::last_os_error(); + // An existing destination is a publication conflict, as for a copy. + if error.kind() == io::ErrorKind::AlreadyExists { + return Err(failed("link companion", error)); + } + return Ok(false); + } + let result = (|| -> Result<()> { + let linked = File::open(to).map_err(|e| failed("open linked companion", e))?; + let linked_metadata = linked + .metadata() + .map_err(|e| failed("inspect linked companion", e))?; + if (linked_metadata.dev(), linked_metadata.ino()) != (metadata.dev(), metadata.ino()) { + return Err(invalid("linked companion is not the payload source file")); + } + // Verify the file published under `to`, not the pre-link source reads. + let mut header = vec![0u8; prefix.len()]; + std::os::unix::fs::FileExt::read_exact_at(&linked, &mut header, 0) + .map_err(|e| failed("read linked companion", e))?; + if header != prefix || linked_metadata.len() != wrapper_len { + return Err(invalid("linked companion is not the payload wrapper")); + } + let slice = Slice::new(Arc::new(linked), prefix.len() as u64, nbytes)?; + if slice.digest()? != id { + return Err(invalid(format!( + "payload {label} bytes do not match the assigned identity" + ))); + } + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(to); + } + result.map(|()| true) +} + +#[cfg(not(target_os = "linux"))] +pub(crate) fn link_companion(_payload: &Slice, _id: PayloadId, _to: &Path, _label: &str) -> Result { + Ok(false) +} + fn write_source( output: &mut W, name: &str, diff --git a/helios/src/core/export.rs b/helios/src/core/export.rs index 5f1d0d31..d9ddc76e 100644 --- a/helios/src/core/export.rs +++ b/helios/src/core/export.rs @@ -1456,114 +1456,14 @@ fn install_payload_bytes( /// Install one payload into a directory core, verifying the published bytes. fn install_directory_payload(root: &Path, payload: &PublishPayload) -> Result<(), crate::Error> { - if link_verified_companion(root, &payload.relative, payload)? { + let path = root.join(&payload.relative); + if container::link_companion(&payload.source, payload.id, &path, &payload.relative)? { return Ok(()); } verify_payload_source(payload)?; write_verified_companion(root, &payload.relative, payload) } -/// Hard-link a source file that is already this payload's deterministic -/// companion wrapper, then hash the linked file. Returns false when the source -/// is not that wrapper or the filesystem refuses the link, and the caller copies. -#[cfg(target_os = "linux")] -fn link_verified_companion( - directory: &Path, - relative: &str, - payload: &PublishPayload, -) -> Result { - use std::os::unix::ffi::OsStrExt; - use std::os::unix::fs::{FileExt, MetadataExt}; - use std::os::unix::io::AsRawFd; - - let nbytes = payload.nbytes(); - let prefix = container::companion_prefix(nbytes)?; - let source = payload.source.file(); - let metadata = source - .metadata() - .map_err(|e| io("inspect payload source", e))?; - let wrapper_len = prefix.len() as u64 + nbytes; - if payload.source.offset() != prefix.len() as u64 - || !metadata.is_file() - || metadata.len() != wrapper_len - { - return Ok(false); - } - let mut header = vec![0u8; prefix.len()]; - source - .read_exact_at(&mut header, 0) - .map_err(|e| io("read payload wrapper", e))?; - if header != prefix { - return Ok(false); - } - let path = directory.join(relative); - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|e| io("create payload directory", e))?; - } - // The descriptor link names the pinned file itself, not a path that could - // have been replaced since the source was opened. - let from = std::ffi::CString::new(format!("/proc/self/fd/{}", source.as_raw_fd())) - .map_err(|_| invalid("payload source descriptor path"))?; - let to = std::ffi::CString::new(path.as_os_str().as_bytes()) - .map_err(|_| invalid(format!("payload path {relative:?} contains NUL")))?; - // SAFETY: Both arguments are NUL-terminated paths that outlive the call. - let status = unsafe { - libc::linkat( - libc::AT_FDCWD, - from.as_ptr(), - libc::AT_FDCWD, - to.as_ptr(), - libc::AT_SYMLINK_FOLLOW, - ) - }; - if status != 0 { - let error = std::io::Error::last_os_error(); - // An existing destination is a publication conflict, as for a copy. - if error.kind() == std::io::ErrorKind::AlreadyExists { - return Err(io("link companion", error)); - } - return Ok(false); - } - let result = (|| -> Result<(), crate::Error> { - let linked = File::open(&path).map_err(|e| io("open linked companion", e))?; - let linked_metadata = linked - .metadata() - .map_err(|e| io("inspect linked companion", e))?; - if (linked_metadata.dev(), linked_metadata.ino()) != (metadata.dev(), metadata.ino()) { - return Err(invalid("linked companion is not the payload source file")); - } - // Verify the file published under this path, not the pre-link source reads. - let mut header = vec![0u8; prefix.len()]; - linked - .read_exact_at(&mut header, 0) - .map_err(|e| io("read linked companion", e))?; - if header != prefix || linked_metadata.len() != wrapper_len { - return Err(invalid("linked companion is not the payload wrapper")); - } - let slice = Slice::new(std::sync::Arc::new(linked), prefix.len() as u64, nbytes)?; - if slice.digest()? != payload.id { - return Err(invalid(format!( - "payload {} bytes do not match the assigned identity", - payload.relative - ))); - } - Ok(()) - })(); - if result.is_err() { - let _ = std::fs::remove_file(&path); - } - result.map(|()| true) -} - -#[cfg(not(target_os = "linux"))] -fn link_verified_companion( - _directory: &Path, - _relative: &str, - _payload: &PublishPayload, -) -> Result { - Ok(false) -} - /// Stream one payload into its companion wrapper, then verify the bytes that /// were actually written against the signed identity. No raw temporary copy is /// created: the source is read once and the written companion is hashed.