Skip to content

Commit 3978c31

Browse files
javieratomicclaude
andcommitted
Link received payload objects into the package cache instead of copying
A rank export receives each cache payload object as a complete companion wrapper in a synced temporary file, then `store_slice` hashed its payload, copied it into a new wrapper in cache staging, synced that copy and verified it again after installation. For a DeepSeek V4.1 Flash TP3 export on three DGX Spark GB10 hosts that is about 200 GB read twice more and written and synced once more on rank zero while the streams wait. `Root::adopt_companion` links the received file itself into staging when its framing is byte-for-byte the wrapper `store_slice` would write (same length and the same leading bytes, from `companion_prefix`) and its payload hashes to the declared identity. Publication is unchanged: the object is made immutable, linked under its identity and verified once more. Any other framing, or a file on another file system, returns `None` and the receiver falls back to `store_slice` as before. An already cached identity returns the verified cached object, as `store_slice` does. Validation: - helios `cargo test --release --features server,runtime-package --lib`: 1181 passed (1179 on main; the only failures are the 7 server::protocol tests that also fail on main); new tests cover the canonical prefix for several lengths, adoption by link (same inode), a non-canonical framing declined, a mismatched identity rejected, and a cached identity reused. - heliopause `cargo test --release`: 166 passed, 0 failed. - Full TP3 exports on three GB10 hosts, same model, recipe flags and cold forward cache: "Receive rank export payloads" took 280.6 s with this change (default streams), against 457.8 s (default streams, parent without concurrent shared insertion) and 477.2 s (parent, 16 transfer streams) without it; "Export native packages" took 1473.2 s, against 1677.4 s for the parent with 16 streams. The exported package served with `flare run`: 262k-token prompt 16/16 retrieval cold (79.9 s) and warm (1.8 s), identical output, decode in line with earlier packages. Limitations: one export measured per configuration; the payload is still hashed by a separate read after receipt instead of while it streams. Signed-off-by: Javier López <59578422+javieratomic@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Design-Signed-Off-By: Claude Opus 5.5; design-review=1c01
1 parent a161b17 commit 3978c31

3 files changed

Lines changed: 134 additions & 11 deletions

File tree

‎heliopause/src/export/transfer.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -735,9 +735,10 @@ fn verify_cached_payload(
735735

736736
/// Adopt a received cache payload object under its signed payload identity.
737737
///
738-
/// The staged file is a complete companion container; its payload bytes are
739-
/// verified against the declared identity by the cache insertion path, so a
740-
/// corrupt or mislabelled object never enters the cache.
738+
/// The staged file is a complete, synced companion container; its payload bytes
739+
/// are verified against the declared identity by the cache insertion path, so a
740+
/// corrupt or mislabelled object never enters the cache. A canonical wrapper is
741+
/// linked into the cache as is; any other is copied through its slice.
741742
fn install_payload_object(
742743
cache: &Path,
743744
id: helios_core::package::PayloadId,
@@ -752,7 +753,10 @@ fn install_payload_object(
752753
"received payload object length differs from its declared group",
753754
));
754755
}
755-
let object = root.store_slice(Some(id), &slice)?;
756+
let object = match root.adopt_companion(id, partial)? {
757+
Some(object) => object,
758+
None => root.store_slice(Some(id), &slice)?,
759+
};
756760
if object.length() != length {
757761
return Err(invalid("received payload object has the wrong length"));
758762
}

‎helios/src/package/cache.rs‎

Lines changed: 102 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,10 @@ thread_local! {
3939
static PAYLOAD_OPENS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
4040
}
4141
use crate::Result;
42+
use crate::os::FileReadAt;
4243
use crate::package::container::{
43-
COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, open_relative,
44-
validate_relative_path, write_companion,
44+
COMPANION_TENSOR, Container, PAYLOAD_FORMAT, Slice, Source, Writer, companion_prefix,
45+
open_relative, validate_relative_path, write_companion,
4546
};
4647

