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