Skip to main content

alien_core/resources/
compute_cluster.rs

1//! ComputeCluster resource for long-running container workloads.
2//!
3//! A ComputeCluster represents the setup-owned compute boundary for containers.
4//! Setup provisions:
5//! - Auto Scaling Groups (AWS), Managed Instance Groups (GCP), or VM Scale Sets (Azure)
6//! - IAM roles/service accounts for machine authentication
7//! - Security groups/firewall rules
8//! - Launch templates/instance configurations
9
10use crate::error::{ErrorData, Result};
11use crate::instance_catalog::{is_same_architecture_machine, Architecture};
12use crate::resource::{ResourceDefinition, ResourceOutputsDefinition, ResourceRef};
13use crate::{PermissionProfile, Platform, ResourceType};
14use alien_error::AlienError;
15use bon::Builder;
16use serde::{Deserialize, Serialize};
17use std::any::Any;
18use std::collections::BTreeMap;
19use std::fmt::Debug;
20
21/// GPU specification for a capacity group.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
24#[serde(rename_all = "camelCase")]
25pub struct GpuSpec {
26    /// GPU type identifier (e.g., "nvidia-a100", "nvidia-t4")
27    #[serde(rename = "type")]
28    pub gpu_type: String,
29    /// Number of GPUs per machine
30    pub count: u32,
31}
32
33/// Machine resource profile for a capacity group.
34///
35/// Represents the hardware specifications for machines in a capacity group.
36/// These are hardware totals (what the instance type advertises), not allocatable
37/// capacity. The managed container scheduler internally subtracts system reserves for planning.
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
39#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
40#[serde(rename_all = "camelCase")]
41pub struct MachineProfile {
42    /// CPU cores per machine (hardware total) - stored as string to preserve precision
43    /// (e.g., "8.0", "4.5")
44    pub cpu: String,
45    /// Memory in bytes (hardware total)
46    pub memory_bytes: u64,
47    /// Ephemeral storage in bytes (hardware total)
48    pub ephemeral_storage_bytes: u64,
49    /// CPU architecture required or provided by this machine profile.
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub architecture: Option<Architecture>,
52    /// GPU specification (optional)
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub gpu: Option<GpuSpec>,
55    /// Additional empty block device each machine must provide, besides the
56    /// disk holding the operating system. On AWS, GCP and Azure the runtime
57    /// controller attaches one to every machine; on Machines, `join` refuses a
58    /// host without one. Kubernetes rejects it.
59    #[serde(default, skip_serializing_if = "Option::is_none")]
60    pub data_disk: Option<DataDiskRequirement>,
61}
62
63/// An additional block device a machine must provide with no partitions,
64/// no filesystem signature, and no mounts, so the workload can format it.
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
67#[serde(rename_all = "camelCase")]
68pub struct DataDiskRequirement {
69    /// Minimum device size in bytes.
70    pub size_bytes: u64,
71}
72
73/// Allowed range and default for a count selected by the installer.
74#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
75#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
76#[serde(rename_all = "camelCase")]
77pub struct ComputeChoiceRange {
78    /// Lowest allowed value.
79    pub min: u32,
80    /// Highest allowed value.
81    pub max: u32,
82    /// Default value recommended when no installer override is supplied.
83    pub default: u32,
84}
85
86impl ComputeChoiceRange {
87    /// Returns whether a selected value is inside the allowed range.
88    pub fn contains(&self, value: u32) -> bool {
89        self.min <= value && value <= self.max
90    }
91}
92
93/// Source-declared scale policy for a capacity group.
94#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
95#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
96#[serde(rename_all = "camelCase", tag = "type")]
97pub enum CapacityGroupScalePolicy {
98    /// A fixed-size pool where the installer can choose the fixed machine count.
99    Fixed {
100        /// Allowed fixed machine count range.
101        machines: ComputeChoiceRange,
102    },
103    /// An autoscaling pool with separately bounded min and max counts.
104    Autoscale {
105        /// Allowed minimum machine count range.
106        min: ComputeChoiceRange,
107        /// Allowed maximum machine count range.
108        max: ComputeChoiceRange,
109    },
110}
111
112impl CapacityGroupScalePolicy {
113    /// Derive the legacy policy represented by selected min/max values.
114    pub fn from_selected_bounds(min_size: u32, max_size: u32) -> Self {
115        if min_size == max_size {
116            Self::Fixed {
117                machines: ComputeChoiceRange {
118                    min: min_size,
119                    max: max_size,
120                    default: min_size,
121                },
122            }
123        } else {
124            Self::Autoscale {
125                min: ComputeChoiceRange {
126                    min: min_size,
127                    max: min_size,
128                    default: min_size,
129                },
130                max: ComputeChoiceRange {
131                    min: max_size,
132                    max: max_size,
133                    default: max_size,
134                },
135            }
136        }
137    }
138
139    /// Default selected min bound.
140    pub fn default_min_size(&self) -> u32 {
141        match self {
142            Self::Fixed { machines } => machines.default,
143            Self::Autoscale { min, .. } => min.default,
144        }
145    }
146
147    /// Default selected max bound.
148    pub fn default_max_size(&self) -> u32 {
149        match self {
150            Self::Fixed { machines } => machines.default,
151            Self::Autoscale { max, .. } => max.default,
152        }
153    }
154}
155
156/// Capacity group definition.
157///
158/// A capacity group represents machines with identical hardware profiles.
159/// Each group becomes a separate Auto Scaling Group (AWS), Managed Instance Group (GCP),
160/// or VM Scale Set (Azure).
161#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
162#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
163#[serde(rename_all = "camelCase")]
164pub struct CapacityGroup {
165    /// Unique identifier for this capacity group (must be lowercase alphanumeric with hyphens)
166    pub group_id: String,
167    /// Provider machine selected at deployment time.
168    /// `alien.ts` should declare portable requirements; preflight materialization
169    /// fills this field from `StackSettings.compute`.
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub instance_type: Option<String>,
172    /// Machine resource profile (auto-derived from instance_type if not specified)
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub profile: Option<MachineProfile>,
175    /// Minimum number of machines (can be 0 for scale-to-zero)
176    pub min_size: u32,
177    /// Maximum number of machines (must be ≤ 10)
178    pub max_size: u32,
179    /// Source-declared scale policy and installer-editable bounds.
180    ///
181    /// Older stacks only have `minSize` and `maxSize`; planners derive an exact
182    /// fixed/autoscale policy from those selected bounds when this field is absent.
183    #[serde(skip_serializing_if = "Option::is_none")]
184    pub scale_policy: Option<CapacityGroupScalePolicy>,
185    /// Require instance types that expose nested virtualization (VT-x/EPT)
186    /// to guest VMs. This is needed by workloads that boot nested VMs inside
187    /// containers.
188    /// When true, the controller's instance-type selector is constrained
189    /// to a vetted nested-virt-capable allowlist.
190    #[serde(skip_serializing_if = "Option::is_none")]
191    pub nested_virtualization: Option<bool>,
192    /// Data disk sizes, in GiB, a deployment may choose from. Present only when
193    /// the source declares a range; `profile.data_disk` then holds the default
194    /// (absent when the default is 0, which means no data disk). Preflight
195    /// writes the deployment's choice into `profile.data_disk`.
196    #[serde(default, skip_serializing_if = "Option::is_none")]
197    pub data_disk_gib_range: Option<ComputeChoiceRange>,
198}
199
200impl CapacityGroup {
201    /// AWS root volume size, including filesystem headroom and runtime overhead.
202    pub fn aws_root_volume_size_gib(&self) -> i32 {
203        self.aws_filesystem_size_gib()
204            .saturating_add(12)
205            .max(20)
206            .min(i32::MAX as u64) as i32
207    }
208
209    /// AWS second EBS volume: the declared data disk, or for a nested
210    /// virtualization pool without one, the volume such pools have always
211    /// received, sized from ephemeral storage.
212    pub fn aws_runtime_data_volume_size_gib(&self) -> Option<i32> {
213        if let Some(size_gib) = self.data_disk_size_gib() {
214            return Some(size_gib.min(i32::MAX as u64) as i32);
215        }
216        (self.nested_virtualization == Some(true))
217            .then(|| self.aws_filesystem_size_gib().max(20).min(i32::MAX as u64) as i32)
218    }
219
220    /// Whether `other` asks for different disks than this group: another root
221    /// storage size or data disk. Runtime controllers replace machines for it.
222    pub fn disks_differ(&self, other: &CapacityGroup) -> bool {
223        let disks = |group: &CapacityGroup| {
224            group
225                .profile
226                .as_ref()
227                .map(|profile| (profile.ephemeral_storage_bytes, profile.data_disk.clone()))
228        };
229        disks(self) != disks(other)
230    }
231
232    /// Data disk sizes a deployment may choose, in GiB: the declared range, or
233    /// just the declared size. `None` when the pool declares no data disk.
234    pub fn data_disk_choices(&self) -> Option<ComputeChoiceRange> {
235        if let Some(range) = &self.data_disk_gib_range {
236            return Some(range.clone());
237        }
238        let size_gib = u32::try_from(self.data_disk_size_gib()?).unwrap_or(u32::MAX);
239        Some(ComputeChoiceRange {
240            min: size_gib,
241            max: size_gib,
242            default: size_gib,
243        })
244    }
245
246    /// Data disk size, in GiB, this pool runs with for a deployment that chose
247    /// `selected` (`None` uses the default). 0 means no data disk.
248    pub fn effective_data_disk_gib(
249        &self,
250        selected: Option<u32>,
251    ) -> std::result::Result<Option<u32>, String> {
252        let Some(choices) = self.data_disk_choices() else {
253            return match selected {
254                None | Some(0) => Ok(None),
255                Some(size) => Err(format!(
256                    "pool '{}' declares no data disk, so dataDiskGib {size} cannot be chosen",
257                    self.group_id
258                )),
259            };
260        };
261        let size = selected.unwrap_or(choices.default);
262        if !choices.contains(size) {
263            return Err(if choices.min == choices.max {
264                format!(
265                    "pool '{}' has a fixed {} GiB data disk; dataDiskGib {size} cannot be chosen",
266                    self.group_id, choices.min
267                )
268            } else {
269                format!(
270                    "pool '{}' data disk must be between {} and {} GiB; dataDiskGib {size} is outside that range",
271                    self.group_id, choices.min, choices.max
272                )
273            });
274        }
275        Ok(Some(size))
276    }
277
278    /// The data disk this pool's machines get for a deployment that chose
279    /// `selected` GiB (`None` uses the default; 0 means none).
280    pub fn selected_data_disk(
281        &self,
282        selected: Option<u32>,
283    ) -> std::result::Result<Option<DataDiskRequirement>, String> {
284        let size_gib = self.effective_data_disk_gib(selected)?;
285        if self.data_disk_gib_range.is_none() {
286            // A fixed size stands as declared, to the byte: "500G" must keep
287            // accepting a 500 GB disk rather than the 466 GiB it rounds up to.
288            return Ok(self
289                .profile
290                .as_ref()
291                .and_then(|profile| profile.data_disk.clone()));
292        }
293        Ok(size_gib
294            .filter(|size_gib| *size_gib > 0)
295            .map(|size_gib| DataDiskRequirement {
296                size_bytes: u64::from(size_gib) * 1024 * 1024 * 1024,
297            }))
298    }
299
300    /// Declared data disk size, rounded up to a whole GiB.
301    pub fn data_disk_size_gib(&self) -> Option<u64> {
302        let data_disk = self.profile.as_ref()?.data_disk.as_ref()?;
303        Some(data_disk.size_bytes.div_ceil(1024 * 1024 * 1024))
304    }
305
306    fn aws_filesystem_size_gib(&self) -> u64 {
307        self.profile
308            .as_ref()
309            .map_or(0, |profile| profile.ephemeral_storage_bytes)
310            .div_ceil(1024 * 1024 * 1024)
311            .saturating_mul(5)
312            .div_ceil(4)
313    }
314}
315
316/// ComputeCluster resource for running long-running container workloads.
317///
318/// A ComputeCluster provides the setup-owned machine boundary for containers.
319/// Alien may manage the worker fleet inside that boundary when setup grants
320/// `compute-cluster/management`.
321///
322/// ## Architecture
323///
324/// - **Setup** creates cloud resources: ASGs/MIGs/VMSSs, IAM roles, security groups
325/// - **Alien** manages allowed fleet operations: machine count, runtime
326///   machine image rollout and, on AWS, the machine type within one CPU architecture
327/// - A node agent runs on each machine from the selected runtime image channel
328///
329/// ## Example
330///
331/// ```rust
332/// use alien_core::{CapacityGroup, ComputeCluster, MachineProfile};
333///
334/// let cluster = ComputeCluster::new("compute".to_string())
335///     .capacity_group(CapacityGroup {
336///         group_id: "general".to_string(),
337///         instance_type: None,
338///         profile: Some(MachineProfile {
339///             cpu: "4.0".to_string(),
340///             memory_bytes: 16 * 1024 * 1024 * 1024,
341///             ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
342///             architecture: None,
343///             gpu: None,
344///             data_disk: None,
345///         }),
346///         min_size: 1,
347///         max_size: 5,
348///         scale_policy: None,
349///         nested_virtualization: None,
350///         data_disk_gib_range: None,
351///     })
352///     .build();
353/// ```
354#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Builder)]
355#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
356#[serde(rename_all = "camelCase", deny_unknown_fields)]
357#[builder(start_fn = new)]
358pub struct ComputeCluster {
359    /// Unique identifier for the container cluster.
360    /// Must contain only alphanumeric characters, hyphens, and underscores.
361    #[builder(start_fn)]
362    pub id: String,
363
364    /// Capacity groups defining the machine pools for this cluster.
365    /// Each group becomes a separate ASG/MIG/VMSS.
366    #[builder(field)]
367    pub capacity_groups: Vec<CapacityGroup>,
368
369    /// Explicit grants for the node identity, keyed by concrete resource ID.
370    /// Independent of workload permission profiles; absent grants no data access.
371    #[serde(default, skip_serializing_if = "Option::is_none")]
372    pub node_permissions: Option<PermissionProfile>,
373
374    /// Cloud platforms on which the explicit node grants apply; absent applies everywhere.
375    #[serde(default, skip_serializing_if = "Option::is_none")]
376    pub node_permissions_platforms: Option<Vec<Platform>>,
377
378    /// Pool reserved for containers created after a deployment is installed.
379    /// If absent, the runtime uses the `general` pool when it exists.
380    #[serde(default, skip_serializing_if = "Option::is_none")]
381    pub dynamic_container_pool: Option<String>,
382
383    /// Concrete provider failure domains selected during setup, keyed by capacity group.
384    /// Empty preserves the existing aggregate layout when no spread policy is configured.
385    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
386    #[builder(default)]
387    pub selected_failure_domains: BTreeMap<String, Vec<String>>,
388
389    /// Requested failure-domain spread keyed by capacity group.
390    /// Empty preserves the existing aggregate layout.
391    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
392    #[builder(default)]
393    pub failure_domain_spread: BTreeMap<String, u8>,
394
395    /// Container CIDR block for internal container networking.
396    /// Auto-generated as "10.244.0.0/16" if not specified.
397    /// Each machine gets a /24 subnet from this range.
398    #[serde(skip_serializing_if = "Option::is_none")]
399    pub container_cidr: Option<String>,
400}
401
402impl ComputeCluster {
403    /// The resource type identifier for ComputeCluster
404    pub const RESOURCE_TYPE: ResourceType = ResourceType::from_static("compute-cluster");
405
406    /// Returns the cluster's unique identifier.
407    pub fn id(&self) -> &str {
408        &self.id
409    }
410
411    /// Returns the container CIDR, defaulting to "10.244.0.0/16" if not specified.
412    pub fn container_cidr(&self) -> &str {
413        self.container_cidr.as_deref().unwrap_or("10.244.0.0/16")
414    }
415}
416
417impl<S: compute_cluster_builder::State> ComputeClusterBuilder<S> {
418    /// Adds a capacity group to the cluster.
419    pub fn capacity_group(mut self, group: CapacityGroup) -> Self {
420        self.capacity_groups.push(group);
421        self
422    }
423}
424
425/// Status of a single capacity group within a ComputeCluster.
426#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
427#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
428#[serde(rename_all = "camelCase")]
429pub struct CapacityGroupStatus {
430    /// Capacity group ID
431    pub group_id: String,
432    /// Current number of machines
433    pub current_machines: u32,
434    /// Desired number of machines (from the managed container capacity plan)
435    pub desired_machines: u32,
436    /// Instance type being used
437    pub instance_type: String,
438}
439
440/// Outputs generated by a successfully provisioned ComputeCluster.
441#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
442#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
443#[serde(rename_all = "camelCase")]
444pub struct ComputeClusterOutputs {
445    /// Managed container cluster ID (workspace/project/deployment/resourceid format)
446    pub cluster_id: String,
447    /// Whether the managed container cluster is ready
448    pub horizon_ready: bool,
449    /// Status of each capacity group
450    pub capacity_group_statuses: Vec<CapacityGroupStatus>,
451    /// Total number of machines across all capacity groups
452    pub total_machines: u32,
453}
454
455impl ResourceOutputsDefinition for ComputeClusterOutputs {
456    fn get_resource_type(&self) -> ResourceType {
457        ComputeCluster::RESOURCE_TYPE.clone()
458    }
459
460    fn as_any(&self) -> &dyn Any {
461        self
462    }
463
464    fn box_clone(&self) -> Box<dyn ResourceOutputsDefinition> {
465        Box::new(self.clone())
466    }
467
468    fn outputs_eq(&self, other: &dyn ResourceOutputsDefinition) -> bool {
469        other.as_any().downcast_ref::<ComputeClusterOutputs>() == Some(self)
470    }
471
472    fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
473        serde_json::to_value(self)
474    }
475}
476
477impl ResourceDefinition for ComputeCluster {
478    fn get_resource_type(&self) -> ResourceType {
479        Self::RESOURCE_TYPE
480    }
481
482    fn id(&self) -> &str {
483        &self.id
484    }
485
486    fn replace_after_failed_create_is_safe(&self) -> bool {
487        // A create that failed waiting for instances may already run some; deleting it terminates
488        // them, and in a capacity-starved zone the capacity may not come back.
489        false
490    }
491
492    fn get_dependencies(&self) -> Vec<ResourceRef> {
493        // ComputeCluster has no static dependencies.
494        // Network dependency is platform-specific:
495        // - AWS/GCP/Azure: Added by ComputeClusterMutation
496        // - Local/Kubernetes: Not needed (Docker/K8s handles networking)
497        // Platform controllers use require_dependency() at runtime to access Network state.
498        Vec::new()
499    }
500
501    fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> {
502        let new_cluster = new_config
503            .as_any()
504            .downcast_ref::<ComputeCluster>()
505            .ok_or_else(|| {
506                AlienError::new(ErrorData::UnexpectedResourceType {
507                    resource_id: self.id.clone(),
508                    expected: Self::RESOURCE_TYPE,
509                    actual: new_config.get_resource_type(),
510                })
511            })?;
512
513        if self.id != new_cluster.id {
514            return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
515                resource_id: self.id.clone(),
516                reason: "the 'id' field is immutable".to_string(),
517            }));
518        }
519
520        // Container CIDR is immutable once set
521        if self.container_cidr.is_some()
522            && new_cluster.container_cidr.is_some()
523            && self.container_cidr != new_cluster.container_cidr
524        {
525            return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
526                resource_id: self.id.clone(),
527                reason: "the 'containerCidr' field is immutable once set".to_string(),
528            }));
529        }
530
531        // Validate capacity groups
532        for new_group in &new_cluster.capacity_groups {
533            if let Some(existing_group) = self
534                .capacity_groups
535                .iter()
536                .find(|g| g.group_id == new_group.group_id)
537            {
538                // The controller rolls the fleet onto another machine of the same architecture;
539                // the stack's images may not run on any other machine. Preflights also limit
540                // this to the installed cloud, which this check can't see.
541                if let (Some(old), Some(new)) = (
542                    existing_group.instance_type.as_deref(),
543                    new_group.instance_type.as_deref(),
544                ) {
545                    if old != new
546                        && ![Platform::Aws, Platform::Gcp, Platform::Azure]
547                            .into_iter()
548                            .any(|platform| is_same_architecture_machine(platform, old, new))
549                    {
550                        return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
551                            resource_id: self.id.clone(),
552                            reason: format!(
553                                "capacity group '{}' can't change machine from '{old}' to '{new}': only a catalog machine in the same cloud and of the same CPU architecture can replace it",
554                                new_group.group_id
555                            ),
556                        }));
557                    }
558                }
559            }
560        }
561
562        Ok(())
563    }
564
565    fn as_any(&self) -> &dyn Any {
566        self
567    }
568
569    fn as_any_mut(&mut self) -> &mut dyn Any {
570        self
571    }
572
573    fn box_clone(&self) -> Box<dyn ResourceDefinition> {
574        Box::new(self.clone())
575    }
576
577    fn resource_eq(&self, other: &dyn ResourceDefinition) -> bool {
578        other.as_any().downcast_ref::<ComputeCluster>() == Some(self)
579    }
580
581    fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
582        serde_json::to_value(self)
583    }
584}
585
586#[cfg(test)]
587mod tests {
588    use super::*;
589
590    #[test]
591    fn aws_disk_sizing_preserves_rounding_minimums_and_saturation() {
592        let gib = 1024 * 1024 * 1024;
593        for (bytes, root, data) in [
594            (0, 20, 20),
595            (gib, 20, 20),
596            (20 * gib, 37, 25),
597            (20 * gib + 1, 39, 27),
598            (u64::MAX, i32::MAX, i32::MAX),
599        ] {
600            let mut group = CapacityGroup {
601                group_id: "general".into(),
602                instance_type: None,
603                profile: Some(MachineProfile {
604                    cpu: "4.0".into(),
605                    memory_bytes: 8 * gib,
606                    ephemeral_storage_bytes: bytes,
607                    architecture: None,
608                    gpu: None,
609                    data_disk: None,
610                }),
611                min_size: 1,
612                max_size: 1,
613                scale_policy: None,
614                nested_virtualization: None,
615                data_disk_gib_range: None,
616            };
617            assert_eq!(group.aws_root_volume_size_gib(), root);
618            assert_eq!(group.aws_runtime_data_volume_size_gib(), None);
619            group.nested_virtualization = Some(true);
620            assert_eq!(group.aws_root_volume_size_gib(), root);
621            assert_eq!(group.aws_runtime_data_volume_size_gib(), Some(data));
622        }
623    }
624
625    /// A fixed size keeps its declared bytes, so a host disk that met it still
626    /// does; a size picked from a range is whole GiB.
627    #[test]
628    fn a_fixed_data_disk_keeps_its_declared_bytes() {
629        let gib: u64 = 1024 * 1024 * 1024;
630        let declared = DataDiskRequirement {
631            size_bytes: 500_000_000_000,
632        };
633        let mut group = CapacityGroup {
634            group_id: "general".into(),
635            instance_type: None,
636            profile: Some(MachineProfile {
637                cpu: "4.0".into(),
638                memory_bytes: 8 * gib,
639                ephemeral_storage_bytes: 20 * gib,
640                architecture: None,
641                gpu: None,
642                data_disk: Some(declared.clone()),
643            }),
644            min_size: 1,
645            max_size: 1,
646            scale_policy: None,
647            nested_virtualization: None,
648            data_disk_gib_range: None,
649        };
650        // 500 GB rounds up to 466 GiB, which is the one size a deployment may name.
651        assert_eq!(group.selected_data_disk(None), Ok(Some(declared.clone())));
652        assert_eq!(group.selected_data_disk(Some(466)), Ok(Some(declared)));
653        assert!(group.selected_data_disk(Some(500)).is_err());
654
655        group.data_disk_gib_range = Some(ComputeChoiceRange {
656            min: 0,
657            max: 1024,
658            default: 466,
659        });
660        assert_eq!(
661            group.selected_data_disk(None),
662            Ok(Some(DataDiskRequirement {
663                size_bytes: 466 * gib
664            }))
665        );
666        assert_eq!(group.selected_data_disk(Some(0)), Ok(None));
667    }
668
669    #[test]
670    fn declared_data_disk_replaces_the_nested_virtualization_volume() {
671        let gib = 1024 * 1024 * 1024;
672        let group = |nested: Option<bool>, data_disk_bytes: u64| CapacityGroup {
673            group_id: "general".into(),
674            instance_type: None,
675            profile: Some(MachineProfile {
676                cpu: "4.0".into(),
677                memory_bytes: 8 * gib,
678                ephemeral_storage_bytes: 20 * gib,
679                architecture: None,
680                gpu: None,
681                data_disk: Some(DataDiskRequirement {
682                    size_bytes: data_disk_bytes,
683                }),
684            }),
685            min_size: 1,
686            max_size: 1,
687            scale_policy: None,
688            nested_virtualization: nested,
689            data_disk_gib_range: None,
690        };
691        for nested in [None, Some(false), Some(true)] {
692            // Root keeps the ephemeral storage formula; the second volume is
693            // exactly the declared size, not the 25 GiB implicit volume.
694            assert_eq!(group(nested, 30 * gib).aws_root_volume_size_gib(), 37);
695            assert_eq!(
696                group(nested, 30 * gib).aws_runtime_data_volume_size_gib(),
697                Some(30)
698            );
699            assert_eq!(
700                group(nested, 30 * gib + 1).aws_runtime_data_volume_size_gib(),
701                Some(31)
702            );
703        }
704        assert_eq!(group(None, 512 * 1024 * 1024).data_disk_size_gib(), Some(1));
705    }
706
707    #[test]
708    fn test_compute_cluster_creation() {
709        let cluster = ComputeCluster::new("compute".to_string())
710            .capacity_group(CapacityGroup {
711                group_id: "general".to_string(),
712                instance_type: Some("m7g.xlarge".to_string()),
713                profile: None,
714                min_size: 1,
715                max_size: 5,
716                scale_policy: None,
717                nested_virtualization: None,
718                data_disk_gib_range: None,
719            })
720            .build();
721
722        assert_eq!(cluster.id(), "compute");
723        assert_eq!(cluster.capacity_groups.len(), 1);
724        assert_eq!(cluster.capacity_groups[0].group_id, "general");
725        assert_eq!(cluster.container_cidr(), "10.244.0.0/16");
726    }
727
728    #[test]
729    fn test_compute_cluster_multiple_capacity_groups() {
730        let cluster = ComputeCluster::new("multi-pool".to_string())
731            .capacity_group(CapacityGroup {
732                group_id: "general".to_string(),
733                instance_type: Some("m7g.xlarge".to_string()),
734                profile: None,
735                min_size: 1,
736                max_size: 3,
737                scale_policy: None,
738                nested_virtualization: None,
739                data_disk_gib_range: None,
740            })
741            .capacity_group(CapacityGroup {
742                group_id: "gpu".to_string(),
743                instance_type: Some("g5.xlarge".to_string()),
744                profile: Some(MachineProfile {
745                    cpu: "4.0".to_string(),
746                    memory_bytes: 17179869184,             // 16 GiB
747                    ephemeral_storage_bytes: 214748364800, // 200 GiB
748                    architecture: None,
749                    gpu: Some(GpuSpec {
750                        gpu_type: "nvidia-a10g".to_string(),
751                        count: 1,
752                    }),
753                    data_disk: None,
754                }),
755                min_size: 0,
756                max_size: 2,
757                scale_policy: None,
758                nested_virtualization: None,
759                data_disk_gib_range: None,
760            })
761            .build();
762
763        assert_eq!(cluster.capacity_groups.len(), 2);
764        assert_eq!(cluster.capacity_groups[0].group_id, "general");
765        assert_eq!(cluster.capacity_groups[1].group_id, "gpu");
766        assert!(cluster.capacity_groups[1]
767            .profile
768            .as_ref()
769            .unwrap()
770            .gpu
771            .is_some());
772    }
773
774    #[test]
775    fn test_compute_cluster_custom_cidr() {
776        let cluster = ComputeCluster::new("custom-net".to_string())
777            .container_cidr("172.30.0.0/16".to_string())
778            .capacity_group(CapacityGroup {
779                group_id: "general".to_string(),
780                instance_type: None,
781                profile: None,
782                min_size: 1,
783                max_size: 5,
784                scale_policy: None,
785                nested_virtualization: None,
786                data_disk_gib_range: None,
787            })
788            .build();
789
790        assert_eq!(cluster.container_cidr(), "172.30.0.0/16");
791    }
792
793    #[test]
794    fn test_compute_cluster_validate_update_immutable_id() {
795        let cluster1 = ComputeCluster::new("cluster-1".to_string())
796            .capacity_group(CapacityGroup {
797                group_id: "general".to_string(),
798                instance_type: None,
799                profile: None,
800                min_size: 1,
801                max_size: 5,
802                scale_policy: None,
803                nested_virtualization: None,
804                data_disk_gib_range: None,
805            })
806            .build();
807
808        let cluster2 = ComputeCluster::new("cluster-2".to_string())
809            .capacity_group(CapacityGroup {
810                group_id: "general".to_string(),
811                instance_type: None,
812                profile: None,
813                min_size: 1,
814                max_size: 5,
815                scale_policy: None,
816                nested_virtualization: None,
817                data_disk_gib_range: None,
818            })
819            .build();
820
821        let result = cluster1.validate_update(&cluster2);
822        assert!(result.is_err());
823    }
824
825    #[test]
826    fn test_compute_cluster_validate_update_scale_change() {
827        let cluster1 = ComputeCluster::new("compute".to_string())
828            .capacity_group(CapacityGroup {
829                group_id: "general".to_string(),
830                instance_type: Some("m7g.xlarge".to_string()),
831                profile: None,
832                min_size: 1,
833                max_size: 5,
834                scale_policy: None,
835                nested_virtualization: None,
836                data_disk_gib_range: None,
837            })
838            .build();
839
840        let cluster2 = ComputeCluster::new("compute".to_string())
841            .capacity_group(CapacityGroup {
842                group_id: "general".to_string(),
843                instance_type: Some("m7g.xlarge".to_string()),
844                profile: None,
845                min_size: 2,
846                max_size: 10,
847                scale_policy: None,
848                nested_virtualization: None,
849                data_disk_gib_range: None,
850            })
851            .build();
852
853        // Scale changes should be allowed
854        let result = cluster1.validate_update(&cluster2);
855        assert!(result.is_ok());
856    }
857
858    #[test]
859    fn machine_changes_stay_within_one_cloud_and_architecture() {
860        let cluster = |machine: &str| {
861            ComputeCluster::new("compute".to_string())
862                .capacity_group(CapacityGroup {
863                    group_id: "general".to_string(),
864                    instance_type: Some(machine.to_string()),
865                    profile: None,
866                    min_size: 1,
867                    max_size: 5,
868                    scale_policy: None,
869                    nested_virtualization: None,
870                    data_disk_gib_range: None,
871                })
872                .build()
873        };
874        for (old, new) in [
875            ("t4g.small", "t4g.medium"),
876            ("e2-medium", "n2-standard-4"),
877            ("Standard_D2s_v5", "Standard_D4s_v5"),
878        ] {
879            cluster(old)
880                .validate_update(&cluster(new))
881                .expect("same cloud and architecture");
882        }
883        for (old, new) in [
884            ("t4g.small", "m7i.large"),
885            ("t4g.small", "t4g.unknown"),
886            ("n2-standard-2", "m7i.large"),
887            ("Standard_D2s_v5", "n2-standard-4"),
888        ] {
889            let error = cluster(old)
890                .validate_update(&cluster(new))
891                .expect_err("needs setup");
892            assert_eq!(error.code, "INVALID_RESOURCE_UPDATE", "{old} -> {new}");
893        }
894    }
895
896    #[test]
897    fn node_permissions_are_inline_optional_and_roundtrip_exactly() {
898        let mut cluster = ComputeCluster::new("compute".to_string()).build();
899        let absent = serde_json::to_value(&cluster).unwrap();
900        assert!(absent.get("nodePermissions").is_none());
901        assert_eq!(
902            serde_json::from_value::<ComputeCluster>(absent).unwrap(),
903            cluster
904        );
905        cluster.node_permissions =
906            Some(PermissionProfile::new().resource("objects", ["storage/data-read"]));
907        let json = serde_json::to_value(&cluster).unwrap();
908        assert_eq!(
909            json["nodePermissions"],
910            serde_json::json!({"objects": ["storage/data-read"]})
911        );
912        assert_eq!(
913            serde_json::from_value::<ComputeCluster>(json).unwrap(),
914            cluster
915        );
916        assert_eq!(ResourceDefinition::get_permissions(&cluster), None);
917    }
918
919    #[test]
920    fn test_compute_cluster_serialization() {
921        let cluster = ComputeCluster::new("test-cluster".to_string())
922            .capacity_group(CapacityGroup {
923                group_id: "general".to_string(),
924                instance_type: Some("m7g.xlarge".to_string()),
925                profile: None,
926                min_size: 1,
927                max_size: 5,
928                scale_policy: None,
929                nested_virtualization: None,
930                data_disk_gib_range: None,
931            })
932            .build();
933
934        let json = serde_json::to_string(&cluster).unwrap();
935        let deserialized: ComputeCluster = serde_json::from_str(&json).unwrap();
936        assert_eq!(cluster, deserialized);
937    }
938}