Module parameter_server 

In-process toy parameter server with sharded embeddings.

ParameterServer is a prototype: it lives entirely inside one Rust process, owns sharded copies of the entity / relation embedding tables, and lets workers (typically super::worker::Worker) pull the latest parameters and push locally-computed gradients back. Two update modes are supported:

1. UpdateMode::Sync — pushes are buffered per-step until every expected worker has pushed; only then is the average applied to the parameters and a new step starts. This is the standard mini-batch SGD contract and gives reproducible convergence. The barrier is per-shard: a worker that pushed for shard A is free to push for shard B without waiting on shard A’s barrier to clear.

2. UpdateMode::Async — pushes are applied immediately with no barrier. This trades a small amount of staleness (workers may be working against a slightly outdated copy of the parameters) for higher throughput. We track a per-shard staleness counter (the number of pushes between the last pull and the next pull) so callers can bound the expected divergence.

The server is bounded by design to 4–8 workers / 4–8 shards for the prototype. Larger setups should use a real RPC-based parameter server.

All public methods are async and use tokio::sync::RwLock so shards can be pulled concurrently from many workers without contention; pushes acquire a write lock only on the affected shard.

Structs

ParameterServer

Sharded parameter server with sync/async update modes.

ParameterServerConfig

Configuration for ParameterServer.

ParameterServerStats

Aggregate stats reported by ParameterServer::stats.

ShardSnapshot

Public-facing snapshot of one shard’s contents.

Enums

UpdateMode

How ParameterServer::push applies gradients.