Skip to main content

alien_core/
instance_catalog.rs

1//! Instance type catalog and selection algorithm for cloud compute infrastructure.
2//!
3//! This module provides:
4//! - A static catalog of known instance types across AWS, GCP, and Azure
5//! - Resource quantity parsing (CPU strings, Kubernetes-style memory/storage quantities)
6//! - An algorithm to select the optimal instance type for a given workload
7//!
8//! The catalog is the single source of truth for instance type specifications.
9//! It is used by the preflights system to automatically populate `CapacityGroup.instance_type`
10//! and `CapacityGroup.profile` based on the containers in a stack.
11
12use crate::{GpuSpec, MachineProfile, Platform};
13use serde::{Deserialize, Serialize};
14
15// ---------------------------------------------------------------------------
16// Resource quantity parsing
17// ---------------------------------------------------------------------------
18
19/// Parse a CPU quantity string to f64.
20///
21/// Accepts plain numbers ("1", "0.5", "2.0") and millicore suffixes ("500m" = 0.5).
22pub fn parse_cpu(s: &str) -> Result<f64, String> {
23    let s = s.trim();
24    if s.is_empty() {
25        return Err("empty CPU string".to_string());
26    }
27
28    if let Some(millis) = s.strip_suffix('m') {
29        let v: f64 = millis
30            .parse()
31            .map_err(|_| format!("invalid CPU millicore value: '{s}'"))?;
32        Ok(v / 1000.0)
33    } else {
34        s.parse().map_err(|_| format!("invalid CPU value: '{s}'"))
35    }
36}
37
38/// Parse a memory or storage quantity string to bytes.
39///
40/// Supports Kubernetes-style binary suffixes (Ki, Mi, Gi, Ti) and
41/// decimal suffixes (k, M, G, T). Plain numbers are interpreted as bytes.
42pub fn parse_memory_bytes(s: &str) -> Result<u64, String> {
43    let s = s.trim();
44    if s.is_empty() {
45        return Err("empty memory/storage string".to_string());
46    }
47
48    for (suffix, multiplier) in [
49        ("Ti", 1024_f64.powi(4)),
50        ("Gi", 1024_f64.powi(3)),
51        ("Mi", 1024_f64.powi(2)),
52        ("Ki", 1024_f64),
53        ("T", 1_000_000_000_000_f64),
54        ("G", 1_000_000_000_f64),
55        ("M", 1_000_000_f64),
56        ("k", 1_000_f64),
57    ] {
58        if let Some(number) = s.strip_suffix(suffix) {
59            let value: f64 = number
60                .parse()
61                .map_err(|_| format!("invalid memory value: '{s}'"))?;
62            let bytes = value * multiplier;
63            // Float-to-integer casts saturate negatives, infinity and NaN instead of
64            // returning an error. Reject them before a sizing request reaches a provider.
65            if !bytes.is_finite() || bytes < 0.0 || bytes >= u64::MAX as f64 {
66                return Err(format!("memory value is out of range: '{s}'"));
67            }
68            return Ok(bytes as u64);
69        }
70    }
71    s.parse()
72        .map_err(|_| format!("invalid memory value: '{s}'"))
73}
74
75// ---------------------------------------------------------------------------
76// Instance type catalog
77// ---------------------------------------------------------------------------
78
79/// Instance family classification.
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
81#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
82#[serde(rename_all = "camelCase")]
83pub enum InstanceFamily {
84    Burstable,
85    GeneralPurpose,
86    ComputeOptimized,
87    MemoryOptimized,
88    StorageOptimized,
89    GpuCompute,
90}
91
92/// CPU architecture.
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
94#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
95#[serde(rename_all = "snake_case")]
96pub enum Architecture {
97    Arm64,
98    X86_64,
99}
100
101/// Default machine architecture for images built for a managed cloud.
102pub fn default_architecture(platform: Platform) -> Option<Architecture> {
103    match platform {
104        Platform::Aws => Some(Architecture::Arm64),
105        Platform::Gcp | Platform::Azure => Some(Architecture::X86_64),
106        Platform::Kubernetes | Platform::Machines | Platform::Local | Platform::Test => None,
107    }
108}
109
110/// Static GPU specification for catalog entries (no heap allocation).
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub struct CatalogGpu {
113    pub gpu_type: &'static str,
114    pub count: u32,
115}
116
117/// A known instance type with its hardware specifications.
118///
119/// All fields are compile-time constants. The catalog is a flat array of these.
120#[derive(Debug, Clone)]
121pub struct InstanceTypeSpec {
122    pub name: &'static str,
123    pub platform: Platform,
124    /// The category the machine's series belongs to.
125    pub family: InstanceFamily,
126    /// The provider series the machine belongs to (`m8i`, `n2-standard`, `Dsv5`).
127    pub series: &'static str,
128    /// How the series is shown to people choosing a machine.
129    pub label: &'static str,
130    pub architecture: Architecture,
131    /// vCPU count (hardware total)
132    pub vcpu: u32,
133    /// Memory in bytes (hardware total)
134    pub memory_bytes: u64,
135    /// Ephemeral storage in bytes (hardware total, NVMe for storage-optimized)
136    pub ephemeral_storage_bytes: u64,
137    /// Local instance disk the provider reports for the machine, in bytes; 0 when it has none.
138    pub local_disk_bytes: u64,
139    /// The provider's machine type for this row. It differs from `name` only for a catalog
140    /// row that adds an accelerator to a machine type, such as `n1-standard-4-t4`.
141    pub machine_type: &'static str,
142    /// GPU specification (for GPU instances)
143    pub gpu: Option<CatalogGpu>,
144    /// Whether the machine exposes nested virtualization (VT-x/EPT) to guest VMs.
145    pub nested_virt: bool,
146    /// Whether automatic selection may pick the machine. The set is fixed so that adding a
147    /// machine to the catalog never changes what an existing stack gets; any other machine is
148    /// used only when someone chooses it.
149    pub auto_select: bool,
150}
151
152impl InstanceTypeSpec {
153    /// Whether ephemeral storage is a provider disk that can be sized at
154    /// deployment time instead of fixed local instance storage.
155    pub fn has_configurable_ephemeral_storage(&self) -> bool {
156        matches!(
157            self.family,
158            InstanceFamily::Burstable
159                | InstanceFamily::GeneralPurpose
160                | InstanceFamily::ComputeOptimized
161                | InstanceFamily::MemoryOptimized
162        )
163    }
164
165    /// Whether this instance type supports nested virtualization.
166    ///
167    /// GCP still requires the instance template to opt in; Azure exposes the
168    /// capability automatically on supported VM sizes.
169    pub fn is_nested_virt_capable(&self) -> bool {
170        self.nested_virt
171    }
172
173    /// Convert this catalog entry into a `MachineProfile` for use in `CapacityGroup`.
174    pub fn to_machine_profile(&self) -> MachineProfile {
175        MachineProfile {
176            cpu: format!("{}.0", self.vcpu),
177            memory_bytes: self.memory_bytes,
178            ephemeral_storage_bytes: self.ephemeral_storage_bytes,
179            architecture: Some(self.architecture),
180            gpu: self.gpu.map(|g| GpuSpec {
181                gpu_type: g.gpu_type.to_string(),
182                count: g.count,
183            }),
184            data_disk: None,
185        }
186    }
187
188    /// Convert this entry to the profile controllers should provision for a
189    /// workload. Configurable cloud disks grow to the requested capacity;
190    /// fixed local disks retain their catalog capacity.
191    pub fn to_machine_profile_for_storage(&self, requested_bytes: u64) -> MachineProfile {
192        let mut profile = self.to_machine_profile();
193        if self.has_configurable_ephemeral_storage() {
194            profile.ephemeral_storage_bytes = profile.ephemeral_storage_bytes.max(requested_bytes);
195        }
196        profile
197    }
198}
199
200// Helpers for readable byte constants
201const KI: u64 = 1024;
202const MI: u64 = KI * 1024;
203const GI: u64 = MI * 1024;
204const TI: u64 = GI * 1024;
205
206/// Maximum workload storage that the current cloud controllers can turn into
207/// their provider-backed root disk after adding 25% filesystem headroom and
208/// 12 GiB of runtime overhead. Keep these bounds in sync with the provider
209/// controller contract; rejecting here is preferable to failing after cloud
210/// resources have already been created.
211pub fn max_configurable_ephemeral_storage_bytes(platform: Platform) -> Option<u64> {
212    let disk_limit_gib = match platform {
213        // AWS gp3 and GCP persistent disks currently top out at 64 TiB.
214        Platform::Aws | Platform::Gcp => 64 * 1024,
215        // Azure managed OS disks currently top out at 4,095 GiB.
216        Platform::Azure => 4_095,
217        Platform::Kubernetes | Platform::Machines | Platform::Local | Platform::Test => {
218            return None;
219        }
220    };
221
222    // Controllers first round the requested bytes up to a whole GiB, then
223    // round the 25% headroom up again. Returning a fractional-GiB bound would
224    // therefore admit `bound + 1 byte` and materialize a disk one GiB over the
225    // provider limit.
226    let max_requested_gib = (disk_limit_gib - 12) * 4 / 5;
227    Some(max_requested_gib * GI)
228}
229
230/// The instance type catalog, generated from the clouds' machine APIs by
231/// `scripts/instance-catalog/generate.py` from `instance-catalog.toml`.
232mod catalog;
233use catalog::CATALOG;
234
235// ---------------------------------------------------------------------------
236// Catalog lookup
237// ---------------------------------------------------------------------------
238
239/// Get all instance types for a given platform.
240pub fn catalog_for_platform(platform: Platform) -> Vec<&'static InstanceTypeSpec> {
241    CATALOG
242        .iter()
243        .filter(|spec| spec.platform == platform)
244        .collect()
245}
246
247/// Find a specific instance type by name and platform.
248pub fn find_instance_type(platform: Platform, name: &str) -> Option<&'static InstanceTypeSpec> {
249    CATALOG
250        .iter()
251        .find(|spec| spec.platform == platform && spec.name == name)
252}
253
254/// Whether a capacity group may move from AWS machine `old` to `new` without setup: both are
255/// catalog machines of one CPU architecture, so the stack's images still run on the new one.
256pub fn is_same_architecture_aws_machine(old: &str, new: &str) -> bool {
257    is_same_architecture_machine(Platform::Aws, old, new)
258}
259
260/// Whether two catalog machines belong to the given cloud and share a CPU architecture.
261pub fn is_same_architecture_machine(platform: Platform, old: &str, new: &str) -> bool {
262    matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure)
263        && find_instance_type(platform, old)
264            .zip(find_instance_type(platform, new))
265            .is_some_and(|(old, new)| old.architecture == new.architecture)
266}
267
268// ---------------------------------------------------------------------------
269// Instance type selection
270// ---------------------------------------------------------------------------
271
272/// Aggregated resource requirements from all containers in a capacity group.
273#[derive(Debug, Clone)]
274pub struct WorkloadRequirements {
275    /// Total CPU needed at desired scale (sum of desired CPU * desired_replicas per container)
276    pub total_cpu_at_desired: f64,
277    /// Total memory needed at desired scale (sum of desired memory * desired_replicas per container)
278    pub total_memory_bytes_at_desired: u64,
279    /// Total CPU needed at maximum scale (sum of desired CPU * max_replicas per container)
280    pub total_cpu_at_max: f64,
281    /// Total memory needed at maximum scale (sum of desired memory * max_replicas per container)
282    pub total_memory_bytes_at_max: u64,
283    /// Largest CPU request among all individual containers (single replica)
284    pub max_cpu_per_container: f64,
285    /// Largest memory request among all individual containers (single replica)
286    pub max_memory_per_container: u64,
287    /// Maximum ephemeral storage any single container requires
288    pub max_ephemeral_storage_bytes: u64,
289    /// GPU requirement (if any container needs GPU)
290    pub gpu: Option<GpuSpec>,
291    /// Required CPU architecture, when source explicitly constrains it.
292    pub architecture: Option<Architecture>,
293    /// If true, only instance types that expose nested virtualization (VT-x/EPT)
294    /// to guest VMs are eligible. Required by workloads that run QEMU/KVM
295    /// inside a container.
296    pub nested_virt: bool,
297}
298
299/// Result of instance type selection.
300#[derive(Debug, Clone)]
301pub struct InstanceSelection {
302    /// Selected instance type name (e.g., "m7g.2xlarge")
303    pub instance_type: &'static str,
304    /// Machine profile derived from the instance type
305    pub profile: MachineProfile,
306    /// Recommended minimum number of machines
307    pub min_machines: u32,
308    /// Recommended maximum number of machines
309    pub max_machines: u32,
310}
311
312/// Ephemeral storage threshold above which storage-optimized instances are selected.
313const STORAGE_OPTIMIZED_THRESHOLD: u64 = 200 * GI;
314
315/// Maximum number of machines per cluster.
316const MAX_MACHINES_PER_CLUSTER: u32 = 10;
317
318/// Preferred vCPU cap for non-GPU/non-storage workloads. A larger indivisible
319/// allocation or bounded fleet can require a bigger machine.
320const MAX_STANDARD_VCPU: u32 = 8;
321
322/// Runtime CPU reserved for system processes on each managed container machine.
323const SYSTEM_RESERVE_CPU: f64 = 0.5;
324
325/// Runtime planning headroom for total desired/max workload.
326const WORKLOAD_HEADROOM_FACTOR: f64 = 1.15;
327
328/// Select the best instance type for a workload on a given platform.
329///
330/// The algorithm:
331/// 1. GPU workloads: Match by GPU type, find smallest instance with enough GPUs.
332/// 2. Storage-heavy workloads (>200Gi ephemeral): Use storage-optimized instances.
333/// 3. All other workloads: Fit the allocation with system reserve and headroom,
334///    preferring at most 8 vCPUs unless the allocation requires more. Use the
335///    workload's preferred family, with general-purpose fallback when needed.
336///
337/// Returns an error if no suitable instance type is found.
338pub fn select_instance_type(
339    platform: Platform,
340    requirements: &WorkloadRequirements,
341) -> Result<InstanceSelection, String> {
342    // An unset architecture accepts either one. Nested virtualization and GPU
343    // models can exist on only one architecture of a platform (AWS nested
344    // virtualization is x86_64 only), so try the platform default first and
345    // then the other architecture.
346    if requirements.architecture.is_none()
347        && (requirements.nested_virt || requirements.gpu.is_some())
348    {
349        if let Some(default) = default_architecture(platform) {
350            let other = match default {
351                Architecture::Arm64 => Architecture::X86_64,
352                Architecture::X86_64 => Architecture::Arm64,
353            };
354            let with_architecture = |architecture| WorkloadRequirements {
355                architecture: Some(architecture),
356                ..requirements.clone()
357            };
358            let selection = select_instance_type_for_architecture(
359                platform,
360                &with_architecture(default),
361            );
362            if selection.is_ok() {
363                return selection;
364            }
365            // Report both reasons: the default architecture often lacks the
366            // feature entirely, while the other one says what to change.
367            return select_instance_type_for_architecture(platform, &with_architecture(other))
368                .map_err(|other_error| {
369                    let default_error = selection.err().unwrap_or_default();
370                    format!("{default:?}: {default_error}; {other:?}: {other_error}")
371                });
372        }
373    }
374    select_instance_type_for_architecture(platform, requirements)
375}
376
377fn select_instance_type_for_architecture(
378    platform: Platform,
379    requirements: &WorkloadRequirements,
380) -> Result<InstanceSelection, String> {
381    let family = select_family(requirements);
382    let family = if requirements.nested_virt && family == InstanceFamily::Burstable {
383        InstanceFamily::GeneralPurpose
384    } else {
385        family
386    };
387    let selection = select_instance_type_in_family(platform, requirements, family);
388    if selection.is_err() && family == InstanceFamily::Burstable {
389        return select_instance_type_in_family(
390            platform,
391            requirements,
392            InstanceFamily::GeneralPurpose,
393        );
394    }
395    selection
396}
397
398fn select_instance_type_in_family(
399    platform: Platform,
400    requirements: &WorkloadRequirements,
401    family: InstanceFamily,
402) -> Result<InstanceSelection, String> {
403    let architecture = requirements
404        .architecture
405        .or_else(|| default_architecture(platform))
406        .ok_or_else(|| format!("platform {platform} has no default compute architecture"))?;
407
408    let candidates: Vec<&InstanceTypeSpec> = CATALOG
409        .iter()
410        .filter(|spec| spec.auto_select && spec.platform == platform && spec.family == family)
411        .filter(|spec| {
412            if requirements.nested_virt {
413                spec.is_nested_virt_capable()
414            } else {
415                platform != Platform::Aws || !spec.is_nested_virt_capable()
416            }
417        })
418        .filter(|spec| spec.architecture == architecture)
419        .collect();
420
421    // A storage-heavy workload may have no fixed-local candidate for the
422    // requested architecture/nested-virtualization contract. That is exactly
423    // when we must consider a provider-backed disk on a general-purpose VM;
424    // rejecting here made the fallback below unreachable.
425    if candidates.is_empty() && family != InstanceFamily::StorageOptimized {
426        let family_has_other_architecture = CATALOG.iter().any(|spec| {
427            spec.auto_select
428                && spec.platform == platform
429                && spec.family == family
430                && if requirements.nested_virt {
431                    spec.is_nested_virt_capable()
432                } else {
433                    platform != Platform::Aws || !spec.is_nested_virt_capable()
434                }
435        });
436        if family_has_other_architecture {
437            return Err(format!(
438                "architecture {architecture:?} is unavailable for this workload on platform {platform}"
439            ));
440        }
441        return Err(if requirements.nested_virt {
442            format!(
443                "no nested-virt-capable {family:?} instance types in catalog for platform {platform}"
444            )
445        } else {
446            format!("no {family:?} instance types in catalog for platform {platform}")
447        });
448    }
449
450    // For GPU workloads, filter by GPU type
451    let candidates = if let Some(ref gpu) = requirements.gpu {
452        let filtered: Vec<&InstanceTypeSpec> = candidates
453            .into_iter()
454            .filter(|spec| {
455                spec.gpu.as_ref().map_or(false, |g| {
456                    g.gpu_type == gpu.gpu_type && g.count >= gpu.count
457                })
458            })
459            .collect();
460        if filtered.is_empty() {
461            return Err(format!(
462                "no instance type for GPU type '{}' x{} on platform {platform}",
463                gpu.gpu_type, gpu.count
464            ));
465        }
466        filtered
467    } else {
468        candidates
469    };
470
471    // For storage workloads, filter by ephemeral storage capacity
472    let candidates = if family == InstanceFamily::StorageOptimized {
473        let filtered: Vec<&InstanceTypeSpec> = candidates
474            .into_iter()
475            .filter(|spec| spec.ephemeral_storage_bytes >= requirements.max_ephemeral_storage_bytes)
476            .collect();
477        if filtered.is_empty() {
478            // Fixed-local NVMe is preferred for large requests, but provider
479            // disks on ordinary cloud machines remain configurable. Fall back
480            // when no fixed-local catalog entry can satisfy the request.
481            CATALOG
482                .iter()
483                .filter(|spec| {
484                    spec.auto_select
485                        && spec.platform == platform
486                        && spec.family == InstanceFamily::GeneralPurpose
487                        && spec.has_configurable_ephemeral_storage()
488                        && max_configurable_ephemeral_storage_bytes(platform)
489                            .is_some_and(|max| requirements.max_ephemeral_storage_bytes <= max)
490                })
491                .filter(|spec| {
492                    if requirements.nested_virt {
493                        spec.is_nested_virt_capable()
494                    } else {
495                        platform != Platform::Aws || !spec.is_nested_virt_capable()
496                    }
497                })
498                .filter(|spec| spec.architecture == architecture)
499                .collect()
500        } else {
501            filtered
502        }
503    } else {
504        candidates
505    };
506
507    if candidates.is_empty() {
508        return Err(format!(
509            "architecture {architecture:?} is unavailable for this workload on platform {platform}"
510        ));
511    }
512
513    // More machines cannot split one replica. Every candidate must fit its
514    // largest allocation after host reserve and workload headroom.
515    let candidates: Vec<_> = candidates
516        .into_iter()
517        .filter(|spec| {
518            allocatable_cpu(spec) >= requirements.max_cpu_per_container * WORKLOAD_HEADROOM_FACTOR
519                && allocatable_memory_bytes(spec) as f64
520                    >= requirements.max_memory_per_container as f64 * WORKLOAD_HEADROOM_FACTOR
521        })
522        .collect();
523    let smallest_replica_machine = candidates.iter().map(|spec| spec.vcpu).min()
524        .ok_or_else(|| format!("no {platform} machine can fit one replica with system reserve and workload headroom"))?;
525
526    // Apply the policy for the candidates we will actually select from. A
527    // storage-heavy request can fall back from fixed-local storage machines to
528    // general-purpose machines with provider-backed disks; those machines must
529    // retain the normal horizontal-scaling cap.
530    let effective_family = candidates[0].family;
531    let vcpu_cap = if effective_family == InstanceFamily::GpuCompute
532        || effective_family == InstanceFamily::StorageOptimized
533    {
534        u32::MAX
535    } else {
536        MAX_STANDARD_VCPU.max(smallest_replica_machine)
537    };
538
539    let desired_target_machines = desired_target_machines(requirements);
540    let target_cpu = requirements
541        .max_cpu_per_container
542        .max(requirements.total_cpu_at_desired / desired_target_machines as f64)
543        * WORKLOAD_HEADROOM_FACTOR;
544    let target_memory = (requirements.max_memory_per_container as f64)
545        .max(requirements.total_memory_bytes_at_desired as f64 / desired_target_machines as f64)
546        * WORKLOAD_HEADROOM_FACTOR;
547
548    // Find the smallest instance whose allocatable capacity meets the workload
549    // target after host reserve and workload headroom. Machine count already
550    // accounts for multiple replicas; requiring space for an arbitrary second
551    // copy here would size the same demand twice.
552    let selected = candidates
553        .iter()
554        .filter(|spec| {
555            spec.vcpu <= vcpu_cap
556                && allocatable_cpu(spec) >= target_cpu
557                && allocatable_memory_bytes(spec) as f64 >= target_memory
558        })
559        .min_by_key(|spec| spec.vcpu)
560        .or_else(|| {
561            // If nothing fits within the cap, pick the largest instance under the cap
562            candidates
563                .iter()
564                .filter(|spec| spec.vcpu <= vcpu_cap)
565                .max_by_key(|spec| spec.vcpu)
566        })
567        .or_else(|| {
568            // Last resort: pick the smallest available instance (for GPU/storage)
569            candidates.iter().min_by_key(|spec| spec.vcpu)
570        })
571        .ok_or_else(|| format!("no instance types available for platform {platform}"))?;
572
573    // Calculate machine counts
574    let max_machines = compute_max_machines(requirements, selected);
575    let min_machines = compute_min_machines(requirements, selected, max_machines);
576
577    Ok(InstanceSelection {
578        instance_type: selected.name,
579        profile: selected.to_machine_profile_for_storage(requirements.max_ephemeral_storage_bytes),
580        min_machines,
581        max_machines,
582    })
583}
584
585/// Select instance family based on workload characteristics.
586///
587/// Uses GeneralPurpose for all standard workloads — widely available across
588/// regions and cost-effective. Only specialized workloads (GPU, large ephemeral
589/// storage) get specialized families. Very small workloads get burstable.
590pub fn select_family(requirements: &WorkloadRequirements) -> InstanceFamily {
591    // GPU workloads always get GPU instances
592    if requirements.gpu.is_some() {
593        return InstanceFamily::GpuCompute;
594    }
595
596    // Large ephemeral storage needs NVMe (storage-optimized)
597    if requirements.max_ephemeral_storage_bytes > STORAGE_OPTIMIZED_THRESHOLD {
598        return InstanceFamily::StorageOptimized;
599    }
600
601    // Very small workloads use burstable instances
602    if requirements.total_cpu_at_max < 2.0 {
603        return InstanceFamily::Burstable;
604    }
605
606    // All other workloads use GeneralPurpose — available everywhere, good pricing
607    InstanceFamily::GeneralPurpose
608}
609
610/// Calculate maximum machines needed to fit the workload with headroom.
611fn compute_max_machines(requirements: &WorkloadRequirements, instance: &InstanceTypeSpec) -> u32 {
612    let cpu_with_headroom = requirements.total_cpu_at_max * WORKLOAD_HEADROOM_FACTOR;
613    let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
614
615    let mem_with_headroom =
616        requirements.total_memory_bytes_at_max as f64 * WORKLOAD_HEADROOM_FACTOR;
617    let mem_machines =
618        (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
619
620    // Take the larger of CPU-based and memory-based, clamped to cluster limit
621    cpu_machines
622        .max(mem_machines)
623        .max(1)
624        .min(MAX_MACHINES_PER_CLUSTER)
625}
626
627/// Calculate minimum machines for HA.
628fn compute_min_machines(
629    requirements: &WorkloadRequirements,
630    instance: &InstanceTypeSpec,
631    max_machines: u32,
632) -> u32 {
633    let cpu_with_headroom = requirements.total_cpu_at_desired * WORKLOAD_HEADROOM_FACTOR;
634    let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
635
636    let mem_with_headroom =
637        requirements.total_memory_bytes_at_desired as f64 * WORKLOAD_HEADROOM_FACTOR;
638    let mem_machines =
639        (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
640
641    cpu_machines
642        .max(mem_machines)
643        .max(1)
644        .min(2)
645        .min(max_machines)
646}
647
648fn desired_target_machines(requirements: &WorkloadRequirements) -> u32 {
649    if requirements.total_cpu_at_desired >= 2.0
650        || requirements.total_memory_bytes_at_desired >= 4 * GI
651    {
652        2
653    } else {
654        1
655    }
656}
657
658pub(crate) fn allocatable_cpu(instance: &InstanceTypeSpec) -> f64 {
659    (instance.vcpu as f64 - SYSTEM_RESERVE_CPU).max(0.25)
660}
661
662pub(crate) fn allocatable_memory_bytes(instance: &InstanceTypeSpec) -> u64 {
663    instance
664        .memory_bytes
665        .saturating_sub(system_reserve_memory_bytes(instance.memory_bytes))
666        .max(256 * MI)
667}
668
669fn system_reserve_memory_bytes(memory_bytes: u64) -> u64 {
670    if memory_bytes < 4 * GI {
671        256 * MI
672    } else if memory_bytes < 16 * GI {
673        512 * MI
674    } else {
675        GI
676    }
677}
678
679// ---------------------------------------------------------------------------
680// Tests
681// ---------------------------------------------------------------------------
682
683#[cfg(test)]
684mod tests {
685    use super::*;
686    use crate::BinaryTarget;
687
688    // -- Parsing tests --
689
690    #[test]
691    fn test_parse_cpu_plain() {
692        assert_eq!(parse_cpu("1").unwrap(), 1.0);
693        assert_eq!(parse_cpu("0.5").unwrap(), 0.5);
694        assert_eq!(parse_cpu("2.0").unwrap(), 2.0);
695        assert_eq!(parse_cpu("16").unwrap(), 16.0);
696    }
697
698    #[test]
699    fn test_parse_cpu_millicore() {
700        assert_eq!(parse_cpu("500m").unwrap(), 0.5);
701        assert_eq!(parse_cpu("250m").unwrap(), 0.25);
702        assert_eq!(parse_cpu("1000m").unwrap(), 1.0);
703        assert_eq!(parse_cpu("100m").unwrap(), 0.1);
704    }
705
706    #[test]
707    fn test_parse_cpu_invalid() {
708        assert!(parse_cpu("").is_err());
709        assert!(parse_cpu("abc").is_err());
710        assert!(parse_cpu("m").is_err());
711    }
712
713    #[test]
714    fn test_parse_memory_binary_suffixes() {
715        assert_eq!(parse_memory_bytes("1Ki").unwrap(), 1024);
716        assert_eq!(parse_memory_bytes("1Mi").unwrap(), 1024 * 1024);
717        assert_eq!(parse_memory_bytes("1Gi").unwrap(), 1024 * 1024 * 1024);
718        assert_eq!(parse_memory_bytes("4Gi").unwrap(), 4 * 1024 * 1024 * 1024);
719        assert_eq!(parse_memory_bytes("512Mi").unwrap(), 512 * 1024 * 1024);
720        assert_eq!(
721            parse_memory_bytes("1Ti").unwrap(),
722            1024u64 * 1024 * 1024 * 1024
723        );
724    }
725
726    #[test]
727    fn test_parse_memory_decimal_suffixes() {
728        assert_eq!(parse_memory_bytes("1k").unwrap(), 1000);
729        assert_eq!(parse_memory_bytes("1M").unwrap(), 1_000_000);
730        assert_eq!(parse_memory_bytes("1G").unwrap(), 1_000_000_000);
731        assert_eq!(parse_memory_bytes("1T").unwrap(), 1_000_000_000_000);
732    }
733
734    #[test]
735    fn test_parse_memory_plain_bytes() {
736        assert_eq!(parse_memory_bytes("1024").unwrap(), 1024);
737        assert_eq!(parse_memory_bytes("0").unwrap(), 0);
738    }
739
740    #[test]
741    fn test_parse_memory_invalid() {
742        assert!(parse_memory_bytes("").is_err());
743        assert!(parse_memory_bytes("abc").is_err());
744        assert!(parse_memory_bytes("Gi").is_err());
745    }
746
747    #[test]
748    fn test_parse_memory_fractional() {
749        assert_eq!(parse_memory_bytes("0.5Gi").unwrap(), GI / 2);
750        assert_eq!(parse_memory_bytes("1.5Gi").unwrap(), GI + GI / 2);
751    }
752
753    // -- Catalog lookup tests --
754
755    #[test]
756    fn test_catalog_has_entries_for_all_cloud_platforms() {
757        assert!(!catalog_for_platform(Platform::Aws).is_empty());
758        assert!(!catalog_for_platform(Platform::Gcp).is_empty());
759        assert!(!catalog_for_platform(Platform::Azure).is_empty());
760    }
761
762    #[test]
763    fn test_catalog_no_entries_for_non_cloud_platforms() {
764        assert!(catalog_for_platform(Platform::Local).is_empty());
765        assert!(catalog_for_platform(Platform::Kubernetes).is_empty());
766    }
767
768    #[test]
769    fn test_find_known_instance_type() {
770        let spec =
771            find_instance_type(Platform::Aws, "m7g.2xlarge").expect("should find m7g.2xlarge");
772        assert_eq!(spec.vcpu, 8);
773        assert_eq!(spec.memory_bytes, 32 * GI);
774        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
775    }
776
777    #[test]
778    fn test_find_aws_c8i_nested_virt_instance_type() {
779        let spec = find_instance_type(Platform::Aws, "c8i.large").expect("should find c8i.large");
780        assert_eq!(spec.vcpu, 2);
781        assert_eq!(spec.memory_bytes, 4 * GI);
782        assert_eq!(spec.family, InstanceFamily::ComputeOptimized);
783        assert_eq!(spec.architecture, Architecture::X86_64);
784        assert!(spec.is_nested_virt_capable());
785    }
786
787    #[test]
788    fn test_find_unknown_instance_type() {
789        assert!(find_instance_type(Platform::Aws, "nonexistent.xlarge").is_none());
790    }
791
792    #[test]
793    fn test_find_wrong_platform() {
794        assert!(find_instance_type(Platform::Gcp, "m7g.2xlarge").is_none());
795    }
796
797    #[test]
798    fn test_to_machine_profile() {
799        let spec = find_instance_type(Platform::Aws, "m7g.2xlarge").unwrap();
800        let profile = spec.to_machine_profile();
801        assert_eq!(profile.cpu, "8.0");
802        assert_eq!(profile.memory_bytes, 32 * GI);
803        assert_eq!(profile.ephemeral_storage_bytes, 20 * GI);
804        assert!(profile.gpu.is_none());
805    }
806
807    #[test]
808    fn test_to_machine_profile_with_gpu() {
809        let spec = find_instance_type(Platform::Aws, "p4d.24xlarge").unwrap();
810        let profile = spec.to_machine_profile();
811        let gpu = profile.gpu.as_ref().expect("should have GPU");
812        assert_eq!(gpu.gpu_type, "nvidia-a100");
813        assert_eq!(gpu.count, 8);
814    }
815
816    // -- Selection algorithm tests --
817
818    #[test]
819    fn small_cpu_with_large_memory_falls_back_from_burstable() {
820        let requirements = WorkloadRequirements {
821            total_cpu_at_desired: 0.5,
822            total_memory_bytes_at_desired: 8 * GI,
823            total_cpu_at_max: 0.5,
824            total_memory_bytes_at_max: 8 * GI,
825            max_cpu_per_container: 0.5,
826            max_memory_per_container: 8 * GI,
827            max_ephemeral_storage_bytes: 0,
828            gpu: None,
829            architecture: Some(Architecture::X86_64),
830            nested_virt: false,
831        };
832        let selected = select_instance_type(Platform::Gcp, &requirements).unwrap();
833        let machine = find_instance_type(Platform::Gcp, selected.instance_type).unwrap();
834        assert!(
835            allocatable_memory_bytes(machine) as f64 >= 8.0 * GI as f64 * WORKLOAD_HEADROOM_FACTOR
836        );
837    }
838
839    #[test]
840    fn test_select_burstable_for_small_workload() {
841        let req = WorkloadRequirements {
842            total_cpu_at_desired: 1.0,
843            total_memory_bytes_at_desired: 2 * GI,
844            total_cpu_at_max: 1.0,
845            total_memory_bytes_at_max: 2 * GI,
846            max_cpu_per_container: 0.5,
847            max_memory_per_container: 1 * GI,
848            max_ephemeral_storage_bytes: 10 * GI,
849            gpu: None,
850            architecture: None,
851            nested_virt: false,
852        };
853        let sel = select_instance_type(Platform::Aws, &req).unwrap();
854        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
855        assert_eq!(spec.family, InstanceFamily::Burstable);
856    }
857
858    #[test]
859    fn test_selects_smallest_burstable_machine_with_real_headroom() {
860        let req = WorkloadRequirements {
861            total_cpu_at_desired: 1.0,
862            total_memory_bytes_at_desired: 2 * GI,
863            total_cpu_at_max: 1.0,
864            total_memory_bytes_at_max: 2 * GI,
865            max_cpu_per_container: 1.0,
866            max_memory_per_container: 2 * GI,
867            max_ephemeral_storage_bytes: 10 * GI,
868            gpu: None,
869            architecture: None,
870            nested_virt: false,
871        };
872
873        let selection = select_instance_type(Platform::Aws, &req).unwrap();
874
875        assert_eq!(selection.instance_type, "t4g.medium");
876        assert_eq!(selection.min_machines, 1);
877        assert_eq!(selection.max_machines, 1);
878    }
879
880    #[test]
881    fn test_select_general_purpose_for_standard_workload() {
882        // Standard workloads always get GeneralPurpose regardless of CPU:memory ratio
883        let req = WorkloadRequirements {
884            total_cpu_at_desired: 20.0,
885            total_memory_bytes_at_desired: 80 * GI,
886            total_cpu_at_max: 20.0,
887            total_memory_bytes_at_max: 80 * GI,
888            max_cpu_per_container: 2.0,
889            max_memory_per_container: 8 * GI,
890            max_ephemeral_storage_bytes: 10 * GI,
891            gpu: None,
892            architecture: None,
893            nested_virt: false,
894        };
895        let sel = select_instance_type(Platform::Aws, &req).unwrap();
896        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
897        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
898    }
899
900    #[test]
901    fn test_select_general_purpose_even_for_cpu_heavy() {
902        // CPU-heavy workloads still get GeneralPurpose (no more ComputeOptimized auto-select)
903        let req = WorkloadRequirements {
904            total_cpu_at_desired: 20.0,
905            total_memory_bytes_at_desired: 20 * GI,
906            total_cpu_at_max: 20.0,
907            total_memory_bytes_at_max: 20 * GI,
908            max_cpu_per_container: 2.0,
909            max_memory_per_container: 2 * GI,
910            max_ephemeral_storage_bytes: 10 * GI,
911            gpu: None,
912            architecture: None,
913            nested_virt: false,
914        };
915        let sel = select_instance_type(Platform::Aws, &req).unwrap();
916        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
917        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
918    }
919
920    #[test]
921    fn test_select_storage_optimized_for_large_ephemeral() {
922        let req = WorkloadRequirements {
923            total_cpu_at_desired: 8.0,
924            total_memory_bytes_at_desired: 32 * GI,
925            total_cpu_at_max: 8.0,
926            total_memory_bytes_at_max: 32 * GI,
927            max_cpu_per_container: 2.0,
928            max_memory_per_container: 8 * GI,
929            max_ephemeral_storage_bytes: 500 * GI,
930            gpu: None,
931            architecture: Some(Architecture::X86_64),
932            nested_virt: false,
933        };
934        let sel = select_instance_type(Platform::Aws, &req).unwrap();
935        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
936        assert_eq!(spec.family, InstanceFamily::StorageOptimized);
937    }
938
939    #[test]
940    fn test_select_configurable_storage_above_fixed_local_catalog() {
941        let req = WorkloadRequirements {
942            total_cpu_at_desired: 8.0,
943            total_memory_bytes_at_desired: 32 * GI,
944            total_cpu_at_max: 8.0,
945            total_memory_bytes_at_max: 32 * GI,
946            max_cpu_per_container: 2.0,
947            max_memory_per_container: 8 * GI,
948            max_ephemeral_storage_bytes: 8_000 * GI,
949            gpu: None,
950            architecture: Some(Architecture::X86_64),
951            nested_virt: false,
952        };
953        let sel = select_instance_type(Platform::Aws, &req).unwrap();
954        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
955        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
956        assert_eq!(sel.profile.ephemeral_storage_bytes, 8_000 * GI);
957    }
958
959    #[test]
960    fn test_configurable_storage_fallback_retains_standard_vcpu_cap() {
961        let req = WorkloadRequirements {
962            total_cpu_at_desired: 70.0,
963            total_memory_bytes_at_desired: 140 * GI,
964            total_cpu_at_max: 70.0,
965            total_memory_bytes_at_max: 140 * GI,
966            max_cpu_per_container: 2.0,
967            max_memory_per_container: 4 * GI,
968            max_ephemeral_storage_bytes: 8_000 * GI,
969            gpu: None,
970            architecture: Some(Architecture::X86_64),
971            nested_virt: false,
972        };
973
974        let sel = select_instance_type(Platform::Aws, &req).unwrap();
975        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
976        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
977        assert!(spec.vcpu <= MAX_STANDARD_VCPU);
978        assert!(sel.max_machines > 1);
979    }
980
981    #[test]
982    fn test_nested_virtualization_storage_falls_back_to_configurable_disk() {
983        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
984            let req = WorkloadRequirements {
985                total_cpu_at_desired: 2.0,
986                total_memory_bytes_at_desired: 8 * GI,
987                total_cpu_at_max: 2.0,
988                total_memory_bytes_at_max: 8 * GI,
989                max_cpu_per_container: 2.0,
990                max_memory_per_container: 8 * GI,
991                max_ephemeral_storage_bytes: 500 * GI,
992                gpu: None,
993                architecture: Some(Architecture::X86_64),
994                nested_virt: true,
995            };
996
997            let sel = select_instance_type(platform, &req)
998                .unwrap_or_else(|error| panic!("{platform} should support fallback: {error}"));
999            let spec = find_instance_type(platform, sel.instance_type).unwrap();
1000            assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1001            assert!(spec.is_nested_virt_capable());
1002            assert_eq!(sel.profile.ephemeral_storage_bytes, 500 * GI);
1003        }
1004    }
1005
1006    #[test]
1007    fn unset_architecture_resolves_to_the_one_offering_nested_virtualization() {
1008        let req = WorkloadRequirements {
1009            total_cpu_at_desired: 2.0,
1010            total_memory_bytes_at_desired: 4 * GI,
1011            total_cpu_at_max: 2.0,
1012            total_memory_bytes_at_max: 4 * GI,
1013            max_cpu_per_container: 2.0,
1014            max_memory_per_container: 4 * GI,
1015            max_ephemeral_storage_bytes: 20 * GI,
1016            gpu: None,
1017            architecture: None,
1018            nested_virt: true,
1019        };
1020
1021        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1022            let selection = select_instance_type(platform, &req)
1023                .unwrap_or_else(|message| panic!("{platform}: {message}"));
1024            let spec = find_instance_type(platform, selection.instance_type)
1025                .expect("selected machine should exist in the catalog");
1026            assert!(spec.is_nested_virt_capable(), "{platform}: {}", spec.name);
1027            assert_eq!(spec.architecture, Architecture::X86_64, "{platform}");
1028            assert_eq!(selection.profile.architecture, Some(Architecture::X86_64));
1029        }
1030
1031        // When neither architecture fits, the error says why for each, so a
1032        // sizing problem is not reported as a missing architecture.
1033        let oversized = WorkloadRequirements {
1034            max_cpu_per_container: 1000.0,
1035            total_cpu_at_desired: 1000.0,
1036            total_cpu_at_max: 1000.0,
1037            ..req.clone()
1038        };
1039        let message = select_instance_type(Platform::Aws, &oversized)
1040            .expect_err("no AWS machine fits 1000 vCPU per replica");
1041        assert!(message.starts_with("Arm64: architecture Arm64 is unavailable"), "{message}");
1042        assert!(
1043            message.contains("; X86_64: no aws machine can fit one replica"),
1044            "{message}"
1045        );
1046
1047        // An explicit architecture is a constraint, not a preference.
1048        let arm = WorkloadRequirements {
1049            architecture: Some(Architecture::Arm64),
1050            ..req
1051        };
1052        assert!(select_instance_type(Platform::Aws, &arm).is_err());
1053    }
1054
1055    #[test]
1056    fn test_configurable_disk_rejects_capacity_above_provider_limit() {
1057        for (platform, requested) in [
1058            (Platform::Aws, 100 * TI),
1059            (Platform::Gcp, 100 * TI),
1060            (Platform::Azure, 8_000 * GI),
1061        ] {
1062            let req = WorkloadRequirements {
1063                total_cpu_at_desired: 2.0,
1064                total_memory_bytes_at_desired: 8 * GI,
1065                total_cpu_at_max: 2.0,
1066                total_memory_bytes_at_max: 8 * GI,
1067                max_cpu_per_container: 2.0,
1068                max_memory_per_container: 8 * GI,
1069                max_ephemeral_storage_bytes: requested,
1070                gpu: None,
1071                architecture: Some(Architecture::X86_64),
1072                nested_virt: true,
1073            };
1074
1075            assert!(select_instance_type(platform, &req).is_err());
1076        }
1077    }
1078
1079    #[test]
1080    fn test_configurable_disk_limits_account_for_controller_rounding() {
1081        for (platform, disk_limit_gib) in [
1082            (Platform::Aws, 64 * 1024),
1083            (Platform::Gcp, 64 * 1024),
1084            (Platform::Azure, 4_095),
1085        ] {
1086            let max = max_configurable_ephemeral_storage_bytes(platform).unwrap();
1087            let materialized_disk_gib = |requested_bytes: u64| {
1088                let requested_gib = requested_bytes.div_ceil(GI);
1089                (requested_gib * 5).div_ceil(4) + 12
1090            };
1091
1092            assert!(materialized_disk_gib(max) <= disk_limit_gib);
1093            assert!(materialized_disk_gib(max + 1) > disk_limit_gib);
1094        }
1095    }
1096
1097    #[test]
1098    fn test_gpu_local_storage_is_not_treated_as_resizable() {
1099        let spec = find_instance_type(Platform::Aws, "g5.xlarge").unwrap();
1100        assert_eq!(spec.family, InstanceFamily::GpuCompute);
1101        assert!(!spec.has_configurable_ephemeral_storage());
1102    }
1103
1104    #[test]
1105    fn test_select_gpu_instance() {
1106        let req = WorkloadRequirements {
1107            total_cpu_at_desired: 8.0,
1108            total_memory_bytes_at_desired: 32 * GI,
1109            total_cpu_at_max: 8.0,
1110            total_memory_bytes_at_max: 32 * GI,
1111            max_cpu_per_container: 4.0,
1112            max_memory_per_container: 16 * GI,
1113            max_ephemeral_storage_bytes: 10 * GI,
1114            gpu: Some(GpuSpec {
1115                gpu_type: "nvidia-a100".to_string(),
1116                count: 1,
1117            }),
1118            architecture: Some(Architecture::X86_64),
1119            nested_virt: false,
1120        };
1121        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1122        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1123        assert_eq!(spec.family, InstanceFamily::GpuCompute);
1124        assert!(spec.gpu.is_some());
1125    }
1126
1127    #[test]
1128    fn test_select_uses_each_cloud_image_target_architecture() {
1129        let req = WorkloadRequirements {
1130            total_cpu_at_desired: 4.0,
1131            total_memory_bytes_at_desired: 16 * GI,
1132            total_cpu_at_max: 4.0,
1133            total_memory_bytes_at_max: 16 * GI,
1134            max_cpu_per_container: 1.0,
1135            max_memory_per_container: 4 * GI,
1136            max_ephemeral_storage_bytes: 10 * GI,
1137            gpu: None,
1138            architecture: None,
1139            nested_virt: false,
1140        };
1141        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1142            let sel = select_instance_type(platform, &req)
1143                .unwrap_or_else(|error| panic!("selection failed for {platform}: {error}"));
1144            let spec = find_instance_type(platform, sel.instance_type)
1145                .expect("selected machine should exist in the catalog");
1146            assert_eq!(
1147                Some(spec.architecture),
1148                default_architecture(platform),
1149                "machine architecture must match the image target for {platform}"
1150            );
1151        }
1152    }
1153
1154    #[test]
1155    fn test_machine_count_reasonable() {
1156        // Single container: 1 CPU, 2Gi, maxReplicas=20
1157        let req = WorkloadRequirements {
1158            total_cpu_at_desired: 20.0,
1159            total_memory_bytes_at_desired: 40 * GI,
1160            total_cpu_at_max: 20.0,
1161            total_memory_bytes_at_max: 40 * GI,
1162            max_cpu_per_container: 1.0,
1163            max_memory_per_container: 2 * GI,
1164            max_ephemeral_storage_bytes: 10 * GI,
1165            gpu: None,
1166            architecture: None,
1167            nested_virt: false,
1168        };
1169        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1170        assert!(sel.min_machines >= 1);
1171        assert!(sel.max_machines <= MAX_MACHINES_PER_CLUSTER);
1172        assert!(sel.max_machines >= sel.min_machines);
1173    }
1174
1175    #[test]
1176    fn test_instance_size_capped_at_8_vcpu() {
1177        // Many small replicas prefer more machines over larger instances.
1178        let req = WorkloadRequirements {
1179            total_cpu_at_desired: 70.0,
1180            total_memory_bytes_at_desired: 140 * GI,
1181            total_cpu_at_max: 70.0,
1182            total_memory_bytes_at_max: 140 * GI,
1183            max_cpu_per_container: 2.0,
1184            max_memory_per_container: 4 * GI,
1185            max_ephemeral_storage_bytes: 10 * GI,
1186            gpu: None,
1187            architecture: None,
1188            nested_virt: false,
1189        };
1190        let sel = select_instance_type(Platform::Gcp, &req).unwrap();
1191        let spec = find_instance_type(Platform::Gcp, sel.instance_type).unwrap();
1192        assert!(
1193            spec.vcpu <= MAX_STANDARD_VCPU,
1194            "selected {} with {} vCPUs, expected <= {}",
1195            spec.name,
1196            spec.vcpu,
1197            MAX_STANDARD_VCPU
1198        );
1199        assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1200        // Should scale horizontally instead
1201        assert!(sel.max_machines > 1);
1202    }
1203
1204    #[test]
1205    fn test_larger_autoscaled_workload_gets_reasonable_instance() {
1206        // Simulates a larger autoscaled workload: 4 containers, each 2 CPU / 4 GiB
1207        // maxReplicas: 10, 10, 10, 5
1208        let req = WorkloadRequirements {
1209            total_cpu_at_desired: 70.0,
1210            total_memory_bytes_at_desired: 140 * GI,
1211            total_cpu_at_max: 70.0,              // 2*10 + 2*10 + 2*10 + 2*5
1212            total_memory_bytes_at_max: 140 * GI, // 4*10 + 4*10 + 4*10 + 4*5
1213            max_cpu_per_container: 2.0,
1214            max_memory_per_container: 4 * GI,
1215            max_ephemeral_storage_bytes: 20 * GI,
1216            gpu: None,
1217            architecture: None,
1218            nested_virt: false,
1219        };
1220        let sel = select_instance_type(Platform::Gcp, &req).unwrap();
1221        // Should pick n2-standard-8 (8 vCPU, 32 GiB) — NOT c3-standard-44
1222        assert_eq!(sel.instance_type, "n2-standard-8");
1223        assert!(sel.max_machines >= 2);
1224    }
1225
1226    /// When `nested_virt` is set on the workload, the selector must
1227    /// restrict to nested-virt-capable families. On AWS that means an m8i
1228    /// (or other 8th-gen Intel) entry, never a Graviton (`*7g`, `t4g`) or
1229    /// burstable. Without this filter the launch template gets created
1230    /// with `CpuOptions.NestedVirtualization=enabled` paired with an
1231    /// instance type AWS rejects at RunInstances.
1232    #[test]
1233    fn test_select_aws_picks_m8i_when_nested_virt_required() {
1234        let req = WorkloadRequirements {
1235            total_cpu_at_desired: 4.0,
1236            total_memory_bytes_at_desired: 8 * GI,
1237            total_cpu_at_max: 4.0,
1238            total_memory_bytes_at_max: 8 * GI,
1239            max_cpu_per_container: 4.0,
1240            max_memory_per_container: 8 * GI,
1241            max_ephemeral_storage_bytes: 10 * GI,
1242            gpu: None,
1243            architecture: Some(Architecture::X86_64),
1244            nested_virt: true,
1245        };
1246        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1247        assert!(
1248            sel.instance_type.starts_with("m8i.")
1249                || sel.instance_type.starts_with("c8i.")
1250                || sel.instance_type.starts_with("r8i."),
1251            "expected an m8i/c8i/r8i instance, got {}",
1252            sel.instance_type
1253        );
1254        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1255        assert!(spec.is_nested_virt_capable());
1256    }
1257
1258    #[test]
1259    fn test_select_gcp_picks_n2_when_nested_virt_required() {
1260        let req = WorkloadRequirements {
1261            total_cpu_at_desired: 4.0,
1262            total_memory_bytes_at_desired: 8 * GI,
1263            total_cpu_at_max: 4.0,
1264            total_memory_bytes_at_max: 8 * GI,
1265            max_cpu_per_container: 4.0,
1266            max_memory_per_container: 8 * GI,
1267            max_ephemeral_storage_bytes: 0,
1268            architecture: Some(Architecture::X86_64),
1269            gpu: None,
1270            nested_virt: true,
1271        };
1272
1273        let selection = select_instance_type(Platform::Gcp, &req).unwrap();
1274        assert_eq!(selection.instance_type, "n2-standard-8");
1275        assert!(find_instance_type(Platform::Gcp, selection.instance_type)
1276            .unwrap()
1277            .is_nested_virt_capable());
1278    }
1279
1280    #[test]
1281    fn test_select_azure_picks_dsv5_when_nested_virt_required() {
1282        let req = WorkloadRequirements {
1283            total_cpu_at_desired: 4.0,
1284            total_memory_bytes_at_desired: 8 * GI,
1285            total_cpu_at_max: 4.0,
1286            total_memory_bytes_at_max: 8 * GI,
1287            max_cpu_per_container: 4.0,
1288            max_memory_per_container: 8 * GI,
1289            max_ephemeral_storage_bytes: 0,
1290            architecture: Some(Architecture::X86_64),
1291            gpu: None,
1292            nested_virt: true,
1293        };
1294
1295        let selection = select_instance_type(Platform::Azure, &req).unwrap();
1296        assert_eq!(selection.instance_type, "Standard_D8s_v5");
1297        assert!(find_instance_type(Platform::Azure, selection.instance_type)
1298            .unwrap()
1299            .is_nested_virt_capable());
1300    }
1301
1302    #[test]
1303    fn test_select_aws_defaults_to_image_target_architecture() {
1304        let req = WorkloadRequirements {
1305            total_cpu_at_desired: 4.0,
1306            total_memory_bytes_at_desired: 8 * GI,
1307            total_cpu_at_max: 4.0,
1308            total_memory_bytes_at_max: 8 * GI,
1309            max_cpu_per_container: 4.0,
1310            max_memory_per_container: 8 * GI,
1311            max_ephemeral_storage_bytes: 10 * GI,
1312            gpu: None,
1313            architecture: None,
1314            nested_virt: false,
1315        };
1316        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1317        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1318        assert_eq!(spec.architecture, Architecture::Arm64);
1319    }
1320
1321    #[test]
1322    fn test_cloud_defaults_match_image_target_architectures() {
1323        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1324            let target = BinaryTarget::defaults_for_platform(platform)
1325                .into_iter()
1326                .next()
1327                .expect("managed cloud should have a default image target");
1328            let image_architecture = match target.oci_arch() {
1329                "arm64" => Architecture::Arm64,
1330                "amd64" => Architecture::X86_64,
1331                architecture => {
1332                    panic!("unsupported managed-cloud image architecture {architecture}")
1333                }
1334            };
1335
1336            assert_eq!(default_architecture(platform), Some(image_architecture));
1337        }
1338    }
1339
1340    /// ARM remains available when the workload or capacity profile declares it.
1341    #[test]
1342    fn test_select_aws_uses_graviton_for_explicit_arm64() {
1343        let req = WorkloadRequirements {
1344            total_cpu_at_desired: 4.0,
1345            total_memory_bytes_at_desired: 8 * GI,
1346            total_cpu_at_max: 4.0,
1347            total_memory_bytes_at_max: 8 * GI,
1348            max_cpu_per_container: 4.0,
1349            max_memory_per_container: 8 * GI,
1350            max_ephemeral_storage_bytes: 10 * GI,
1351            gpu: None,
1352            architecture: Some(Architecture::Arm64),
1353            nested_virt: false,
1354        };
1355        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1356        let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1357        assert_eq!(spec.architecture, Architecture::Arm64);
1358    }
1359
1360    #[test]
1361    fn test_select_rejects_explicit_architecture_missing_from_cloud_catalog() {
1362        let req = WorkloadRequirements {
1363            total_cpu_at_desired: 1.0,
1364            total_memory_bytes_at_desired: 2 * GI,
1365            total_cpu_at_max: 1.0,
1366            total_memory_bytes_at_max: 2 * GI,
1367            max_cpu_per_container: 1.0,
1368            max_memory_per_container: 2 * GI,
1369            max_ephemeral_storage_bytes: 10 * GI,
1370            gpu: None,
1371            architecture: Some(Architecture::Arm64),
1372            nested_virt: false,
1373        };
1374
1375        let error = select_instance_type(Platform::Gcp, &req)
1376            .expect_err("GCP catalog has no ARM64 machine");
1377
1378        assert!(error.contains("architecture Arm64 is unavailable"));
1379    }
1380
1381    #[test]
1382    fn test_profile_has_required_fields() {
1383        let req = WorkloadRequirements {
1384            total_cpu_at_desired: 4.0,
1385            total_memory_bytes_at_desired: 16 * GI,
1386            total_cpu_at_max: 4.0,
1387            total_memory_bytes_at_max: 16 * GI,
1388            max_cpu_per_container: 1.0,
1389            max_memory_per_container: 4 * GI,
1390            max_ephemeral_storage_bytes: 10 * GI,
1391            gpu: None,
1392            architecture: None,
1393            nested_virt: false,
1394        };
1395        let sel = select_instance_type(Platform::Aws, &req).unwrap();
1396        assert!(!sel.profile.cpu.is_empty());
1397        assert!(sel.profile.memory_bytes > 0);
1398        assert!(sel.profile.ephemeral_storage_bytes > 0);
1399    }
1400
1401    #[test]
1402    fn test_error_for_unsupported_gpu_type() {
1403        let req = WorkloadRequirements {
1404            total_cpu_at_desired: 8.0,
1405            total_memory_bytes_at_desired: 32 * GI,
1406            total_cpu_at_max: 8.0,
1407            total_memory_bytes_at_max: 32 * GI,
1408            max_cpu_per_container: 4.0,
1409            max_memory_per_container: 16 * GI,
1410            max_ephemeral_storage_bytes: 10 * GI,
1411            gpu: Some(GpuSpec {
1412                gpu_type: "amd-mi300".to_string(),
1413                count: 1,
1414            }),
1415            architecture: None,
1416            nested_virt: false,
1417        };
1418        let result = select_instance_type(Platform::Aws, &req);
1419        assert!(result.is_err());
1420    }
1421
1422    /// Workload shapes for the auto-selection matrix: name, CPU and memory per replica,
1423    /// desired and max replicas, ephemeral storage, GPU type and count.
1424    type SelectionShape = (
1425        &'static str,
1426        f64,
1427        u64,
1428        u32,
1429        u32,
1430        u64,
1431        Option<(&'static str, u32)>,
1432    );
1433
1434    #[rustfmt::skip]
1435    const SELECTION_SHAPES: &[SelectionShape] = &[
1436        ("tiny", 0.25, 512 * MI, 1, 1, 0, None),
1437        ("burstable", 1.0, 2 * GI, 1, 1, 10 * GI, None),
1438        ("burstable-big-memory", 0.5, 8 * GI, 1, 1, 0, None),
1439        ("small", 1.0, 2 * GI, 2, 4, 10 * GI, None),
1440        ("medium", 4.0, 8 * GI, 1, 1, 10 * GI, None),
1441        ("general", 2.0, 8 * GI, 4, 10, 20 * GI, None),
1442        ("many-replicas", 2.0, 4 * GI, 35, 35, 20 * GI, None),
1443        ("large-replica", 12.0, 48 * GI, 1, 2, 20 * GI, None),
1444        ("huge-replica", 60.0, 240 * GI, 1, 1, 20 * GI, None),
1445        ("memory-heavy", 4.0, 60 * GI, 2, 4, 20 * GI, None),
1446        ("storage-300g", 2.0, 8 * GI, 2, 2, 300 * GI, None),
1447        ("storage-1t", 4.0, 16 * GI, 2, 4, 1000 * GI, None),
1448        ("storage-2t", 4.0, 16 * GI, 2, 4, 2000 * GI, None),
1449        ("storage-8t", 2.0, 8 * GI, 2, 2, 8000 * GI, None),
1450        ("storage-100t", 2.0, 8 * GI, 1, 1, 100 * TI, None),
1451        ("gpu-t4", 2.0, 8 * GI, 1, 1, 10 * GI, Some(("nvidia-t4", 1))),
1452        ("gpu-a10g", 2.0, 8 * GI, 1, 1, 10 * GI, Some(("nvidia-a10g", 1))),
1453        ("gpu-a100", 4.0, 16 * GI, 1, 1, 10 * GI, Some(("nvidia-a100", 1))),
1454        ("gpu-a100x8", 16.0, 128 * GI, 1, 1, 10 * GI, Some(("nvidia-a100", 8))),
1455        ("gpu-h100x8", 16.0, 128 * GI, 1, 1, 10 * GI, Some(("nvidia-h100", 8))),
1456    ];
1457
1458    /// One line per case: what `select_instance_type` picks for every platform, architecture,
1459    /// nested-virtualization setting and workload shape above.
1460    fn render_selection_matrix() -> String {
1461        let mut table = String::new();
1462        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1463            for architecture in [None, Some(Architecture::Arm64), Some(Architecture::X86_64)] {
1464                for nested_virt in [false, true] {
1465                    for &(shape, cpu, memory, desired, max, storage, gpu) in SELECTION_SHAPES {
1466                        let requirements = WorkloadRequirements {
1467                            total_cpu_at_desired: cpu * desired as f64,
1468                            total_memory_bytes_at_desired: memory * desired as u64,
1469                            total_cpu_at_max: cpu * max as f64,
1470                            total_memory_bytes_at_max: memory * max as u64,
1471                            max_cpu_per_container: cpu,
1472                            max_memory_per_container: memory,
1473                            max_ephemeral_storage_bytes: storage,
1474                            gpu: gpu.map(|(gpu_type, count)| GpuSpec {
1475                                gpu_type: gpu_type.to_string(),
1476                                count,
1477                            }),
1478                            architecture,
1479                            nested_virt,
1480                        };
1481                        let pick = match select_instance_type(platform, &requirements) {
1482                            Ok(selection) => format!(
1483                                "{} min={} max={} memory={} storage={}",
1484                                selection.instance_type,
1485                                selection.min_machines,
1486                                selection.max_machines,
1487                                selection.profile.memory_bytes,
1488                                selection.profile.ephemeral_storage_bytes
1489                            ),
1490                            Err(error) => format!("error: {error}"),
1491                        };
1492                        let architecture = architecture
1493                            .map_or("default".to_string(), |architecture| {
1494                                format!("{architecture:?}")
1495                            });
1496                        table.push_str(&format!(
1497                            "{platform} {architecture} nested={nested_virt} {shape} -> {pick}\n"
1498                        ));
1499                    }
1500                }
1501            }
1502        }
1503        table
1504    }
1505
1506    /// Matrix cases whose pick differs from the baseline on purpose, with the pick they get.
1507    /// AWS reports an NVIDIA A10G on g5, so a T4 request has no AWS machine and an A10G
1508    /// request gets g5.
1509    #[rustfmt::skip]
1510    const ACCEPTED_PICK_CHANGES: &[(&str, &str)] = &[
1511        (
1512            "aws default nested=false gpu-t4",
1513            "error: Arm64: architecture Arm64 is unavailable for this workload on platform aws; \
1514             X86_64: no instance type for GPU type 'nvidia-t4' x1 on platform aws",
1515        ),
1516        (
1517            "aws default nested=false gpu-a10g",
1518            "g5.xlarge min=1 max=1 memory=17179869184 storage=268435456000",
1519        ),
1520        (
1521            "aws X86_64 nested=false gpu-t4",
1522            "error: no instance type for GPU type 'nvidia-t4' x1 on platform aws",
1523        ),
1524        (
1525            "aws X86_64 nested=false gpu-a10g",
1526            "g5.xlarge min=1 max=1 memory=17179869184 storage=268435456000",
1527        ),
1528    ];
1529
1530    /// Automatic picks, with the memory and disk they provision, are what existing stacks get
1531    /// on their next deploy, so a catalog change must leave them alone. The baseline holds the
1532    /// picks every case gets; only the cases in `ACCEPTED_PICK_CHANGES` may differ.
1533    #[test]
1534    fn automatic_picks_match_the_baseline_except_accepted_changes() {
1535        let baseline = include_str!("instance_catalog/auto_selection_baseline.txt");
1536        let rendered = render_selection_matrix();
1537        assert_eq!(baseline.lines().count(), rendered.lines().count());
1538        fn split(line: &str) -> (&str, &str) {
1539            line.split_once(" -> ").unwrap_or((line, ""))
1540        }
1541        for (before, now) in baseline.lines().zip(rendered.lines()) {
1542            let ((case, _), (now_case, pick)) = (split(before), split(now));
1543            assert_eq!(case, now_case);
1544            match ACCEPTED_PICK_CHANGES
1545                .iter()
1546                .find(|(accepted, _)| *accepted == case)
1547            {
1548                Some((_, accepted_pick)) => {
1549                    assert_ne!(
1550                        before, now,
1551                        "{case} no longer changes; drop it from the list"
1552                    );
1553                    assert_eq!(pick, *accepted_pick, "{case}");
1554                }
1555                None => assert_eq!(before, now, "automatic pick changed"),
1556            }
1557        }
1558    }
1559
1560    type PinnedMachine = (
1561        Platform,
1562        &'static str,
1563        u32,
1564        u64,
1565        u64,
1566        Option<(&'static str, u32)>,
1567        bool,
1568    );
1569
1570    /// Machines existing stacks may run, with the vCPU, memory, disk, GPU and nested
1571    /// virtualization they are planned with. Each must stay in the catalog with those values.
1572    #[rustfmt::skip]
1573    const PINNED_MACHINES: &[PinnedMachine] = &[
1574        (Platform::Aws, "t4g.micro", 2, GI, 20 * GI, None, false),
1575        (Platform::Aws, "t4g.small", 2, 2 * GI, 20 * GI, None, false),
1576        (Platform::Aws, "t4g.medium", 2, 4 * GI, 20 * GI, None, false),
1577        (Platform::Aws, "t4g.large", 2, 8 * GI, 20 * GI, None, false),
1578        (Platform::Aws, "t3.xlarge", 4, 16 * GI, 20 * GI, None, false),
1579        (Platform::Aws, "t4g.xlarge", 4, 16 * GI, 20 * GI, None, false),
1580        (Platform::Aws, "m7g.medium", 1, 4 * GI, 20 * GI, None, false),
1581        (Platform::Aws, "m7i.large", 2, 8 * GI, 20 * GI, None, false),
1582        (Platform::Aws, "m7g.large", 2, 8 * GI, 20 * GI, None, false),
1583        (Platform::Aws, "m8i.large", 2, 8 * GI, 20 * GI, None, true),
1584        (Platform::Aws, "m7i.xlarge", 4, 16 * GI, 20 * GI, None, false),
1585        (Platform::Aws, "m7g.xlarge", 4, 16 * GI, 20 * GI, None, false),
1586        (Platform::Aws, "m8i.xlarge", 4, 16 * GI, 20 * GI, None, true),
1587        (Platform::Aws, "m7i.2xlarge", 8, 32 * GI, 20 * GI, None, false),
1588        (Platform::Aws, "m7g.2xlarge", 8, 32 * GI, 20 * GI, None, false),
1589        (Platform::Aws, "m8i.2xlarge", 8, 32 * GI, 20 * GI, None, true),
1590        (Platform::Aws, "m7i.4xlarge", 16, 64 * GI, 20 * GI, None, false),
1591        (Platform::Aws, "m7g.4xlarge", 16, 64 * GI, 20 * GI, None, false),
1592        (Platform::Aws, "m8i.4xlarge", 16, 64 * GI, 20 * GI, None, true),
1593        (Platform::Aws, "c7g.medium", 1, 2 * GI, 20 * GI, None, false),
1594        (Platform::Aws, "c8g.medium", 1, 2 * GI, 20 * GI, None, false),
1595        (Platform::Aws, "c7g.large", 2, 4 * GI, 20 * GI, None, false),
1596        (Platform::Aws, "c8g.large", 2, 4 * GI, 20 * GI, None, false),
1597        (Platform::Aws, "c8i.large", 2, 4 * GI, 20 * GI, None, true),
1598        (Platform::Aws, "c7g.xlarge", 4, 8 * GI, 20 * GI, None, false),
1599        (Platform::Aws, "c8g.xlarge", 4, 8 * GI, 20 * GI, None, false),
1600        (Platform::Aws, "c8i.xlarge", 4, 8 * GI, 20 * GI, None, true),
1601        (Platform::Aws, "c7g.2xlarge", 8, 16 * GI, 20 * GI, None, false),
1602        (Platform::Aws, "c8g.2xlarge", 8, 16 * GI, 20 * GI, None, false),
1603        (Platform::Aws, "c8i.2xlarge", 8, 16 * GI, 20 * GI, None, true),
1604        (Platform::Aws, "c7g.4xlarge", 16, 32 * GI, 20 * GI, None, false),
1605        (Platform::Aws, "c8g.4xlarge", 16, 32 * GI, 20 * GI, None, false),
1606        (Platform::Aws, "c8i.4xlarge", 16, 32 * GI, 20 * GI, None, true),
1607        (Platform::Aws, "r7g.medium", 1, 8 * GI, 20 * GI, None, false),
1608        (Platform::Aws, "r7g.large", 2, 16 * GI, 20 * GI, None, false),
1609        (Platform::Aws, "r6a.large", 2, 16 * GI, 20 * GI, None, false),
1610        (Platform::Aws, "r7i.large", 2, 16 * GI, 20 * GI, None, false),
1611        (Platform::Aws, "r7g.xlarge", 4, 32 * GI, 20 * GI, None, false),
1612        (Platform::Aws, "r6a.xlarge", 4, 32 * GI, 20 * GI, None, false),
1613        (Platform::Aws, "r7i.xlarge", 4, 32 * GI, 20 * GI, None, false),
1614        (Platform::Aws, "r7g.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1615        (Platform::Aws, "r6a.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1616        (Platform::Aws, "r7i.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1617        (Platform::Aws, "r7g.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1618        (Platform::Aws, "r6a.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1619        (Platform::Aws, "r7i.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1620        (Platform::Aws, "i4i.xlarge", 4, 32 * GI, 937 * GI, None, false),
1621        (Platform::Aws, "i4i.2xlarge", 8, 64 * GI, 1875 * GI, None, false),
1622        (Platform::Aws, "i4i.4xlarge", 16, 128 * GI, 3750 * GI, None, false),
1623        (Platform::Aws, "i4i.8xlarge", 32, 256 * GI, 7500 * GI, None, false),
1624        (Platform::Aws, "g5.xlarge", 4, 16 * GI, 250 * GI, Some(("nvidia-t4", 1)), false),
1625        (Platform::Aws, "g5.2xlarge", 8, 32 * GI, 450 * GI, Some(("nvidia-t4", 1)), false),
1626        (Platform::Aws, "p4d.24xlarge", 96, 1152 * GI, 8000 * GI, Some(("nvidia-a100", 8)), false),
1627        (Platform::Aws, "p5.48xlarge", 192, 2048 * GI, 8000 * GI, Some(("nvidia-h100", 8)), false),
1628        (Platform::Gcp, "e2-micro", 2, GI, 20 * GI, None, false),
1629        (Platform::Gcp, "e2-small", 2, 2 * GI, 20 * GI, None, false),
1630        (Platform::Gcp, "e2-medium", 2, 4 * GI, 20 * GI, None, false),
1631        (Platform::Gcp, "n2-standard-2", 2, 8 * GI, 20 * GI, None, true),
1632        (Platform::Gcp, "n2-standard-4", 4, 16 * GI, 20 * GI, None, true),
1633        (Platform::Gcp, "n2-standard-8", 8, 32 * GI, 20 * GI, None, true),
1634        (Platform::Gcp, "n2-standard-16", 16, 64 * GI, 20 * GI, None, true),
1635        (Platform::Gcp, "c3-standard-4", 4, 8 * GI, 20 * GI, None, false),
1636        (Platform::Gcp, "c3-standard-8", 8, 16 * GI, 20 * GI, None, false),
1637        (Platform::Gcp, "n2-highmem-2", 2, 16 * GI, 20 * GI, None, false),
1638        (Platform::Gcp, "n2-highmem-4", 4, 32 * GI, 20 * GI, None, false),
1639        (Platform::Gcp, "n2-highmem-8", 8, 64 * GI, 20 * GI, None, false),
1640        (Platform::Gcp, "n2-highmem-16", 16, 128 * GI, 20 * GI, None, false),
1641        (Platform::Gcp, "n2-highmem-32", 32, 256 * GI, 20 * GI, None, false),
1642        (Platform::Gcp, "c3d-standard-8", 8, 32 * GI, 480 * GI, None, false),
1643        (Platform::Gcp, "c3d-standard-16", 16, 64 * GI, 960 * GI, None, false),
1644        (Platform::Gcp, "c3d-standard-30", 30, 120 * GI, 1920 * GI, None, false),
1645        (Platform::Gcp, "n1-standard-4-t4", 4, 15 * GI, 100 * GI, Some(("nvidia-t4", 1)), false),
1646        (Platform::Gcp, "a2-highgpu-1g", 12, 85 * GI, 100 * GI, Some(("nvidia-a100", 1)), false),
1647        (Platform::Gcp, "a2-highgpu-8g", 96, 1360 * GI, 100 * GI, Some(("nvidia-a100", 8)), false),
1648        (Platform::Gcp, "a3-highgpu-8g", 208, 1872 * GI, 100 * GI, Some(("nvidia-h100", 8)), false),
1649        (Platform::Azure, "Standard_B1s", 1, GI, 20 * GI, None, false),
1650        (Platform::Azure, "Standard_B2s", 2, 4 * GI, 20 * GI, None, false),
1651        (Platform::Azure, "Standard_B2ms", 2, 8 * GI, 20 * GI, None, false),
1652        (Platform::Azure, "Standard_B4ms", 4, 16 * GI, 20 * GI, None, false),
1653        (Platform::Azure, "Standard_D2s_v4", 2, 8 * GI, 20 * GI, None, false),
1654        (Platform::Azure, "Standard_D2s_v5", 2, 8 * GI, 20 * GI, None, true),
1655        (Platform::Azure, "Standard_D4s_v5", 4, 16 * GI, 20 * GI, None, true),
1656        (Platform::Azure, "Standard_D8s_v5", 8, 32 * GI, 20 * GI, None, true),
1657        (Platform::Azure, "Standard_D16s_v5", 16, 64 * GI, 20 * GI, None, true),
1658        (Platform::Azure, "Standard_F2s_v2", 2, 4 * GI, 20 * GI, None, true),
1659        (Platform::Azure, "Standard_F4s_v2", 4, 8 * GI, 20 * GI, None, true),
1660        (Platform::Azure, "Standard_F8s_v2", 8, 16 * GI, 20 * GI, None, true),
1661        (Platform::Azure, "Standard_F16s_v2", 16, 32 * GI, 20 * GI, None, true),
1662        (Platform::Azure, "Standard_E2s_v5", 2, 16 * GI, 20 * GI, None, true),
1663        (Platform::Azure, "Standard_E4s_v5", 4, 32 * GI, 20 * GI, None, true),
1664        (Platform::Azure, "Standard_E8s_v5", 8, 64 * GI, 20 * GI, None, true),
1665        (Platform::Azure, "Standard_E16s_v5", 16, 128 * GI, 20 * GI, None, true),
1666        (Platform::Azure, "Standard_L8s_v3", 8, 64 * GI, 1788 * GI, None, false),
1667        (Platform::Azure, "Standard_L16s_v3", 16, 128 * GI, 3576 * GI, None, false),
1668        (Platform::Azure, "Standard_L32s_v3", 32, 256 * GI, 7154 * GI, None, false),
1669        (Platform::Azure, "Standard_NC4as_T4_v3", 4, 28 * GI, 176 * GI, Some(("nvidia-t4", 1)), false),
1670        (Platform::Azure, "Standard_NC24ads_A100_v4", 24, 220 * GI, 958 * GI, Some(("nvidia-a100", 1)), false),
1671        (Platform::Azure, "Standard_NC96ads_A100_v4", 96, 880 * GI, 3916 * GI, Some(("nvidia-a100", 4)), false),
1672        (Platform::Azure, "Standard_ND96isr_H100_v5", 96, 1900 * GI, 1000 * GI, Some(("nvidia-h100", 8)), false),
1673    ];
1674
1675    /// Pinned machines whose GPU type the AWS API reports differently, with that type.
1676    const GPU_TYPE_CORRECTIONS: &[(&str, &str)] =
1677        &[("g5.xlarge", "nvidia-a10g"), ("g5.2xlarge", "nvidia-a10g")];
1678
1679    #[test]
1680    fn every_pinned_machine_is_still_in_the_catalog_with_its_values() {
1681        for &(platform, name, vcpu, memory_bytes, storage_bytes, gpu, nested_virt) in
1682            PINNED_MACHINES
1683        {
1684            let spec = find_instance_type(platform, name)
1685                .unwrap_or_else(|| panic!("{platform} machine {name} left the catalog"));
1686            let gpu = gpu.map(|(gpu_type, count)| {
1687                let corrected = GPU_TYPE_CORRECTIONS
1688                    .iter()
1689                    .find(|(corrected, _)| *corrected == name);
1690                (corrected.map_or(gpu_type, |(_, gpu_type)| *gpu_type), count)
1691            });
1692            assert_eq!(
1693                (
1694                    spec.vcpu,
1695                    spec.memory_bytes,
1696                    spec.ephemeral_storage_bytes,
1697                    spec.gpu.map(|gpu| (gpu.gpu_type, gpu.count)),
1698                    spec.nested_virt,
1699                ),
1700                (vcpu, memory_bytes, storage_bytes, gpu, nested_virt),
1701                "{platform} {name}"
1702            );
1703        }
1704    }
1705
1706    #[test]
1707    fn automatic_selection_considers_exactly_the_pinned_machines() {
1708        let pinned: std::collections::HashSet<_> = PINNED_MACHINES
1709            .iter()
1710            .map(|&(platform, name, ..)| (platform, name))
1711            .collect();
1712        let auto_select: std::collections::HashSet<_> = CATALOG
1713            .iter()
1714            .filter(|spec| spec.auto_select)
1715            .map(|spec| (spec.platform, spec.name))
1716            .collect();
1717        assert_eq!(auto_select, pinned);
1718    }
1719
1720    #[test]
1721    fn catalog_lists_no_bare_metal_machine() {
1722        for spec in CATALOG {
1723            assert!(!spec.name.contains("metal"), "{}", spec.name);
1724        }
1725    }
1726
1727    /// The generator writes each series as one run of rows, smallest first.
1728    #[test]
1729    fn test_catalog_instance_types_sorted_by_vcpu_within_series() {
1730        let mut finished = std::collections::HashSet::new();
1731        for window in CATALOG.windows(2) {
1732            let (previous, next) = (&window[0], &window[1]);
1733            if (previous.platform, previous.series) == (next.platform, next.series) {
1734                assert!(
1735                    previous.vcpu <= next.vcpu,
1736                    "{} ({}) before {} ({})",
1737                    previous.name,
1738                    previous.vcpu,
1739                    next.name,
1740                    next.vcpu
1741                );
1742            } else {
1743                finished.insert((previous.platform, previous.series));
1744                assert!(
1745                    !finished.contains(&(next.platform, next.series)),
1746                    "series {} is split",
1747                    next.series
1748                );
1749            }
1750        }
1751    }
1752}