Function receive_model_data

cake-core/src/cake/sharding/mod.rs:940–1099

Receive model data from master and write to the cache directory.

(
    stream: &mut TcpStream,
    cache_dir: &Path,
    layers: &[String],
    on_progress: Option<&SetupProgressFn>,
)

/// Receive model data from master and write to the cache directory.
async fn receive_model_data(
    stream: &mut TcpStream,
    cache_dir: &Path,
    layers: &[String],
    on_progress: Option<&SetupProgressFn>,
) -> Result<()> {
    let overall_start = Instant::now();
    let mut overall_bytes: u64 = 0;
    let mut current_file: Option<(String, std::fs::File, Instant, u64)> = None;

    let layer_range = if layers.is_empty() {
        "(none)".to_string()
    } else {
        format!(
            "{} — {} ({} layers)",
            layers.first().unwrap(),
            layers.last().unwrap(),
            layers.len()
        )
    };

    log::info!("receiving model data [{}] ...", layer_range);

    let mut read_buf = Vec::with_capacity(MODEL_DATA_CHUNK_SIZE + 1024);
    loop {
        let (_, msg) = Message::from_reader_buf(stream, &mut read_buf).await?;

        match msg {
            Message::ModelDataChunk {
                filename,
                offset,
                total_size,
                compressed,
                checksum,
                data,
            } => {
                // Verify CRC32 checksum before writing
                let actual_crc = crc32fast::hash(&data);
                if actual_crc != checksum {
                    return Err(anyhow!(
                        "checksum mismatch for {} at offset {}: expected {:#x}, got {:#x}",
                        filename, offset, checksum, actual_crc
                    ));
                }

                // Decompress if compressed
                let data = if compressed {
                    zstd::decode_all(data.as_slice())
                        .map_err(|e| anyhow!("zstd decompress failed for {} at offset {}: {}", filename, offset, e))?
                } else {
                    data
                };
                // Open new file if needed
                let file = if let Some((ref name, ref mut file, _, _)) = current_file {
                    if name == &filename {
                        file
                    } else {
                        // Close previous file, log stats