Struct GradientServer 

pub struct GradientServer { /* private fields */ }

Gradient server running on the coordinator node.

Implementations

impl GradientServer

pub fn bind(config: DistributedConfig) -> Result<Self, String>

Create and bind the gradient server.

Errors

Returns error if binding fails.

pub fn wait_for_workers(&mut self) -> Result<(), String>

Wait for all expected workers to connect.

Blocks until expect_workers workers have sent JoinRequest messages.

Errors

Returns error if any connection fails or timeout is exceeded.

pub fn set_total_samples(&mut self, n: usize)

Set total sample count for sharding

pub fn send_shard_assignments(&mut self, step: u64) -> Result<(), String>

Send shard assignments to all workers for a given step.

Errors

Returns error if any send fails.

pub fn collect_and_reduce( &mut self, step: u64, ) -> Result<AllReduceResult, String>

Collect gradients from all workers and compute AllReduce.

Contract: F-DP-003

If any gradient contains NaN/Inf, returns an error (Jidoka halt).

Errors

Returns error on communication failure or non-finite gradient.

pub fn broadcast_averaged( &mut self, step: u64, result: &AllReduceResult, ) -> Result<(), String>

Broadcast averaged gradients to all workers.

Errors

Returns error if any send fails.

pub fn shutdown_workers(&mut self)

Send shutdown message to all workers.

pub fn worker_count(&self) -> usize

Number of connected workers

pub fn collect_and_reduce_block( &mut self, step: u64, block_idx: u32, ) -> Result<BlockAllReduceResult, String>

Collect and reduce per-block gradients from all workers.

Waits for BlockGradientPayload from each worker for the specified block index, averages them, and returns the result.

Contract: C-DDP-001

Output equals arithmetic mean of all workers’ block gradients. Jidoka halt on NaN/Inf gradients (F-DP-003).

Errors

Returns error on communication failure, step mismatch, or NaN gradient.

pub fn broadcast_averaged_block( &mut self, step: u64, result: &BlockAllReduceResult, ) -> Result<(), String>

Broadcast averaged block gradient to all workers.

Errors

Returns error if any send fails.

pub fn collect_and_reduce_non_block( &mut self, step: u64, expected_component: u8, ) -> Result<NonBlockAllReduceResult, String>

Collect and reduce non-block gradient from all workers.

Used for LM head, final norm, and embedding gradients.

Errors

Returns error on communication failure or NaN gradient.

pub fn broadcast_averaged_non_block( &mut self, step: u64, result: &NonBlockAllReduceResult, ) -> Result<(), String>

Broadcast averaged non-block gradient to all workers.

impl GradientServer

pub fn local_addr(&self) -> SocketAddr

Get the local address this server is listening on.

Useful when binding to port 0 (OS-assigned) in tests.