1use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ScaleConfig {
10 pub device_count: usize,
12
13 pub points_per_device: usize,
15
16 pub batch_size: usize,
18
19 pub memory_per_device: usize,
21
22 pub concurrent_connections: usize,
24
25 pub target_ops_per_sec: u64,
27}
28
29impl Default for ScaleConfig {
30 fn default() -> Self {
31 Self::medium()
32 }
33}
34
35impl ScaleConfig {
36 pub fn devices(count: usize) -> Self {
38 Self {
39 device_count: count,
40 points_per_device: 100,
41 batch_size: (count / 100).max(100),
42 memory_per_device: 4096,
43 concurrent_connections: (count / 100).max(10),
44 target_ops_per_sec: (count * 10) as u64,
45 }
46 }
47
48 pub fn tiny() -> Self {
50 Self {
51 device_count: 10,
52 points_per_device: 10,
53 batch_size: 10,
54 memory_per_device: 1024,
55 concurrent_connections: 5,
56 target_ops_per_sec: 100,
57 }
58 }
59
60 pub fn small() -> Self {
62 Self {
63 device_count: 100,
64 points_per_device: 50,
65 batch_size: 50,
66 memory_per_device: 2048,
67 concurrent_connections: 20,
68 target_ops_per_sec: 1000,
69 }
70 }
71
72 pub fn medium() -> Self {
74 Self {
75 device_count: 1_000,
76 points_per_device: 100,
77 batch_size: 100,
78 memory_per_device: 4096,
79 concurrent_connections: 100,
80 target_ops_per_sec: 10_000,
81 }
82 }
83
84 pub fn large() -> Self {
86 Self {
87 device_count: 10_000,
88 points_per_device: 100,
89 batch_size: 500,
90 memory_per_device: 4096,
91 concurrent_connections: 500,
92 target_ops_per_sec: 100_000,
93 }
94 }
95
96 pub fn xlarge() -> Self {
98 Self {
99 device_count: 50_000,
100 points_per_device: 100,
101 batch_size: 1000,
102 memory_per_device: 4096,
103 concurrent_connections: 1000,
104 target_ops_per_sec: 500_000,
105 }
106 }
107
108 pub fn max() -> Self {
110 Self {
111 device_count: 100_000,
112 points_per_device: 100,
113 batch_size: 2000,
114 memory_per_device: 4096,
115 concurrent_connections: 2000,
116 target_ops_per_sec: 1_000_000,
117 }
118 }
119
120 pub fn with_device_count(mut self, count: usize) -> Self {
122 self.device_count = count;
123 self
124 }
125
126 pub fn with_points_per_device(mut self, count: usize) -> Self {
128 self.points_per_device = count;
129 self
130 }
131
132 pub fn with_batch_size(mut self, size: usize) -> Self {
134 self.batch_size = size;
135 self
136 }
137
138 pub fn with_memory_per_device(mut self, bytes: usize) -> Self {
140 self.memory_per_device = bytes;
141 self
142 }
143
144 pub fn total_points(&self) -> usize {
146 self.device_count * self.points_per_device
147 }
148
149 pub fn estimated_memory(&self) -> usize {
151 self.device_count * self.memory_per_device
152 }
153
154 pub fn estimated_memory_mb(&self) -> f64 {
156 self.estimated_memory() as f64 / 1024.0 / 1024.0
157 }
158
159 pub fn description(&self) -> String {
161 format!(
162 "{} devices, {} points/device ({} total), ~{:.1} MB",
163 self.device_count,
164 self.points_per_device,
165 self.total_points(),
166 self.estimated_memory_mb()
167 )
168 }
169}
170
171#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
173pub enum ScalePreset {
174 Tiny,
176 Small,
178 Medium,
180 Large,
182 XLarge,
184 Max,
186}
187
188impl ScalePreset {
189 pub fn to_config(self) -> ScaleConfig {
191 match self {
192 Self::Tiny => ScaleConfig::tiny(),
193 Self::Small => ScaleConfig::small(),
194 Self::Medium => ScaleConfig::medium(),
195 Self::Large => ScaleConfig::large(),
196 Self::XLarge => ScaleConfig::xlarge(),
197 Self::Max => ScaleConfig::max(),
198 }
199 }
200
201 pub fn all() -> &'static [ScalePreset] {
203 &[
204 Self::Tiny,
205 Self::Small,
206 Self::Medium,
207 Self::Large,
208 Self::XLarge,
209 Self::Max,
210 ]
211 }
212}
213
214impl From<ScalePreset> for ScaleConfig {
215 fn from(preset: ScalePreset) -> Self {
216 preset.to_config()
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn test_scale_presets() {
226 assert_eq!(ScaleConfig::tiny().device_count, 10);
227 assert_eq!(ScaleConfig::small().device_count, 100);
228 assert_eq!(ScaleConfig::medium().device_count, 1_000);
229 assert_eq!(ScaleConfig::large().device_count, 10_000);
230 assert_eq!(ScaleConfig::xlarge().device_count, 50_000);
231 assert_eq!(ScaleConfig::max().device_count, 100_000);
232 }
233
234 #[test]
235 fn test_scale_devices() {
236 let config = ScaleConfig::devices(5000);
237 assert_eq!(config.device_count, 5000);
238 }
239
240 #[test]
241 fn test_total_points() {
242 let config = ScaleConfig::devices(100).with_points_per_device(50);
243 assert_eq!(config.total_points(), 5000);
244 }
245
246 #[test]
247 fn test_estimated_memory() {
248 let config = ScaleConfig::devices(1024).with_memory_per_device(1024);
250 assert_eq!(config.estimated_memory(), 1024 * 1024);
251 assert!((config.estimated_memory_mb() - 1.0).abs() < 0.01);
252 }
253
254 #[test]
255 fn test_scale_preset_conversion() {
256 let config: ScaleConfig = ScalePreset::Large.into();
257 assert_eq!(config.device_count, 10_000);
258 }
259
260 #[test]
261 fn test_description() {
262 let config = ScaleConfig::tiny();
263 let desc = config.description();
264 assert!(desc.contains("10 devices"));
265 }
266}