alien_core/stack_authoring/
pool.rs

1//! A compute pool as a stack file declares it.
2
3use serde::{Deserialize, Serialize};
4
5use super::{
6    failure_domain_spread, invalid, ComputeChoice, ComputePool, CpuAmount, DataDiskSize,
7    PoolRequirements, PoolScale,
8};
9use crate::error::Result;
10use crate::instance_catalog::Architecture;
11use crate::GpuSpec;
12
13/// Machine counts an autoscaling pool moves between.
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
16#[serde(rename_all = "camelCase", deny_unknown_fields)]
17pub struct AuthoredAutoscaling {
18    /// Fewest machines.
19    pub min: ComputeChoice,
20    /// Most machines.
21    pub max: ComputeChoice,
22}
23
24/// A pool of machines: what each needs, and how many run.
25///
26/// Set `machines` for a fixed count or `autoscaling` for a range, not both.
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
29#[serde(rename_all = "camelCase", deny_unknown_fields)]
30pub struct AuthoredPool {
31    /// CPU cores per machine: a number, or a quantity such as `"500m"`.
32    pub cpu: CpuAmount,
33    /// Memory per machine, for example `"16Gi"`.
34    pub memory: String,
35    /// Size of the disk holding the operating system. `"20Gi"` when unset.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub ephemeral_storage: Option<String>,
38    /// CPU architecture the machines must have.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub architecture: Option<Architecture>,
41    /// Whether the machines must support nested virtualization.
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub nested_virtualization: Option<bool>,
44    /// GPUs each machine must have.
45    #[serde(default, skip_serializing_if = "Option::is_none")]
46    pub gpu: Option<GpuSpec>,
47    /// An extra empty data disk each machine must provide.
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub data_disk: Option<DataDiskSize>,
50    /// A fixed number of machines, or a range a deployment chooses from.
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub machines: Option<ComputeChoice>,
53    /// Machines scale between a minimum and a maximum.
54    #[serde(default, skip_serializing_if = "Option::is_none")]
55    pub autoscaling: Option<AuthoredAutoscaling>,
56    /// How many failure domains the pool spreads across, 1 to 255.
57    #[serde(default, skip_serializing_if = "Option::is_none")]
58    pub failure_domain_spread: Option<i64>,
59}
60
61impl AuthoredPool {
62    /// The pool, once its scale is known to be one of the two forms.
63    pub fn to_compute_pool(&self) -> Result<ComputePool> {
64        let scale = match (&self.machines, &self.autoscaling) {
65            (Some(machines), None) => PoolScale::Fixed {
66                machines: machines.clone(),
67            },
68            (None, Some(autoscaling)) => PoolScale::Autoscale {
69                min: autoscaling.min.clone(),
70                max: autoscaling.max.clone(),
71            },
72            (Some(_), Some(_)) => {
73                return Err(invalid(
74                    "autoscaling",
75                    "set either `machines` (a fixed count) or `autoscaling`, not both",
76                ))
77            }
78            (None, None) => {
79                return Err(invalid(
80                    "machines",
81                    "a pool needs `machines`, a fixed count, or `autoscaling` with a minimum and a maximum",
82                ))
83            }
84        };
85        Ok(ComputePool {
86            requirements: PoolRequirements {
87                cpu: self.cpu.clone(),
88                memory: self.memory.clone(),
89                ephemeral_storage: self.ephemeral_storage.clone(),
90                architecture: self.architecture,
91                nested_virtualization: self.nested_virtualization,
92                gpu: self.gpu.clone(),
93                data_disk: self.data_disk.clone(),
94            },
95            scale,
96        })
97    }
98
99    /// The failure-domain spread, validated.
100    pub fn failure_domain_spread(&self) -> Result<Option<u8>> {
101        self.failure_domain_spread
102            .map(|spread| failure_domain_spread("failureDomainSpread", spread))
103            .transpose()
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::{CapacityGroupScalePolicy, ComputeChoiceRange, ErrorData};
111    use serde_json::json;
112
113    fn field(error: alien_error::AlienError<ErrorData>) -> String {
114        match error.error.expect("error data") {
115            ErrorData::StackDefinitionInvalid { field, .. } => field,
116            other => panic!("unexpected error {other:?}"),
117        }
118    }
119
120    #[test]
121    fn a_fixed_pool_becomes_a_capacity_group() {
122        let pool: AuthoredPool = serde_json::from_value(json!({
123            "cpu": 2, "memory": "4Gi", "machines": 3, "architecture": "x86_64",
124            "dataDisk": { "min": "0Gi", "default": "100Gi", "max": "1Ti" },
125            "failureDomainSpread": 2
126        }))
127        .unwrap();
128        let group = pool
129            .to_compute_pool()
130            .unwrap()
131            .to_capacity_group("workers")
132            .unwrap();
133        assert_eq!(group.group_id, "workers");
134        assert_eq!((group.min_size, group.max_size), (3, 3));
135        let profile = group.profile.unwrap();
136        assert_eq!(profile.cpu, "2");
137        assert_eq!(profile.memory_bytes, 4 * 1024 * 1024 * 1024);
138        assert_eq!(profile.ephemeral_storage_bytes, 20 * 1024 * 1024 * 1024);
139        assert_eq!(group.data_disk_gib_range.unwrap().max, 1024);
140        assert_eq!(pool.failure_domain_spread().unwrap(), Some(2));
141    }
142
143    #[test]
144    fn an_autoscaling_pool_takes_ranges() {
145        let pool: AuthoredPool = serde_json::from_value(json!({
146            "cpu": "500m", "memory": "1Gi",
147            "autoscaling": { "min": 1, "max": { "min": 2, "max": 10, "default": 4 } }
148        }))
149        .unwrap();
150        let group = pool
151            .to_compute_pool()
152            .unwrap()
153            .to_capacity_group("burst")
154            .unwrap();
155        assert_eq!((group.min_size, group.max_size), (1, 4));
156        match group.scale_policy.unwrap() {
157            CapacityGroupScalePolicy::Autoscale { min, max } => {
158                assert_eq!(
159                    min,
160                    ComputeChoiceRange {
161                        min: 1,
162                        max: 1,
163                        default: 1
164                    }
165                );
166                assert_eq!(max.max, 10);
167            }
168            other => panic!("expected autoscale, got {other:?}"),
169        }
170    }
171
172    #[test]
173    fn a_pool_needs_exactly_one_scale() {
174        let base = json!({ "cpu": 1, "memory": "1Gi" });
175        let none: AuthoredPool = serde_json::from_value(base.clone()).unwrap();
176        assert_eq!(field(none.to_compute_pool().unwrap_err()), "machines");
177        let mut both = base;
178        both["machines"] = json!(1);
179        both["autoscaling"] = json!({ "min": 1, "max": 2 });
180        let both: AuthoredPool = serde_json::from_value(both).unwrap();
181        assert_eq!(field(both.to_compute_pool().unwrap_err()), "autoscaling");
182    }
183
184    #[test]
185    fn the_failure_domain_spread_is_bounded() {
186        for spread in [0, 256, -1] {
187            let pool: AuthoredPool = serde_json::from_value(
188                json!({ "cpu": 1, "memory": "1Gi", "machines": 1, "failureDomainSpread": spread }),
189            )
190            .unwrap();
191            assert_eq!(
192                field(pool.failure_domain_spread().unwrap_err()),
193                "failureDomainSpread"
194            );
195        }
196    }
197}
alien_core/stack_authoring/pool.rs

alien_core/stack_authoring/
pool.rs
