Function master_setup

cake-core/src/cake/sharding/mod.rs:162–506

Run the full zero-config master setup. Discovers workers via mDNS, computes layer assignments based on VRAM, connects to each worker with mutual authentication, pushes model data as needed, and returns a `Topology` ready for normal inference.

(
    cluster_key: &str,
    model_path: &Path,
    discovery_timeout: Duration,
    min_workers: usize,
)

/// connects to each worker with mutual authentication, pushes model data
/// as needed, and returns a `Topology` ready for normal inference.
pub async fn master_setup(
    cluster_key: &str,
    model_path: &Path,
    discovery_timeout: Duration,
    min_workers: usize,
) -> Result<Topology> {
    // Read config.json and compute a fingerprint for cache keying
    let config_path = model_path.join("config.json");
    let config_data = std::fs::read_to_string(&config_path)
        .map_err(|e| anyhow!("failed to read {}: {}", config_path.display(), e))?;
    let model_hash = {
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(config_data.as_bytes());
        let result = hasher.finalize();
        hex::encode(&result[..4])
    };
    let config_json: serde_json::Value = serde_json::from_str(&config_data)?;
    let num_layers = config_json
        .get("num_hidden_layers")
        .and_then(|v| v.as_u64())
        .or_else(|| {
            // Some models (e.g. Qwen3.5) nest config under text_config
            config_json
                .get("text_config")
                .and_then(|tc| tc.get("num_hidden_layers"))
                .and_then(|v| v.as_u64())
        })
        .ok_or_else(|| anyhow!("num_hidden_layers not found in config.json"))? as usize;

    // Derive layer naming prefix from architecture (needed early for layer size estimation)
    let layer_prefix = default::layer_prefix_for_config(&config_json);

    log::info!(
        "model has {} transformer layers (prefix: {})",
        num_layers,
        &layer_prefix,
    );

    // Detect master GPU and free VRAM concurrently with the discovery window
    // (nvidia-smi can take ~1-2s; hide that cost inside the discovery timeout).
    let master_gpus = discovery::detect_gpus();
    let master_tflops: f64 = master_gpus.iter().map(|g| g.tflops as f64).sum();
    let free_gpu_fut = tokio::task::spawn_blocking(detect_free_gpu_memory);

    // Discover workers
    let workers = discovery::discover_workers(cluster_key, discovery_timeout, min_workers).await?;
    if workers.is_empty() {
        log::warn!("no workers discovered — all layers will be loaded locally");
        return Ok(Topology::new());
    }

    // nvidia-smi result is now ready (ran during the discovery window)
    let master_free_from_smi = free_gpu_fut.await.unwrap_or(0);

    // Estimate per-layer size for VRAM-aware capping.
    // Uses weight_map tensor-count fractions to exclude non-layer weights
    // (visual encoder, MTP heads, embeddings, lm_head, FP8 scale_inv, etc.).