optirs_core/distributed/
parameter_server.rs

1use super::averaging::{AveragingStrategy, ParameterAverager};
2use crate::error::{OptimError, Result};
3use scirs2_core::ndarray::{Array, Dimension, ScalarOperand};
4use scirs2_core::numeric::Float;
5use std::collections::HashMap;
6use std::fmt::Debug;
7
8/// Synchronous parameter server for distributed training
9#[derive(Debug)]
10pub struct ParameterServer<A: Float, D: Dimension> {
11    /// Parameter averager
12    averager: ParameterAverager<A, D>,
13    /// Current global parameters
14    global_parameters: Vec<Array<A, D>>,
15    /// Node update counters
16    update_counts: HashMap<usize, usize>,
17    /// Expected updates per round
18    expected_updates_per_round: usize,
19    /// Current round number
20    current_round: usize,
21    /// Synchronization barrier
22    pending_updates: HashMap<usize, Vec<Array<A, D>>>,
23}
24
25impl<A: Float + ScalarOperand + Debug + Send + Sync, D: Dimension + Send + Sync>
26    ParameterServer<A, D>
27{
28    /// Create a new parameter server
29    pub fn new(
30        strategy: AveragingStrategy,
31        numnodes: usize,
32        expected_updates_per_round: usize,
33    ) -> Self {
34        Self {
35            averager: ParameterAverager::new(strategy, numnodes),
36            global_parameters: Vec::new(),
37            update_counts: HashMap::new(),
38            expected_updates_per_round,
39            current_round: 0,
40            pending_updates: HashMap::new(),
41        }
42    }
43
44    /// Initialize with global parameters
45    pub fn initialize(&mut self, initialparams: &[Array<A, D>]) -> Result<()> {
46        if self.expected_updates_per_round == 0
47            || self.expected_updates_per_round > self.averager.numnodes()
48        {
49            return Err(OptimError::InvalidConfig(format!(
50                "expected_updates_per_round ({}) must be in [1, numnodes={}]",
51                self.expected_updates_per_round,
52                self.averager.numnodes()
53            )));
54        }
55
56        self.averager.initialize(initialparams)?;
57        self.global_parameters = initialparams.to_vec();
58
59        // Initialize update counts
60        for nodeid in 0..self.averager.numnodes() {
61            self.update_counts.insert(nodeid, 0);
62        }
63
64        Ok(())
65    }
66
67    /// Submit parameter update from a node
68    ///
69    /// A node that has already submitted for the current (not-yet-aggregated)
70    /// round is rejected rather than silently overwritten -- `pending_updates`
71    /// is a map keyed by node id, so a resubmission would otherwise vanish
72    /// without raising the round's completion count, corrupting the barrier.
73    pub fn submit_update(&mut self, nodeid: usize, parameters: Vec<Array<A, D>>) -> Result<bool> {
74        if nodeid >= self.averager.numnodes() {
75            return Err(OptimError::InvalidConfig(format!(
76                "Node ID {} exceeds number of nodes {}",
77                nodeid,
78                self.averager.numnodes()
79            )));
80        }
81
82        if self.pending_updates.contains_key(&nodeid) {
83            return Err(OptimError::InvalidState(format!(
84                "Node {} already submitted an update for the current round (round {}); \
85                 call force_aggregation() to close the round before resubmitting",
86                nodeid,
87                self.current_round + 1
88            )));
89        }
90
91        // Store the update
92        self.pending_updates.insert(nodeid, parameters);
93        *self.update_counts.entry(nodeid).or_insert(0) += 1;
94
95        // Check if we have enough updates for this round
96        let ready_for_aggregation = self.pending_updates.len() >= self.expected_updates_per_round;
97
98        if ready_for_aggregation {
99            self.aggregate_and_update()?;
100        }
101
102        Ok(ready_for_aggregation)
103    }
104
105    /// Force aggregation with current pending updates
106    pub fn force_aggregation(&mut self) -> Result<()> {
107        if !self.pending_updates.is_empty() {
108            self.aggregate_and_update()?;
109        }
110        Ok(())
111    }
112
113    /// Internal aggregation and update
114    fn aggregate_and_update(&mut self) -> Result<()> {
115        // Convert pending updates to the format expected by averager
116        let node_params: Vec<(usize, Vec<Array<A, D>>)> = self.pending_updates.drain().collect();
117
118        // Perform averaging
119        self.averager.average_parameters(&node_params)?;
120
121        // Update global parameters
122        self.global_parameters = self.averager.get_averaged_parameters_cloned();
123
124        // Increment round
125        self.current_round += 1;
126
127        Ok(())
128    }
129
130    /// Get current global parameters
131    pub fn get_global_parameters(&self) -> &[Array<A, D>] {
132        &self.global_parameters
133    }
134
135    /// Get cloned global parameters
136    pub fn get_global_parameters_cloned(&self) -> Vec<Array<A, D>> {
137        self.global_parameters.clone()
138    }
139
140    /// Get current round number
141    pub fn current_round(&self) -> usize {
142        self.current_round
143    }
144
145    /// Get update count for a node
146    pub fn get_update_count(&self, nodeid: usize) -> usize {
147        self.update_counts.get(&nodeid).copied().unwrap_or(0)
148    }
149
150    /// Get number of pending updates
151    pub fn pending_updates_count(&self) -> usize {
152        self.pending_updates.len()
153    }
154
155    /// Set node weight for weighted averaging
156    pub fn set_node_weight(&mut self, nodeid: usize, weight: A) -> Result<()> {
157        self.averager.set_node_weight(nodeid, weight)
158    }
159
160    /// Reset server state
161    pub fn reset(&mut self) {
162        self.averager.reset();
163        self.update_counts.clear();
164        self.pending_updates.clear();
165        self.current_round = 0;
166
167        for nodeid in 0..self.averager.numnodes() {
168            self.update_counts.insert(nodeid, 0);
169        }
170    }
171}
optirs_core/distributed/parameter_server.rs

optirs_core/distributed/
parameter_server.rs
