somatize_core/strategy.rs
1//! Training strategies for distributed execution.
2//!
3//! A [`TrainingStrategy`] is a graph-level attribute that controls HOW the
4//! Scheduler distributes work across workers and HOW workers coordinate
5//! during training (gradient aggregation, state sync, communication).
6//!
7//! Only the description lives here. *Running* one — sharding inputs,
8//! calling workers in a round loop, aggregating gradients — is execution,
9//! and is in `somatize_runtime::strategy` along with the traits that
10//! describe it.
11//!
12//! Subgraphs inherit the parent's strategy unless overridden.
13
14use crate::filter::RemoteTarget;
15use crate::graph::NodeId;
16use serde::{Deserialize, Serialize};
17
18/// Training strategy — graph-level attribute, inherited by subgraphs.
19#[derive(Debug, Clone, Default, Serialize, Deserialize)]
20#[serde(tag = "type")]
21#[non_exhaustive]
22pub enum TrainingStrategy {
23 /// All nodes execute locally (default).
24 #[default]
25 Local,
26
27 /// Replicate the entire graph on N workers, each sees a data shard.
28 /// Gradients are aggregated after each step.
29 DataParallel {
30 /// Number of graph replicas — one worker and one data shard each.
31 num_replicas: usize,
32 /// How the replicas' gradients are combined after each step.
33 aggregation: GradientAggregation,
34 },
35
36 /// Arbitrary model partitioning: each Partition maps a set of
37 /// node IDs to a worker target. Any topology is supported.
38 ModelParallel {
39 /// Which nodes run where. Nodes not covered by any
40 /// [`Partition`] stay on the default (local) target.
41 partitions: Vec<Partition>,
42 /// How activations and gradients move between partitions.
43 communication: CommunicationProtocol,
44 },
45
46 /// Federated learning: data stays on workers, only model updates
47 /// are shared. The coordinator aggregates after each round.
48 Federated {
49 /// Total number of participating clients (the pool
50 /// [`ClientSelection`] draws from each round).
51 num_clients: usize,
52 /// Number of train→aggregate rounds to run.
53 rounds: usize,
54 /// How client updates are combined into the global model.
55 aggregation: FederatedAggregation,
56 /// Which clients participate in each round.
57 client_selection: ClientSelection,
58 },
59
60 /// Population-Based Training: evolutionary hyperparameter optimization.
61 /// Each generation trains a population, evaluates, then evolves.
62 PopulationBased {
63 /// Number of concurrently trained population members.
64 population_size: usize,
65 /// Number of train→evaluate→exploit/explore cycles.
66 generations: usize,
67 /// How underperformers copy from top performers.
68 exploit: ExploitStrategy,
69 /// How copied hyperparameters are mutated afterwards.
70 explore: ExploreStrategy,
71 },
72
73 /// User-defined strategy with a registered coordinator.
74 Custom {
75 /// Name identifying the user-provided coordinator. The
76 /// built-in strategy executor refuses this variant outright —
77 /// it never falls back to `Local` — so running it means
78 /// supplying your own coordination logic.
79 coordinator: String,
80 /// Opaque configuration passed through to the coordinator.
81 config: serde_json::Value,
82 },
83}
84
85/// How gradients are aggregated across workers in data-parallel training.
86#[derive(Debug, Clone, Serialize, Deserialize)]
87#[serde(tag = "method")]
88#[non_exhaustive]
89pub enum GradientAggregation {
90 /// All workers exchange gradients (ring or tree reduction).
91 AllReduce,
92 /// A central parameter server collects and distributes updates.
93 ParameterServer,
94 /// Decentralized gossip-based aggregation.
95 Decentralized {
96 /// Name of the gossip topology (e.g. `"ring"`), interpreted by
97 /// the executing backend.
98 topology: String,
99 },
100}
101
102/// A partition maps a set of node IDs to a worker target.
103///
104/// Used in `ModelParallel` to define which nodes run on which worker.
105/// The user has full control over the partitioning.
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct Partition {
108 /// The nodes this partition claims. Need not be contiguous in the
109 /// graph — any subset can be pinned to a target.
110 pub node_ids: Vec<NodeId>,
111 /// Worker the nodes run on, by id or by capability tag.
112 pub target: RemoteTarget,
113}
114
115/// How model-parallel partitions communicate activations and gradients.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117#[serde(tag = "protocol")]
118#[non_exhaustive]
119pub enum CommunicationProtocol {
120 /// Intermediate values flow via DataStore (S3, shared disk).
121 DataStore,
122 /// Direct point-to-point streaming between workers.
123 Direct,
124 /// Pipeline parallelism with micro-batching for overlap.
125 Pipeline {
126 /// Rows per micro-batch: smaller batches overlap partitions
127 /// more at the cost of per-message overhead.
128 micro_batch_size: usize,
129 },
130}
131
132/// Aggregation method for federated learning rounds.
133#[derive(Debug, Clone, Serialize, Deserialize)]
134#[serde(tag = "method")]
135#[non_exhaustive]
136pub enum FederatedAggregation {
137 /// Federated Averaging: weighted mean of client updates.
138 FedAvg,
139 /// FedProx: adds proximal term to prevent client drift.
140 FedProx {
141 /// Proximal term weight: larger values pull clients harder
142 /// toward the global model.
143 mu: f64,
144 },
145 /// FedYogi: adaptive federated optimization.
146 FedYogi {
147 /// First-moment (momentum) decay rate.
148 beta1: f64,
149 /// Second-moment decay rate.
150 beta2: f64,
151 /// Adaptivity floor: bounds how large the effective per-round
152 /// step can grow.
153 tau: f64,
154 },
155}
156
157/// How clients are selected per federated round.
158#[derive(Debug, Clone, Serialize, Deserialize)]
159#[serde(tag = "method")]
160#[non_exhaustive]
161pub enum ClientSelection {
162 /// All available clients participate.
163 All,
164 /// Random subset of clients.
165 Random {
166 /// Fraction of `num_clients` sampled each round, in `(0, 1]`.
167 fraction: f64,
168 },
169 /// Only clients matching specific tags.
170 ByCapability {
171 /// Capability tags a client must carry to be eligible.
172 required_tags: Vec<String>,
173 },
174}
175
176/// PBT exploit strategy: how underperformers learn from top performers.
177#[derive(Debug, Clone, Serialize, Deserialize)]
178#[serde(tag = "method")]
179#[non_exhaustive]
180pub enum ExploitStrategy {
181 /// Bottom fraction copies weights+hyperparams from top fraction.
182 Truncation {
183 /// Fraction of the population treated as top/bottom (the
184 /// `PbtRunner` clamps it to at most half the population).
185 fraction: f64,
186 },
187 /// Each member is compared to a random other; loser copies winner.
188 Binary {
189 /// Intended fitness margin the winner must exceed before the
190 /// loser copies it. The current `PbtRunner` copies on any
191 /// strict fitness loss and does not read this field yet.
192 threshold: f64,
193 },
194}
195
196/// PBT explore strategy: how hyperparameters are mutated after exploit.
197#[derive(Debug, Clone, Serialize, Deserialize)]
198#[serde(tag = "method")]
199#[non_exhaustive]
200pub enum ExploreStrategy {
201 /// Multiply each hyperparameter by a random factor in [1-factor, 1+factor].
202 Perturbation {
203 /// Half-width of the multiplicative jitter (0.2 → factors in
204 /// [0.8, 1.2]).
205 factor: f64,
206 },
207 /// Resample hyperparameters from the original search space.
208 Resample,
209}
210
211#[cfg(test)]
212mod tests {
213 use super::*;
214
215 #[test]
216 fn default_is_local() {
217 assert!(matches!(
218 TrainingStrategy::default(),
219 TrainingStrategy::Local
220 ));
221 }
222
223 #[test]
224 fn serde_roundtrip_data_parallel() {
225 let strategy = TrainingStrategy::DataParallel {
226 num_replicas: 4,
227 aggregation: GradientAggregation::AllReduce,
228 };
229 let json = serde_json::to_string(&strategy).unwrap();
230 let parsed: TrainingStrategy = serde_json::from_str(&json).unwrap();
231 assert!(matches!(
232 parsed,
233 TrainingStrategy::DataParallel {
234 num_replicas: 4,
235 ..
236 }
237 ));
238 }
239
240 #[test]
241 fn serde_roundtrip_model_parallel() {
242 let strategy = TrainingStrategy::ModelParallel {
243 partitions: vec![
244 Partition {
245 node_ids: vec!["embed".into(), "backbone".into()],
246 target: RemoteTarget::Tag("gpu-0".into()),
247 },
248 Partition {
249 node_ids: vec!["head_a".into()],
250 target: RemoteTarget::Tag("gpu-1".into()),
251 },
252 ],
253 communication: CommunicationProtocol::Pipeline {
254 micro_batch_size: 4,
255 },
256 };
257 let json = serde_json::to_string(&strategy).unwrap();
258 let parsed: TrainingStrategy = serde_json::from_str(&json).unwrap();
259 assert!(matches!(parsed, TrainingStrategy::ModelParallel { .. }));
260 }
261
262 #[test]
263 fn serde_roundtrip_federated() {
264 let strategy = TrainingStrategy::Federated {
265 num_clients: 10,
266 rounds: 50,
267 aggregation: FederatedAggregation::FedProx { mu: 0.01 },
268 client_selection: ClientSelection::Random { fraction: 0.3 },
269 };
270 let json = serde_json::to_string(&strategy).unwrap();
271 let parsed: TrainingStrategy = serde_json::from_str(&json).unwrap();
272 assert!(matches!(
273 parsed,
274 TrainingStrategy::Federated {
275 num_clients: 10,
276 rounds: 50,
277 ..
278 }
279 ));
280 }
281
282 #[test]
283 fn serde_roundtrip_pbt() {
284 let strategy = TrainingStrategy::PopulationBased {
285 population_size: 20,
286 generations: 50,
287 exploit: ExploitStrategy::Truncation { fraction: 0.2 },
288 explore: ExploreStrategy::Perturbation { factor: 0.2 },
289 };
290 let json = serde_json::to_string(&strategy).unwrap();
291 let parsed: TrainingStrategy = serde_json::from_str(&json).unwrap();
292 assert!(matches!(
293 parsed,
294 TrainingStrategy::PopulationBased {
295 population_size: 20,
296 ..
297 }
298 ));
299 }
300}