prepare_weights/
prepare_weights.rs

1//! Measure page preparation plus a first host read or completed device upload.
2//! Use saved fresh processes in alternating order; this does not evict caches.
3use hrx::artifacts::safetensors::FileView;
4use std::time::Instant;
5
6fn main() -> Result<(), Box<dyn std::error::Error>> {
7    let mut args: Vec<_> = std::env::args().skip(1).collect();
8    let upload = args.last().is_some_and(|arg| arg == "--upload");
9    if upload {
10        args.pop();
11    }
12    if args.len() != 3 || !matches!(args[2].as_str(), "none" | "advice" | "populate") {
13        return Err("usage: prepare_weights FILE TENSOR none|advice|populate [--upload]".into());
14    }
15    // The command requires an immutable checkpoint for the mapping's lifetime.
16    let file = unsafe { FileView::map(&args[0]) }?;
17    let tensor = file.get(&args[1])?;
18    let mut stream = if upload {
19        Some(hrx::Stream::open()?)
20    } else {
21        None
22    };
23    let buffer = stream
24        .as_ref()
25        .map(|stream| stream.allocate(tensor.bytes.len()))
26        .transpose()?;
27    let start = Instant::now();
28    let prepared = match args[2].as_str() {
29        "populate" => format!("{:?}", file.prepare_bytes(tensor.bytes)?),
30        "advice" => {
31            file.will_need(tensor);
32            "Advised".into()
33        }
34        _ => "None".into(),
35    };
36    let preparation_ms = start.elapsed().as_secs_f64() * 1000.;
37    let checksum: u64;
38    let total_ms;
39    if let (Some(stream), Some(buffer)) = (&mut stream, &buffer) {
40        stream.upload_blocking(buffer.binding(), tensor.bytes)?;
41        total_ms = start.elapsed().as_secs_f64() * 1000.;
42        // Validation is outside timing, after the upload fence has retired.
43        let mut actual = vec![0; tensor.bytes.len()];
44        stream.read_blocking(buffer.binding(), &mut actual)?;
45        if actual != tensor.bytes {
46            return Err("upload bytes differ from checkpoint".into());
47        }
48        checksum = actual.iter().map(|&v| u64::from(v)).sum();
49    } else {
50        checksum = tensor.bytes.iter().map(|&v| u64::from(v)).sum();
51        total_ms = start.elapsed().as_secs_f64() * 1000.;
52    }
53    println!(
54        "{}",
55        serde_json::json!({"mode":args[2], "preparation":prepared,
56        "bytes":tensor.bytes.len(), "preparation_ms":preparation_ms,
57        "total_ms":total_ms, "checksum":checksum, "upload":upload,
58        "allocation_policy":"preallocated destination; fresh process",
59        "timed_scope":if upload { "page preparation plus completed upload" } else { "page preparation plus complete host read" }, "cache_state":"externally controlled"})
60    );
61    Ok(())
62}
prepare_weights/prepare_weights.rs

prepare_weights/
prepare_weights.rs
