Enum WireMessage 

pub enum WireMessage {
    JoinRequest {
        node_id: String,
        gpu_count: u32,
        backend: String,
    },
    JoinAccepted {
        worker_id: u32,
        total_workers: u32,
    },
    ShardAssignment {
        step: u64,
        shard_start: usize,
        shard_end: usize,
    },
    GradientPayload {
        step: u64,
        worker_id: u32,
        gradients: Vec<f32>,
        loss: f32,
        correct: usize,
        total: usize,
    },
    AveragedGradient {
        step: u64,
        gradients: Vec<f32>,
        global_loss: f32,
    },
    Heartbeat {
        node_id: String,
        timestamp_ms: u64,
    },
    Shutdown,
    BlockGradientPayload {
        step: u64,
        worker_id: u32,
        block_idx: u32,
        num_blocks: u32,
        gradients: Vec<f32>,
        component_sizes: Vec<u32>,
    },
    AveragedBlockGradient {
        step: u64,
        block_idx: u32,
        gradients: Vec<f32>,
        component_sizes: Vec<u32>,
    },
    NonBlockGradientPayload {
        step: u64,
        worker_id: u32,
        component: u8,
        gradients: Vec<f32>,
    },
    AveragedNonBlockGradient {
        step: u64,
        component: u8,
        gradients: Vec<f32>,
    },
}

Messages sent between coordinator and workers over TCP.

Each message is length-prefixed: [u32 big-endian length][payload bytes]. Payload is bincode-serialized WireMessage.

Variants

JoinRequest

Worker → Coordinator: request to join training

Fields

node_id: String

gpu_count: u32

backend: String

JoinAccepted

Coordinator → Worker: join accepted with assigned worker ID

Fields

worker_id: u32

total_workers: u32

ShardAssignment

Coordinator → Worker: here is your shard for this step

Fields

step: u64

shard_start: usize

shard_end: usize

GradientPayload

Worker → Coordinator: gradient data for this step

Fields

step: u64

worker_id: u32

gradients: Vec<f32>

Serialized f32 gradient vector (LoRA params + classifier head)

loss: f32

correct: usize

total: usize

AveragedGradient

Coordinator → Worker: averaged gradient for optimizer step

Fields

step: u64

gradients: Vec<f32>

global_loss: f32

Heartbeat

Bidirectional: heartbeat ping/pong

Fields

node_id: String

timestamp_ms: u64

Shutdown

Coordinator → Worker: training complete, shut down

BlockGradientPayload

Worker → Coordinator: gradient data for a single transformer block

Sent after backward pass for block[block_idx]. Contains 9 gradient components (w_q, w_k, w_v, w_o, gate, up, down, input_norm, post_attn_norm) concatenated into a single f32 vector with per-component sizes.

Fields

step: u64

worker_id: u32

block_idx: u32

num_blocks: u32

gradients: Vec<f32>

Concatenated f32 gradient vector for all components

component_sizes: Vec<u32>

Element count for each component (e.g., [HH, HD_kv, HD_kv, HH, HI, HI, I*H, H, H])

AveragedBlockGradient

Coordinator → Worker: averaged gradient for a single transformer block

Fields

step: u64

block_idx: u32

gradients: Vec<f32>

Concatenated averaged f32 gradient vector

component_sizes: Vec<u32>

Element count for each component (same layout as BlockGradientPayload)

NonBlockGradientPayload

Worker → Coordinator: gradient for a non-block component (LM head, final norm, embedding)

Fields

step: u64

worker_id: u32

component: u8

Component ID: 0=lm_head, 1=final_norm, 2=embedding

gradients: Vec<f32>

Flattened f32 gradient

AveragedNonBlockGradient

Coordinator → Worker: averaged gradient for a non-block component

Fields

step: u64

component: u8

Component ID: 0=lm_head, 1=final_norm, 2=embedding

gradients: Vec<f32>

Averaged f32 gradient

Implementations

impl WireMessage

pub fn to_bytes(&self) -> Vec<u8>

Serialize this message to bytes (length-prefixed).

Format: [4 bytes big-endian length][payload]

pub fn from_payload(payload: &[u8]) -> Result<Self, String>

Deserialize from a complete payload (without length prefix).

Errors

Returns error if payload is malformed.

Trait Implementations

impl Clone for WireMessage

fn clone(&self) -> WireMessage

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for WireMessage

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.