Struct ParameterServer 

pub struct ParameterServer { /* private fields */ }

Sharded parameter server with sync/async update modes.

See module-level docs for semantics.

Implementations

impl ParameterServer

pub fn new( config: ParameterServerConfig, entity_ids: Vec<String>, relation_ids: Vec<String>, shard_manager: ModelShardManager, ) -> Result<Self>

Build a new parameter server.

entity_ids and relation_ids must list every entity / relation the server should own; the server hashes each entity ID into a shard via the supplied ModelShardManager. Embedding rows are initialised to small uniform values in [-0.05, 0.05] for reproducibility.

pub fn num_shards(&self) -> usize

Number of shards.

pub fn config(&self) -> &ParameterServerConfig

Configuration snapshot.

pub fn shard_manager(&self) -> &ModelShardManager

Shard manager view (read-only).

pub async fn pull(&self, shard_id: usize) -> Result<ShardSnapshot>

Pull a snapshot of shard_id.

Returns owned data so workers can compute gradients without holding any lock. Panics-free: an out-of-range shard_id returns an Err.

pub async fn push( &self, shard_id: usize, worker_id: u32, rows: Vec<(usize, Vec<f32>)>, ) -> Result<()>

Push entity-row gradients to shard_id.

rows is (row_index_inside_shard, gradient_row) pairs. In UpdateMode::Sync the push is buffered until expected_workers pushes have accumulated for this shard or the per-step barrier timeout fires. In UpdateMode::Async the gradient is applied immediately and the per-shard staleness counter is incremented.

pub async fn push_relation( &self, worker_id: u32, rows: Vec<(usize, Vec<f32>)>, ) -> Result<()>

Apply gradients to relation rows.

Relation gradients are always averaged across the most-recent push regardless of update_mode — they’re a tiny fully-replicated table.

pub async fn stats(&self) -> ParameterServerStats

Snapshot of current stats.

pub async fn shard_steps(&self) -> Vec<u64>

Total per-shard step counts (for tests).