feat(distributed): make zero-config --pp-peer bring-up work (beacons and stage push-assignment) #2658

Depends on #2655

Problem / Background

Zero-config pipeline bring-up over LAN discovery cannot work today. mlxcel-server --pp-auto N --cluster-discovery=mdns always waits the 30 s DEFAULT_DISCOVERY_TIMEOUT (src/distributed/cluster_init.rs:52) and then fails with "cluster discovery timed out ... fewer peer mlxcel-server processes running than expected" (cluster_init.rs:596-604), because no process sends a beacon. --pp-peer cannot join without a hand-copied TOML and a hand-guessed --node-id. The decision is to make the mdns flow work, not to remove it; the removal previously planned in #2655 item c is withdrawn and its verified facts moved here.

This depends on #2655 because both change src/server/startup.rs, src/distributed/cluster_init.rs, and the cluster handshake, which #2655 bumps to protocol version 2 with a typed identity rejection. Build on the merged #2655 handshake; do not fork it.

Current Behavior

• Coordinator: run_zero_config_bring_up (src/server/startup.rs:2749-2812) resolves the secret, calls discover_peers (:2760; receiver at cluster_init.rs:499-610), then plan_cluster (:2775) and write_plan_toml (:2781), and points startup.distributed_config at the written file (:2791). Nothing is sent to any peer.
• Peer: the --pp-peer branch (startup.rs:2798-2808) only ensure!s --distributed-config ("Future work will remove this once the coordinator push-assigns stages (follow-up)") and --node-id. The id must equal the stage-{i} that plan_cluster derives from the lexicographic order of peer addresses (cluster_init.rs:200-245), which the operator has to guess.
• broadcast_beacon_loop (cluster_init.rs:619-666) is the peer-side sender. Its only caller is cluster_init_tests.rs:409; it has had no production caller since it was added in 535987f. It sends to the hard-coded limited broadcast 255.255.255.255 (:642). The DiscoveryBeacon doc (cluster_init.rs:435-442) already describes the intended flow: a peer broadcasts "until the coordinator confirms receipt". Beacons are HMAC-signed with timestamp, nonce, version, and source-IP checks (sign_beacon at discovery_beacon.rs:121, BeaconVerifier at :182, from fix(distributed): authenticate cluster transports and correlate stage replies #2650).
• --pp-peer help (src/bin/mlx_server.rs:1438-1446) promises registration "against a known coordinator address" or by broadcast, then a stage service "using the stage assignment the coordinator returns". None of it exists, and in static mode it is the coordinator that knows the peer addresses (--cluster-peers), not the reverse. src/main.rs:2597-2633 carries the same flags for mlxcel serve.
• The pipeline transport already runs the mutual HMAC handshake over a raw TcpStream: client_handshake and server_handshake (src/distributed/cluster_handshake.rs:184, :241), with the client's advertised identity bound to its source IP (bind_advertised_identity, :91-107).

Facts found while verifying that the design below accounts for:

• A remote stage computes its own layer range from its local checkpoint (resolve_in_process_stage_assignments_for_model inside resolve_remote_pipeline_topology, startup.rs:2560-2590). The TOML carries no layer range, and the partition is a pure function of the model config (src/distributed/pipeline/local_runtime.rs:109-140).
• --pp-auto N needs N peer processes, because the coordinator node has no stage (cluster_init.rs:206, :218-224). --pp-auto 2 takes two --pp-peer processes, not one.
• discover_peers binds its UDP socket to the coordinator control IP (startup.rs:2763, bound at cluster_init.rs:522). A UDP socket bound to a unicast address does not receive datagrams sent to 255.255.255.255 on Linux or macOS, and the default control IP is the --host default 127.0.0.1 (mlx_server.rs:432, startup.rs:2672-2677), so a sender alone would still not be heard. cluster_init_tests.rs:406-407 already notes that limited broadcast is unreliable and uses a direct unicast sender instead.
• With --host 0.0.0.0 the coordinator node is recorded as 0.0.0.0:19000, and a stage's PeerMatcher then accepts commands from any authenticated member on port 19000 (src/distributed/peer_identity.rs:79-91).
• The generated TOML records the coordinator's absolute secret_file path (startup.rs:2699, :2713), which need not exist on a peer host.
• A peer has no flag that names its own listen address.

Proposed Solution

A new module src/distributed/stage_assignment.rs (tests in stage_assignment_tests.rs), re-exported from src/distributed/mod.rs:99-103, owns the assignment protocol. startup.rs calls it from both branches of run_zero_config_bring_up; everything after that function is unchanged, so there is one stage startup path.

1. Assignment channel and wire format

The coordinator dials the peer's stage address and runs client_handshake advertising its control address; the peer runs server_handshake advertising its listener address. Then both sides exchange frames: u32 big-endian length followed by JSON, at most 64 KiB, every read bounded by ASSIGNMENT_IO_TIMEOUT = 10 s. The exchange has two phases, so a bring-up that fails on one peer leaves no peer committed:

1. coordinator to peer: AssignmentFrame::Offer(StageAssignmentOffer)
2. peer to coordinator: AssignmentFrame::Accepted { node_id } or AssignmentFrame::Rejected { reason: AssignmentRejectReason, detail: String }
3. only after every peer answered Accepted, coordinator to each peer: AssignmentFrame::Commit
4. peer to coordinator: AssignmentFrame::Committed { node_id }

pub const ASSIGNMENT_PROTOCOL_VERSION: u16 = 1;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelIdentity {
    pub config_sha256: String,            // hex SHA-256 of the config.json bytes
    pub weight_files: Vec<(String, u64)>, // every *.safetensors file name and byte size, sorted by name
    pub num_layers: usize,                // resolve_in_process_pipeline_num_layers
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StageAssignmentOffer {
    pub protocol: u16,                       // ASSIGNMENT_PROTOCOL_VERSION
    pub mlxcel_version: String,              // env!("CARGO_PKG_VERSION")
    pub cluster_name: String,
    pub node_id: String,                     // "stage-{i}" from plan_cluster
    pub stage_index: u32,
    pub num_stages: u32,
    pub layer_range: std::ops::Range<usize>, // coordinator's partition, re-derived and compared by the peer
    pub coordinator_control_addr: SocketAddr,
    pub model: ModelIdentity,
    pub cluster_toml: String,                // render_deterministic_toml output with [cluster] secret_file = None
}

#[non_exhaustive]
pub enum AssignmentRejectReason {
    ProtocolMismatch, VersionMismatch, ClusterNameMismatch, ModelMismatch,
    LayerRangeMismatch, AddressMismatch, InvalidClusterConfig,
}

ModelIdentity::from_model_dir(&Path) -> Result<ModelIdentity> uses sha2 (already in Cargo.toml:198) and reads no weight bytes. The layer range is not an instruction (each stage derives its own, see above); it is a second identity check, so a checkpoint whose config differs in a way that moves the partition is refused even if the hash check were weakened later.

2. Peer side (--pp-peer, both discovery modes)

await_stage_assignment(config: PeerJoinConfig) -> Result<CommittedAssignment>, where PeerJoinConfig { listen_addr, discovery, discovery_port, discovery_target: IpAddr, cluster_name, node_id_hint: Option<String>, model_dir, auth: Arc<ClusterAuth>, assignment_timeout, commit_timeout } and CommittedAssignment { node_id, stage_index, cluster: ClusterConfig, coordinator_identity: String }:

1. Startup resolves the secret from --cluster-secret-file, MLXCEL_CLUSTER_SECRET_FILE, MLXCEL_CLUSTER_SECRET (ClusterSecretSources::from_env(cli, None)) and fails closed before any socket is bound. The peer computes its ModelIdentity once.
2. Bind a tokio::net::TcpListener on the peer listen address, taken from --cluster-control-addr (the flag gains a peer meaning, see 5). Static mode requires it with a non-zero port, since it must equal the coordinator's --cluster-peers entry. Mdns defaults it to 0.0.0.0:0 and uses the OS-assigned port.
3. Mdns only: spawn broadcast_beacon_loop with peer_addr set to the listener's local address (BeaconVerifier substitutes the source IP for an unspecified IP), node_id set to --node-id when given, interval 1 s, and the target from --cluster-discovery-target (see 4).
4. Accept connections until PEER_ASSIGNMENT_TIMEOUT = 600 s, then fail with "no stage assignment received within 600s" naming the listen address, discovery mode, and cluster name. Run at most 8 handshakes concurrently (excess sockets are closed at accept and logged with the same rate limit as handshake_warn_allowed in tcp_transport.rs). A failed handshake is logged and does not end the wait.
5. Verify the offer in this order and answer Rejected with the first failing reason: protocol version; mlxcel_version equals ours; cluster_name equals ours (--cluster-name, default mlxcel-cluster); model equals ours; layer_range equals the entry for stage_index from our own resolve_in_process_stage_assignments_for_model(model_dir, num_layers, Some(num_stages), None); cluster_toml parses with ClusterConfig::from_toml and holds node_id as a pipeline_stage node with stage == stage_index whose port equals our listener port; coordinator_control_addr agrees with the handshake-bound identity (same port, and same IP unless the offered IP is unspecified).
6. ProtocolMismatch, VersionMismatch, ModelMismatch, and LayerRangeMismatch end the process with the typed error AssignmentRefused { reason, detail }, because a retry against the same coordinator cannot succeed. Every other rejection, and a connection that drops, returns to step 4.
7. After Accepted, wait up to ASSIGNMENT_COMMIT_TIMEOUT = 60 s for Commit. On timeout or EOF discard the offer and return to step 4 (beacons keep running). On Commit, reply Committed, flush, stop the beacon loop (await its task for at most 1 s), drop the listener, and accept nothing more: exactly one committed assignment per run.
8. Startup then rewrites the received config: the coordinator node's address becomes the handshake-bound identity when the offered IP was unspecified; our own node's address becomes the listener's local address (the address the stage service can bind); secret_file stays None. It writes the result with render_deterministic_toml to --cluster-config-out (default .mlxcel/cluster-<node_id>.toml) and sets startup.distributed_config and startup.node_id = Some(node_id), mirroring the coordinator branch at startup.rs:2791-2794. A user --node-id that differs from the assigned id is overridden with an info log. --distributed-config given together with --pp-peer keeps today's behavior as a manual override: skip the assignment wait entirely.

3. Coordinator side (--pp-auto N)

Between plan_cluster (startup.rs:2775) and write_plan_toml, call push_stage_assignments(plan: &ClusterInitPlan, coordinator_control_addr: SocketAddr, model_dir: &Path, auth: &ClusterAuth, timeout: Duration) -> Result<()>. It computes ModelIdentity and the partition once, builds one offer per pipeline_stage node, and drives every peer concurrently: TCP connect retried every 500 ms (a static peer may start after the coordinator), handshake, Offer, wait for Accepted; then Commit to all and wait for Committed. The whole push is bounded by ASSIGNMENT_ACK_TIMEOUT = 60 s. Any failure returns AssignmentPushFailed { failures: Vec<PeerAssignmentFailure { node_id, addr, reason }> }, whose message names every failing peer with its reason (timeout plus the last connect error, handshake rejection, or the peer's Rejected reason and detail). No Commit is sent when any peer failed the offer phase. The existing path then continues unchanged.

Two coordinator fixes ride along. Bind the mdns receiver to Ipv4Addr::UNSPECIFIED (Ipv6Addr::UNSPECIFIED for a v6 control address) instead of the control IP, so broadcast datagrams arrive. After discovery, fail when the control address is loopback and any resolved peer is not, with a message naming --cluster-control-addr and --host, because stages reply to the control address.

4. Broadcast target

Add --cluster-discovery-target <IP> (both CLIs, peer side, default 255.255.255.255) and pass it to broadcast_beacon_loop as a new target: IpAddr parameter that replaces the hard-coded Ipv4Addr::BROADCAST. It accepts a directed subnet broadcast (192.168.1.255) or a unicast coordinator IP, including 127.0.0.1. It is needed for three reasons: limited broadcast leaves only through the interface of the default route, so a host with a separate cluster NIC or a Thunderbolt bridge needs a directed broadcast; loopback carries no broadcast, so the loopback integration and end-to-end tests cannot exercise the real sender without it; and a unicast target still works on networks that drop broadcast. Rejected: a coordinator-side setting, since the coordinator only listens.

5. CLI help and docs

Rewrite the help for --pp-auto, --pp-peer, --cluster-discovery, --cluster-control-addr, and the new --cluster-discovery-target in src/bin/mlx_server.rs:1427-1493 and src/main.rs:2597-2633: --pp-auto N pushes assignments to N peer processes; --pp-peer listens on --cluster-control-addr, announces itself by beacon in mdns mode, and waits for the assignment, with --distributed-config and --node-id optional; --cluster-discovery drops the "future zeroconf" wording and states HMAC-signed UDP beacons on --cluster-discovery-port; --cluster-control-addr means the control address on the coordinator and the stage listen address on a peer. In docs/distributed.md, replace :138-140 and :224-225 with a "Zero-config bring-up" subsection under "Server / multi-host path" that walks through two hosts for both static and mdns (secret distribution, coordinator plus two peers, --host 0.0.0.0 or --cluster-control-addr, UDP 49555 and the TCP stage ports in the firewall), and extend the Discovery beacons bullet (:195-200) to cover the assignment channel.

Scope

In scope: src/distributed/stage_assignment.rs and its tests (new); src/distributed/cluster_init.rs (broadcast_beacon_loop target parameter); src/distributed/mod.rs (re-exports); src/server/startup.rs (both run_zero_config_bring_up branches, receiver bind, loopback check, new cluster_discovery_target field at :403-416); src/server/cli_input.rs:291-303, src/bin/mlx_server.rs, src/main.rs (flag and help); docs/distributed.md; tests listed below.

Out of scope: DNS-SD or real mDNS (the mode keeps its name); TLS or per-frame MAC; NAT traversal; handshake protocol changes (owned by #2655); waiting for peer model load before the first request (pre-existing 30 s stage_timeout); remote parity of llama, gemma3, gemma4 (#2654).

Implementation Notes

• Reuse: client_handshake and server_handshake as merged by hardening(distributed): Bound pre-auth slots, type handshake rejections, extend cluster auth #2655; sign_beacon, BeaconVerifier, discover_peers, broadcast_beacon_loop; plan_cluster, render_deterministic_toml, ClusterConfig::from_toml; resolve_in_process_pipeline_num_layers and resolve_in_process_stage_assignments_for_model; ClusterSecretSources. The peer feeds the written TOML and node id into resolve_distributed_startup (startup.rs:2814) like a manual run; do not add a second stage startup path.
• Constraints: Linux and macOS (native Windows is already refused by ServerDistributedFlags::ensure_supported). The stage service rebinds the listener's port after the assignment listener is dropped; tokio sets SO_REUSEADDR on Unix. Rejected alternative: handing the bound listener to RemoteStageServiceHandle::spawn (src/distributed/pipeline/remote_service.rs:792), which binds on its own thread and runtime and would need a non-Clone field in RemoteStageServiceConfig.
• Edge cases: a static peer starting after the coordinator is covered by the connect retry; an mdns peer must beacon within the coordinator's 30 s discovery window, otherwise the coordinator fails and the peer keeps waiting until its own timeout; peers beyond N that were not selected time out with the step 4 message; a second coordinator after commit finds the listener closed; a coordinator that dies between Accepted and Commit leaves the peer waiting again; repeated beacons from one peer collapse in discover_peers (PeerKey); after commit the coordinator's first stage connect may race the peer's rebind, which Transport::connect already tolerates with a warning and a lazy reconnect.
• Error handling: the peer exits non-zero on AssignmentRefused and on its timeout; the coordinator exits non-zero on AssignmentPushFailed, before writing the TOML. Neither side logs the secret or the offer's cluster_toml above debug.

Acceptance Criteria

• --pp-peer no longer requires --distributed-config or --node-id; with static discovery it requires --cluster-control-addr with a concrete port.
• The coordinator pushes and commits assignments to every planned peer over the authenticated handshake and fails, naming each peer that did not commit, within ASSIGNMENT_ACK_TIMEOUT.
• A coordinator holding a different secret cannot assign a stage or a model; the peer keeps waiting and still accepts the real coordinator afterwards.
• A peer refuses a model identity, layer range, or mlxcel version mismatch with AssignmentRefused, and commits exactly one assignment per run.
• Every wait is bounded: handshake 5 s, frame I/O 10 s, commit 60 s, coordinator push 60 s, peer assignment 600 s.
• --cluster-discovery-target exists in both CLIs, and help text and docs/distributed.md describe the working flow with a two-host walkthrough for static and mdns.
• Integrated into the real CLI: mlxcel-server --pp-auto 2 --cluster-discovery=mdns plus two mlxcel-server --pp-peer --cluster-discovery=mdns processes serve a /v1/chat/completions request end to end. The work is not complete until this passes through the real binaries.
• Unit tests in src/distributed/stage_assignment_tests.rs: assignment_frames_round_trip_and_reject_oversize (each frame variant encodes and decodes equal; a 64 KiB + 1 frame and a truncated frame are errors); model_identity_changes_with_config_or_weight_files (tempdir model dirs); peer_commits_exactly_one_assignment (a second push after commit fails with AssignmentPushFailed); peer_refuses_model_mismatch_with_typed_error (peer returns AssignmentRefused { reason: ModelMismatch, .. }, coordinator error names stage-0); peer_refuses_version_mismatch (raw client after client_handshake sends an offer with mlxcel_version = "0.0.0"); peer_discards_offer_without_commit (raw coordinator sends Offer and drops; a later real push commits); peer_times_out_without_assignment (200 ms timeout); peer_ignores_coordinator_with_wrong_secret; coordinator_names_peers_that_did_not_commit (one listening peer and one dead address, 300 ms timeout; the error lists only the dead node id and address, and the live peer receives no Commit); peer_substitutes_bound_identity_for_unspecified_coordinator.
• cluster_init_tests::broadcast_beacon_loop_resigns_each_tick_and_resolves_against_coordinator sends with target 127.0.0.1 and drops the direct sender, so the loop itself must reach discover_peers.
• src/server/startup_tests.rs: pp_peer_static_requires_listen_addr_not_distributed_config (error names --cluster-control-addr; on current code it names --distributed-config) and pp_auto_rejects_loopback_control_with_remote_peers.
• Integration tests in tests/zero_config_cluster_init.rs, no weights (a tempdir with config.json holding num_hidden_layers = 4 and an empty model.safetensors): zero_config_static_bring_up_without_toml and zero_config_mdns_bring_up_without_toml run a coordinator (discover_peers, plan_cluster, push_stage_assignments) and two await_stage_assignment peers on loopback, the mdns one through broadcast_beacon_loop with target 127.0.0.1; both assert node ids stage-0/stage-1, a parseable TOML with a loopback coordinator address and no secret_file. These fail on current code.
• Ignored real-model tests in the same file, zero_config_cli_static_matches_single_node and zero_config_cli_mdns_matches_single_node, spawn three mlxcel-server processes (common::repo_binary_path("mlxcel-server")) on loopback with qwen3-0.6b-4bit and assert that greedy output equals a single-node mlxcel-server on the same prompt. Use qwen3 because llama, gemma3, and gemma4 remote parity is broken (fix(distributed): remote pipeline parity fails for llama, gemma3, gemma4 on GB10 #2654); the existing llama-based ignored test in this file is affected by that too.

Verification

Run once at the end, on the final head. On macOS drop --features cuda.

cargo test --profile test-fast --features cuda -p mlxcel --lib -- --test-threads=1 distributed::
cargo test --profile test-fast --features cuda -p mlxcel --lib -- --test-threads=1 server::startup
cargo test --profile test-fast --features cuda -p mlxcel --test zero_config_cluster_init -- --test-threads=1
cargo test --profile test-fast --features cuda -p mlxcel --test cli_help_consistency
cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests -- -D warnings
cargo fmt --check

Real model on the GB10 host. Build first outside the lock (a cargo build inside gpu-lock run hands the lock to the sccache daemon; sccache --stop-server frees it), then:

cargo test --profile test-fast --features cuda -p mlxcel --test zero_config_cluster_init --no-run
gpu-lock run --tag zero-config-e2e -- env MLXCEL_SDPA_DETERMINISTIC=1 \
  cargo test --profile test-fast --features cuda -p mlxcel --test zero_config_cluster_init -- \
  --ignored --test-threads=1 zero_config_cli_

Manual CLI check under the same lock, with MLXCEL_SDPA_DETERMINISTIC=1 and one shared --cluster-secret-file on every process: start two peers with mlxcel-server -m models/qwen3-0.6b-4bit --pp-peer --cluster-discovery=mdns --cluster-discovery-target 127.0.0.1 --cluster-control-addr 127.0.0.1:0 --cluster-config-out <distinct path>, then the coordinator with mlxcel-server -m models/qwen3-0.6b-4bit --pp-auto 2 --cluster-discovery=mdns --cluster-control-addr 127.0.0.1:19000 --port 18080. A pass is: both peers log their committed stage-0/stage-1, /v1/chat/completions on port 18080 returns 200 with temperature: 0, and its text equals a single-node mlxcel-server -m models/qwen3-0.6b-4bit --port 18090 answer to the same request. Repeat with --cluster-discovery=static, --cluster-peers on the coordinator, and fixed peer ports. Record the commands and outputs in the PR body.

Related: #2637, #2650, #2654, #2655.