4748
const MAX_NAME: usize = 128;
@@ -1498,11 +1499,7 @@ impl Root {
14981499
)
14991500
.into());
15001501
}
1501-
if let Some(object) = self.payload(id)? {
1502-
if object.length() != slice.len() {
1503-
return Err(format!("cached object {id} has a different length").into());
1504-
}
1505-
verify_slice(object.slice(), id, None)?;
1502+
if let Some(object) = self.cached_slice(id, slice.len())? {
15061503
return Ok(object);
15071504
}
15081505
let transaction = Transaction::begin(self)?;
@@ -1524,6 +1521,59 @@ impl Root {
15241521
result
15251522
}
15261523

1524+
/// Adopt a complete, durable companion wrapper file as a payload object
1525+
/// without copying it.
1526+
///
1527+
/// The file is linked into staging only when its framing is byte-for-byte
1528+
/// the wrapper `store_slice` would write and its payload hashes to
1529+
/// `declared`; publication verifies the installed object again. `None`
1530+
/// means the file is not that canonical wrapper or lives on another file
1531+
/// system, and the caller stores its slice instead. The caller must have
1532+
/// synced the file and must not modify it afterwards: it becomes the object.
1533+
pub fn adopt_companion(&self, declared: PayloadId, path: &Path) -> Result<Option<PayloadObject>> {
1534+
let slice = Container::open(path)?.tensor(COMPANION_TENSOR)?;
1535+
let prefix = companion_prefix(slice.len())?;
1536+
let file = File::open(path).map_err(failure)?;
1537+
if file.metadata().map_err(failure)?.len() != prefix.len() as u64 + slice.len() {
1538+
return Ok(None);
1539+
}
1540+
let mut framing = vec![0; prefix.len()];
1541+
file.read_exact_at(&mut framing, 0).map_err(failure)?;
1542+
if framing != prefix {
1543+
return Ok(None);
1544+
}
1545+
let id = slice.digest()?;
1546+
if id != declared {
1547+
return Err(format!(
1548+
"payload digest {id} does not match its declared identity {declared}"
1549+
)
1550+
.into());
1551+
}
1552+
if let Some(object) = self.cached_slice(id, slice.len())? {
1553+
return Ok(Some(object));
1554+
}
1555+
let transaction = Transaction::begin(self)?;
1556+
let staged = transaction.staging_file("wrapped")?;
1557+
match fs::hard_link(path, &staged) {
1558+
Ok(()) => {}
1559+
Err(error) if error.kind() == io::ErrorKind::CrossesDevices => return Ok(None),
1560+
Err(error) => return Err(failure(error)),
1561+
}
1562+
self.publish_payload(&transaction, &staged, id, slice.len()).map(Some)
1563+
}
1564+
1565+
/// A cached object for `id`, verified, when one is already installed.
1566+
fn cached_slice(&self, id: PayloadId, length: u64) -> Result<Option<PayloadObject>> {
1567+
let Some(object) = self.payload(id)? else {
1568+
return Ok(None);
1569+
};
1570+
if object.length() != length {
1571+
return Err(format!("cached object {id} has a different length").into());
1572+
}
1573+
verify_slice(object.slice(), id, None)?;
1574+
Ok(Some(object))
1575+
}
1576+
15271577
/// Insert an optional retention source while distinguishing a corrupt
15281578
/// source from publication failure. A missing/corrupt source yields
15291579
/// `Ok(None)`; object-store or staging I/O errors are returned.
@@ -4756,6 +4806,51 @@ mod tests {
47564806
assert_eq!(root.collect().unwrap().objects, 1);
47574807
}
47584808

4809+
#[test]
4810+
#[cfg(unix)]
4811+
fn adopt_companion_links_canonical_wrappers_and_declines_others() {
4812+
use std::os::unix::fs::MetadataExt;
4813+
let temporary = Temporary::new();
4814+
let root = temporary.root();
4815+
let bytes: Vec<u8> = (0..2 * 1024 * 1024 + 77).map(|i| (i % 241) as u8).collect();
4816+
let raw = temporary.0.join("slice.bin");
4817+
fs::write(&raw, &bytes).unwrap();
4818+
let slice = Slice::new(Arc::new(File::open(&raw).unwrap()), 0, bytes.len() as u64).unwrap();
4819+
let id = PayloadId::of(&bytes);
4820+
4821+
// A framing that is not the canonical wrapper is left to store_slice.
4822+
let framed = temporary.0.join("framed.safetensors");
4823+
let mut writer = Writer::new(File::create(&framed).unwrap(), PAYLOAD_FORMAT).unwrap();
4824+
writer.add(COMPANION_TENSOR, bytes.len() as u64, Source::Slice(slice.clone())).unwrap();
4825+
writer.add("extra.tail", 3, Source::Bytes(vec![1, 2, 3])).unwrap();
4826+
writer.finish().unwrap();
4827+
assert!(root.adopt_companion(id, &framed).unwrap().is_none());
4828+
4829+
let received = temporary.0.join("received.safetensors");
4830+
write_companion(&mut File::create(&received).unwrap(), &slice).unwrap();
4831+
assert!(root.adopt_companion(PayloadId::of(b"other"), &received).is_err());
4832+
assert!(root.payload(id).unwrap().is_none());
4833+
4834+
// The canonical wrapper becomes the object itself: no copy is written.
4835+
let object = root.adopt_companion(id, &received).unwrap().unwrap();
4836+
assert_eq!(object.id(), id);
4837+
assert_eq!(object.length(), bytes.len() as u64);
4838+
assert_eq!(
4839+
fs::metadata(object.path()).unwrap().ino(),
4840+
fs::metadata(&received).unwrap().ino()
4841+
);
4842+
fs::remove_file(&received).unwrap();
4843+
let mut copy = vec![0; bytes.len()];
4844+
object.slice().read_at(0, &mut copy).unwrap();
4845+
assert_eq!(copy, bytes);
4846+
4847+
// A second arrival of the same payload returns the cached object.
4848+
let again_path = temporary.0.join("again.safetensors");
4849+
write_companion(&mut File::create(&again_path).unwrap(), &slice).unwrap();
4850+
let again = root.adopt_companion(id, &again_path).unwrap().unwrap();
4851+
assert_eq!(again.path(), object.path());
4852+
}
4853+
47594854
#[test]
47604855
#[cfg(unix)]
47614856
fn legacy_import_streams_large_companions_and_publishes_once() {

‎helios/src/package/container.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,15 @@ pub fn write_companion(output: &mut impl Write, payload: &Slice) -> Result<()> {
537537
writer.finish()
538538
}
539539

540+
/// The exact bytes `write_companion` writes ahead of a `len`-byte payload.
541+
pub fn companion_prefix(len: u64) -> Result<Vec<u8>> {
542+
let layout = [(COMPANION_TENSOR.to_owned(), 0, len, Source::Zero(len))];
543+
let header = header_bytes(PAYLOAD_FORMAT, &layout)?;
544+
let mut prefix = (header.len() as u64).to_le_bytes().to_vec();
545+
prefix.extend_from_slice(&header);
546+
Ok(prefix)
547+
}
548+
540549
/// Maps a companion path to its packed tensor name.
541550
pub fn packed_tensor_name(relative: &str) -> Result<String> {
542551
validate_relative_path(relative)?;
@@ -1207,6 +1216,21 @@ mod tests {
12071216
);
12081217
}
12091218

1219+
#[test]
1220+
fn companion_prefix_is_the_wrapper_ahead_of_its_payload() {
1221+
let scratch = Scratch::new("prefix");
1222+
for length in [1usize, 1234, 4096, 4097] {
1223+
let payload: Vec<u8> = (0..length as u32).map(|value| (value * 13) as u8).collect();
1224+
let (_, container) = container_with(&scratch, "payload.safetensors", &payload);
1225+
let slice = container.tensor(COMPANION_TENSOR).unwrap();
1226+
let mut wrapper = Vec::new();
1227+
write_companion(&mut wrapper, &slice).unwrap();
1228+
let prefix = companion_prefix(length as u64).unwrap();
1229+
assert_eq!(wrapper.len(), prefix.len() + length);
1230+
assert_eq!(&wrapper[..prefix.len()], prefix.as_slice());
1231+
}
1232+
}
1233+
12101234
#[test]
12111235
fn duplicate_header_keys_are_rejected() {
12121236
let scratch = Scratch::new("duplicate");

0 commit comments

Comments
 (0)