1use crate::{GpuSpec, MachineProfile, Platform};
13use serde::{Deserialize, Serialize};
14
15pub 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
38pub 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 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#[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#[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
102pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub struct CatalogGpu {
114 pub gpu_type: &'static str,
115 pub count: u32,
116}
117
118#[derive(Debug, Clone)]
122pub struct InstanceTypeSpec {
123 pub name: &'static str,
124 pub platform: Platform,
125 pub family: InstanceFamily,
127 pub series: &'static str,
129 pub label: &'static str,
131 pub architecture: Architecture,
132 pub vcpu: u32,
134 pub memory_bytes: u64,
136 pub ephemeral_storage_bytes: u64,
138 pub local_disk_bytes: u64,
140 pub machine_type: &'static str,
143 pub gpu: Option<CatalogGpu>,
145 pub nested_virt: bool,
147 pub auto_select: bool,
151}
152
153impl InstanceTypeSpec {
154 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 pub fn is_nested_virt_capable(&self) -> bool {
171 self.nested_virt
172 }
173
174 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 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
201const KI: u64 = 1024;
203const MI: u64 = KI * 1024;
204const GI: u64 = MI * 1024;
205const TI: u64 = GI * 1024;
206
207pub fn max_configurable_ephemeral_storage_bytes(platform: Platform) -> Option<u64> {
213 let disk_limit_gib = match platform {
214 Platform::Aws | Platform::Gcp => 64 * 1024,
216 Platform::Azure => 4_095,
218 Platform::Kubernetes | Platform::Machines | Platform::Local | Platform::Test => {
219 return None;
220 }
221 };
222
223 let max_requested_gib = (disk_limit_gib - 12) * 4 / 5;
228 Some(max_requested_gib * GI)
229}
230
231mod catalog;
234use catalog::CATALOG;
235
236pub fn catalog_for_platform(platform: Platform) -> Vec<&'static InstanceTypeSpec> {
242 CATALOG
243 .iter()
244 .filter(|spec| spec.platform == platform)
245 .collect()
246}
247
248pub 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
255pub fn is_same_architecture_aws_machine(old: &str, new: &str) -> bool {
258 is_same_architecture_machine(Platform::Aws, old, new)
259}
260
261pub 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#[derive(Debug, Clone)]
275pub struct WorkloadRequirements {
276 pub total_cpu_at_desired: f64,
278 pub total_memory_bytes_at_desired: u64,
280 pub total_cpu_at_max: f64,
282 pub total_memory_bytes_at_max: u64,
284 pub max_cpu_per_container: f64,
286 pub max_memory_per_container: u64,
288 pub max_ephemeral_storage_bytes: u64,
290 pub gpu: Option<GpuSpec>,
292 pub architecture: Option<Architecture>,
294 pub nested_virt: bool,
298}
299
300#[derive(Debug, Clone)]
302pub struct InstanceSelection {
303 pub instance_type: &'static str,
305 pub profile: MachineProfile,
307 pub min_machines: u32,
309 pub max_machines: u32,
311}
312
313const STORAGE_OPTIMIZED_THRESHOLD: u64 = 200 * GI;
315
316const MAX_MACHINES_PER_CLUSTER: u32 = 10;
318
319const MAX_STANDARD_VCPU: u32 = 8;
322
323const SYSTEM_RESERVE_CPU: f64 = 0.5;
325
326const WORKLOAD_HEADROOM_FACTOR: f64 = 1.15;
328
329pub fn select_instance_type(
340 platform: Platform,
341 requirements: &WorkloadRequirements,
342) -> Result<InstanceSelection, String> {
343 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 =
360 select_instance_type_for_architecture(platform, &with_architecture(default));
361 if selection.is_ok() {
362 return selection;
363 }
364 return select_instance_type_for_architecture(platform, &with_architecture(other))
367 .map_err(|other_error| {
368 let default_error = selection.err().unwrap_or_default();
369 format!("{default:?}: {default_error}; {other:?}: {other_error}")
370 });
371 }
372 }
373 select_instance_type_for_architecture(platform, requirements)
374}
375
376fn select_instance_type_for_architecture(
377 platform: Platform,
378 requirements: &WorkloadRequirements,
379) -> Result<InstanceSelection, String> {
380 let family = select_family(requirements);
381 let family = if requirements.nested_virt && family == InstanceFamily::Burstable {
382 InstanceFamily::GeneralPurpose
383 } else {
384 family
385 };
386 let selection = select_instance_type_in_family(platform, requirements, family);
387 if selection.is_err() && family == InstanceFamily::Burstable {
388 return select_instance_type_in_family(
389 platform,
390 requirements,
391 InstanceFamily::GeneralPurpose,
392 );
393 }
394 selection
395}
396
397fn select_instance_type_in_family(
398 platform: Platform,
399 requirements: &WorkloadRequirements,
400 family: InstanceFamily,
401) -> Result<InstanceSelection, String> {
402 let architecture = requirements
403 .architecture
404 .or_else(|| default_architecture(platform))
405 .ok_or_else(|| format!("platform {platform} has no default compute architecture"))?;
406
407 let candidates: Vec<&InstanceTypeSpec> = CATALOG
408 .iter()
409 .filter(|spec| spec.auto_select && spec.platform == platform && spec.family == family)
410 .filter(|spec| {
411 if requirements.nested_virt {
412 spec.is_nested_virt_capable()
413 } else {
414 platform != Platform::Aws || !spec.is_nested_virt_capable()
415 }
416 })
417 .filter(|spec| spec.architecture == architecture)
418 .collect();
419
420 if candidates.is_empty() && family != InstanceFamily::StorageOptimized {
425 let family_has_other_architecture = CATALOG.iter().any(|spec| {
426 spec.auto_select
427 && spec.platform == platform
428 && spec.family == family
429 && if requirements.nested_virt {
430 spec.is_nested_virt_capable()
431 } else {
432 platform != Platform::Aws || !spec.is_nested_virt_capable()
433 }
434 });
435 if family_has_other_architecture {
436 return Err(format!(
437 "architecture {architecture:?} is unavailable for this workload on platform {platform}"
438 ));
439 }
440 return Err(if requirements.nested_virt {
441 format!(
442 "no nested-virt-capable {family:?} instance types in catalog for platform {platform}"
443 )
444 } else {
445 format!("no {family:?} instance types in catalog for platform {platform}")
446 });
447 }
448
449 let candidates = if let Some(ref gpu) = requirements.gpu {
451 let filtered: Vec<&InstanceTypeSpec> = candidates
452 .into_iter()
453 .filter(|spec| {
454 spec.gpu.as_ref().map_or(false, |g| {
455 g.gpu_type == gpu.gpu_type && g.count >= gpu.count
456 })
457 })
458 .collect();
459 if filtered.is_empty() {
460 return Err(format!(
461 "no instance type for GPU type '{}' x{} on platform {platform}",
462 gpu.gpu_type, gpu.count
463 ));
464 }
465 filtered
466 } else {
467 candidates
468 };
469
470 let candidates = if family == InstanceFamily::StorageOptimized {
472 let filtered: Vec<&InstanceTypeSpec> = candidates
473 .into_iter()
474 .filter(|spec| spec.ephemeral_storage_bytes >= requirements.max_ephemeral_storage_bytes)
475 .collect();
476 if filtered.is_empty() {
477 CATALOG
481 .iter()
482 .filter(|spec| {
483 spec.auto_select
484 && spec.platform == platform
485 && spec.family == InstanceFamily::GeneralPurpose
486 && spec.has_configurable_ephemeral_storage()
487 && max_configurable_ephemeral_storage_bytes(platform)
488 .is_some_and(|max| requirements.max_ephemeral_storage_bytes <= max)
489 })
490 .filter(|spec| {
491 if requirements.nested_virt {
492 spec.is_nested_virt_capable()
493 } else {
494 platform != Platform::Aws || !spec.is_nested_virt_capable()
495 }
496 })
497 .filter(|spec| spec.architecture == architecture)
498 .collect()
499 } else {
500 filtered
501 }
502 } else {
503 candidates
504 };
505
506 if candidates.is_empty() {
507 return Err(format!(
508 "architecture {architecture:?} is unavailable for this workload on platform {platform}"
509 ));
510 }
511
512 let candidates: Vec<_> = candidates
515 .into_iter()
516 .filter(|spec| {
517 allocatable_cpu(spec) >= requirements.max_cpu_per_container * WORKLOAD_HEADROOM_FACTOR
518 && allocatable_memory_bytes(spec) as f64
519 >= requirements.max_memory_per_container as f64 * WORKLOAD_HEADROOM_FACTOR
520 })
521 .collect();
522 let smallest_replica_machine = candidates.iter().map(|spec| spec.vcpu).min()
523 .ok_or_else(|| format!("no {platform} machine can fit one replica with system reserve and workload headroom"))?;
524
525 let effective_family = candidates[0].family;
530 let vcpu_cap = if effective_family == InstanceFamily::GpuCompute
531 || effective_family == InstanceFamily::StorageOptimized
532 {
533 u32::MAX
534 } else {
535 MAX_STANDARD_VCPU.max(smallest_replica_machine)
536 };
537
538 let desired_target_machines = desired_target_machines(requirements);
539 let target_cpu = requirements
540 .max_cpu_per_container
541 .max(requirements.total_cpu_at_desired / desired_target_machines as f64)
542 * WORKLOAD_HEADROOM_FACTOR;
543 let target_memory = (requirements.max_memory_per_container as f64)
544 .max(requirements.total_memory_bytes_at_desired as f64 / desired_target_machines as f64)
545 * WORKLOAD_HEADROOM_FACTOR;
546
547 let selected = candidates
552 .iter()
553 .filter(|spec| {
554 spec.vcpu <= vcpu_cap
555 && allocatable_cpu(spec) >= target_cpu
556 && allocatable_memory_bytes(spec) as f64 >= target_memory
557 })
558 .min_by_key(|spec| spec.vcpu)
559 .or_else(|| {
560 candidates
562 .iter()
563 .filter(|spec| spec.vcpu <= vcpu_cap)
564 .max_by_key(|spec| spec.vcpu)
565 })
566 .or_else(|| {
567 candidates.iter().min_by_key(|spec| spec.vcpu)
569 })
570 .ok_or_else(|| format!("no instance types available for platform {platform}"))?;
571
572 let max_machines = compute_max_machines(requirements, selected);
574 let min_machines = compute_min_machines(requirements, selected, max_machines);
575
576 Ok(InstanceSelection {
577 instance_type: selected.name,
578 profile: selected.to_machine_profile_for_storage(requirements.max_ephemeral_storage_bytes),
579 min_machines,
580 max_machines,
581 })
582}
583
584pub fn select_family(requirements: &WorkloadRequirements) -> InstanceFamily {
590 if requirements.gpu.is_some() {
592 return InstanceFamily::GpuCompute;
593 }
594
595 if requirements.max_ephemeral_storage_bytes > STORAGE_OPTIMIZED_THRESHOLD {
597 return InstanceFamily::StorageOptimized;
598 }
599
600 if requirements.total_cpu_at_max < 2.0 {
602 return InstanceFamily::Burstable;
603 }
604
605 InstanceFamily::GeneralPurpose
607}
608
609fn compute_max_machines(requirements: &WorkloadRequirements, instance: &InstanceTypeSpec) -> u32 {
611 let cpu_with_headroom = requirements.total_cpu_at_max * WORKLOAD_HEADROOM_FACTOR;
612 let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
613
614 let mem_with_headroom =
615 requirements.total_memory_bytes_at_max as f64 * WORKLOAD_HEADROOM_FACTOR;
616 let mem_machines =
617 (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
618
619 cpu_machines
621 .max(mem_machines)
622 .max(1)
623 .min(MAX_MACHINES_PER_CLUSTER)
624}
625
626fn compute_min_machines(
628 requirements: &WorkloadRequirements,
629 instance: &InstanceTypeSpec,
630 max_machines: u32,
631) -> u32 {
632 let cpu_with_headroom = requirements.total_cpu_at_desired * WORKLOAD_HEADROOM_FACTOR;
633 let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
634
635 let mem_with_headroom =
636 requirements.total_memory_bytes_at_desired as f64 * WORKLOAD_HEADROOM_FACTOR;
637 let mem_machines =
638 (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
639
640 cpu_machines
641 .max(mem_machines)
642 .max(1)
643 .min(2)
644 .min(max_machines)
645}
646
647fn desired_target_machines(requirements: &WorkloadRequirements) -> u32 {
648 if requirements.total_cpu_at_desired >= 2.0
649 || requirements.total_memory_bytes_at_desired >= 4 * GI
650 {
651 2
652 } else {
653 1
654 }
655}
656
657pub(crate) fn allocatable_cpu(instance: &InstanceTypeSpec) -> f64 {
658 (instance.vcpu as f64 - SYSTEM_RESERVE_CPU).max(0.25)
659}
660
661pub(crate) fn allocatable_memory_bytes(instance: &InstanceTypeSpec) -> u64 {
662 instance
663 .memory_bytes
664 .saturating_sub(system_reserve_memory_bytes(instance.memory_bytes))
665 .max(256 * MI)
666}
667
668fn system_reserve_memory_bytes(memory_bytes: u64) -> u64 {
669 if memory_bytes < 4 * GI {
670 256 * MI
671 } else if memory_bytes < 16 * GI {
672 512 * MI
673 } else {
674 GI
675 }
676}
677
678#[cfg(test)]
683mod tests {
684 use super::*;
685 use crate::BinaryTarget;
686
687 #[test]
690 fn test_parse_cpu_plain() {
691 assert_eq!(parse_cpu("1").unwrap(), 1.0);
692 assert_eq!(parse_cpu("0.5").unwrap(), 0.5);
693 assert_eq!(parse_cpu("2.0").unwrap(), 2.0);
694 assert_eq!(parse_cpu("16").unwrap(), 16.0);
695 }
696
697 #[test]
698 fn test_parse_cpu_millicore() {
699 assert_eq!(parse_cpu("500m").unwrap(), 0.5);
700 assert_eq!(parse_cpu("250m").unwrap(), 0.25);
701 assert_eq!(parse_cpu("1000m").unwrap(), 1.0);
702 assert_eq!(parse_cpu("100m").unwrap(), 0.1);
703 }
704
705 #[test]
706 fn test_parse_cpu_invalid() {
707 assert!(parse_cpu("").is_err());
708 assert!(parse_cpu("abc").is_err());
709 assert!(parse_cpu("m").is_err());
710 }
711
712 #[test]
713 fn test_parse_memory_binary_suffixes() {
714 assert_eq!(parse_memory_bytes("1Ki").unwrap(), 1024);
715 assert_eq!(parse_memory_bytes("1Mi").unwrap(), 1024 * 1024);
716 assert_eq!(parse_memory_bytes("1Gi").unwrap(), 1024 * 1024 * 1024);
717 assert_eq!(parse_memory_bytes("4Gi").unwrap(), 4 * 1024 * 1024 * 1024);
718 assert_eq!(parse_memory_bytes("512Mi").unwrap(), 512 * 1024 * 1024);
719 assert_eq!(
720 parse_memory_bytes("1Ti").unwrap(),
721 1024u64 * 1024 * 1024 * 1024
722 );
723 }
724
725 #[test]
726 fn test_parse_memory_decimal_suffixes() {
727 assert_eq!(parse_memory_bytes("1k").unwrap(), 1000);
728 assert_eq!(parse_memory_bytes("1M").unwrap(), 1_000_000);
729 assert_eq!(parse_memory_bytes("1G").unwrap(), 1_000_000_000);
730 assert_eq!(parse_memory_bytes("1T").unwrap(), 1_000_000_000_000);
731 }
732
733 #[test]
734 fn test_parse_memory_plain_bytes() {
735 assert_eq!(parse_memory_bytes("1024").unwrap(), 1024);
736 assert_eq!(parse_memory_bytes("0").unwrap(), 0);
737 }
738
739 #[test]
740 fn test_parse_memory_invalid() {
741 assert!(parse_memory_bytes("").is_err());
742 assert!(parse_memory_bytes("abc").is_err());
743 assert!(parse_memory_bytes("Gi").is_err());
744 }
745
746 #[test]
747 fn test_parse_memory_fractional() {
748 assert_eq!(parse_memory_bytes("0.5Gi").unwrap(), GI / 2);
749 assert_eq!(parse_memory_bytes("1.5Gi").unwrap(), GI + GI / 2);
750 }
751
752 #[test]
755 fn test_catalog_has_entries_for_all_cloud_platforms() {
756 assert!(!catalog_for_platform(Platform::Aws).is_empty());
757 assert!(!catalog_for_platform(Platform::Gcp).is_empty());
758 assert!(!catalog_for_platform(Platform::Azure).is_empty());
759 }
760
761 #[test]
762 fn test_catalog_no_entries_for_non_cloud_platforms() {
763 assert!(catalog_for_platform(Platform::Local).is_empty());
764 assert!(catalog_for_platform(Platform::Kubernetes).is_empty());
765 }
766
767 #[test]
768 fn test_find_known_instance_type() {
769 let spec =
770 find_instance_type(Platform::Aws, "m7g.2xlarge").expect("should find m7g.2xlarge");
771 assert_eq!(spec.vcpu, 8);
772 assert_eq!(spec.memory_bytes, 32 * GI);
773 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
774 }
775
776 #[test]
777 fn test_find_aws_c8i_nested_virt_instance_type() {
778 let spec = find_instance_type(Platform::Aws, "c8i.large").expect("should find c8i.large");
779 assert_eq!(spec.vcpu, 2);
780 assert_eq!(spec.memory_bytes, 4 * GI);
781 assert_eq!(spec.family, InstanceFamily::ComputeOptimized);
782 assert_eq!(spec.architecture, Architecture::X86_64);
783 assert!(spec.is_nested_virt_capable());
784 }
785
786 #[test]
787 fn test_find_unknown_instance_type() {
788 assert!(find_instance_type(Platform::Aws, "nonexistent.xlarge").is_none());
789 }
790
791 #[test]
792 fn test_find_wrong_platform() {
793 assert!(find_instance_type(Platform::Gcp, "m7g.2xlarge").is_none());
794 }
795
796 #[test]
797 fn test_to_machine_profile() {
798 let spec = find_instance_type(Platform::Aws, "m7g.2xlarge").unwrap();
799 let profile = spec.to_machine_profile();
800 assert_eq!(profile.cpu, "8.0");
801 assert_eq!(profile.memory_bytes, 32 * GI);
802 assert_eq!(profile.ephemeral_storage_bytes, 20 * GI);
803 assert!(profile.gpu.is_none());
804 }
805
806 #[test]
807 fn test_to_machine_profile_with_gpu() {
808 let spec = find_instance_type(Platform::Aws, "p4d.24xlarge").unwrap();
809 let profile = spec.to_machine_profile();
810 let gpu = profile.gpu.as_ref().expect("should have GPU");
811 assert_eq!(gpu.gpu_type, "nvidia-a100");
812 assert_eq!(gpu.count, 8);
813 }
814
815 #[test]
818 fn small_cpu_with_large_memory_falls_back_from_burstable() {
819 let requirements = WorkloadRequirements {
820 total_cpu_at_desired: 0.5,
821 total_memory_bytes_at_desired: 8 * GI,
822 total_cpu_at_max: 0.5,
823 total_memory_bytes_at_max: 8 * GI,
824 max_cpu_per_container: 0.5,
825 max_memory_per_container: 8 * GI,
826 max_ephemeral_storage_bytes: 0,
827 gpu: None,
828 architecture: Some(Architecture::X86_64),
829 nested_virt: false,
830 };
831 let selected = select_instance_type(Platform::Gcp, &requirements).unwrap();
832 let machine = find_instance_type(Platform::Gcp, selected.instance_type).unwrap();
833 assert!(
834 allocatable_memory_bytes(machine) as f64 >= 8.0 * GI as f64 * WORKLOAD_HEADROOM_FACTOR
835 );
836 }
837
838 #[test]
839 fn test_select_burstable_for_small_workload() {
840 let req = WorkloadRequirements {
841 total_cpu_at_desired: 1.0,
842 total_memory_bytes_at_desired: 2 * GI,
843 total_cpu_at_max: 1.0,
844 total_memory_bytes_at_max: 2 * GI,
845 max_cpu_per_container: 0.5,
846 max_memory_per_container: 1 * GI,
847 max_ephemeral_storage_bytes: 10 * GI,
848 gpu: None,
849 architecture: None,
850 nested_virt: false,
851 };
852 let sel = select_instance_type(Platform::Aws, &req).unwrap();
853 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
854 assert_eq!(spec.family, InstanceFamily::Burstable);
855 }
856
857 #[test]
858 fn test_selects_smallest_burstable_machine_with_real_headroom() {
859 let req = WorkloadRequirements {
860 total_cpu_at_desired: 1.0,
861 total_memory_bytes_at_desired: 2 * GI,
862 total_cpu_at_max: 1.0,
863 total_memory_bytes_at_max: 2 * GI,
864 max_cpu_per_container: 1.0,
865 max_memory_per_container: 2 * GI,
866 max_ephemeral_storage_bytes: 10 * GI,
867 gpu: None,
868 architecture: None,
869 nested_virt: false,
870 };
871
872 let selection = select_instance_type(Platform::Aws, &req).unwrap();
873
874 assert_eq!(selection.instance_type, "t4g.medium");
875 assert_eq!(selection.min_machines, 1);
876 assert_eq!(selection.max_machines, 1);
877 }
878
879 #[test]
880 fn test_select_general_purpose_for_standard_workload() {
881 let req = WorkloadRequirements {
883 total_cpu_at_desired: 20.0,
884 total_memory_bytes_at_desired: 80 * GI,
885 total_cpu_at_max: 20.0,
886 total_memory_bytes_at_max: 80 * GI,
887 max_cpu_per_container: 2.0,
888 max_memory_per_container: 8 * GI,
889 max_ephemeral_storage_bytes: 10 * GI,
890 gpu: None,
891 architecture: None,
892 nested_virt: false,
893 };
894 let sel = select_instance_type(Platform::Aws, &req).unwrap();
895 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
896 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
897 }
898
899 #[test]
900 fn test_select_general_purpose_even_for_cpu_heavy() {
901 let req = WorkloadRequirements {
903 total_cpu_at_desired: 20.0,
904 total_memory_bytes_at_desired: 20 * GI,
905 total_cpu_at_max: 20.0,
906 total_memory_bytes_at_max: 20 * GI,
907 max_cpu_per_container: 2.0,
908 max_memory_per_container: 2 * GI,
909 max_ephemeral_storage_bytes: 10 * GI,
910 gpu: None,
911 architecture: None,
912 nested_virt: false,
913 };
914 let sel = select_instance_type(Platform::Aws, &req).unwrap();
915 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
916 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
917 }
918
919 #[test]
920 fn test_select_storage_optimized_for_large_ephemeral() {
921 let req = WorkloadRequirements {
922 total_cpu_at_desired: 8.0,
923 total_memory_bytes_at_desired: 32 * GI,
924 total_cpu_at_max: 8.0,
925 total_memory_bytes_at_max: 32 * GI,
926 max_cpu_per_container: 2.0,
927 max_memory_per_container: 8 * GI,
928 max_ephemeral_storage_bytes: 500 * GI,
929 gpu: None,
930 architecture: Some(Architecture::X86_64),
931 nested_virt: false,
932 };
933 let sel = select_instance_type(Platform::Aws, &req).unwrap();
934 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
935 assert_eq!(spec.family, InstanceFamily::StorageOptimized);
936 }
937
938 #[test]
939 fn test_select_configurable_storage_above_fixed_local_catalog() {
940 let req = WorkloadRequirements {
941 total_cpu_at_desired: 8.0,
942 total_memory_bytes_at_desired: 32 * GI,
943 total_cpu_at_max: 8.0,
944 total_memory_bytes_at_max: 32 * GI,
945 max_cpu_per_container: 2.0,
946 max_memory_per_container: 8 * GI,
947 max_ephemeral_storage_bytes: 8_000 * GI,
948 gpu: None,
949 architecture: Some(Architecture::X86_64),
950 nested_virt: false,
951 };
952 let sel = select_instance_type(Platform::Aws, &req).unwrap();
953 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
954 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
955 assert_eq!(sel.profile.ephemeral_storage_bytes, 8_000 * GI);
956 }
957
958 #[test]
959 fn test_configurable_storage_fallback_retains_standard_vcpu_cap() {
960 let req = WorkloadRequirements {
961 total_cpu_at_desired: 70.0,
962 total_memory_bytes_at_desired: 140 * GI,
963 total_cpu_at_max: 70.0,
964 total_memory_bytes_at_max: 140 * GI,
965 max_cpu_per_container: 2.0,
966 max_memory_per_container: 4 * GI,
967 max_ephemeral_storage_bytes: 8_000 * GI,
968 gpu: None,
969 architecture: Some(Architecture::X86_64),
970 nested_virt: false,
971 };
972
973 let sel = select_instance_type(Platform::Aws, &req).unwrap();
974 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
975 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
976 assert!(spec.vcpu <= MAX_STANDARD_VCPU);
977 assert!(sel.max_machines > 1);
978 }
979
980 #[test]
981 fn test_nested_virtualization_storage_falls_back_to_configurable_disk() {
982 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
983 let req = WorkloadRequirements {
984 total_cpu_at_desired: 2.0,
985 total_memory_bytes_at_desired: 8 * GI,
986 total_cpu_at_max: 2.0,
987 total_memory_bytes_at_max: 8 * GI,
988 max_cpu_per_container: 2.0,
989 max_memory_per_container: 8 * GI,
990 max_ephemeral_storage_bytes: 500 * GI,
991 gpu: None,
992 architecture: Some(Architecture::X86_64),
993 nested_virt: true,
994 };
995
996 let sel = select_instance_type(platform, &req)
997 .unwrap_or_else(|error| panic!("{platform} should support fallback: {error}"));
998 let spec = find_instance_type(platform, sel.instance_type).unwrap();
999 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1000 assert!(spec.is_nested_virt_capable());
1001 assert_eq!(sel.profile.ephemeral_storage_bytes, 500 * GI);
1002 }
1003 }
1004
1005 #[test]
1006 fn unset_architecture_resolves_to_the_one_offering_nested_virtualization() {
1007 let req = WorkloadRequirements {
1008 total_cpu_at_desired: 2.0,
1009 total_memory_bytes_at_desired: 4 * GI,
1010 total_cpu_at_max: 2.0,
1011 total_memory_bytes_at_max: 4 * GI,
1012 max_cpu_per_container: 2.0,
1013 max_memory_per_container: 4 * GI,
1014 max_ephemeral_storage_bytes: 20 * GI,
1015 gpu: None,
1016 architecture: None,
1017 nested_virt: true,
1018 };
1019
1020 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1021 let selection = select_instance_type(platform, &req)
1022 .unwrap_or_else(|message| panic!("{platform}: {message}"));
1023 let spec = find_instance_type(platform, selection.instance_type)
1024 .expect("selected machine should exist in the catalog");
1025 assert!(spec.is_nested_virt_capable(), "{platform}: {}", spec.name);
1026 assert_eq!(spec.architecture, Architecture::X86_64, "{platform}");
1027 assert_eq!(selection.profile.architecture, Some(Architecture::X86_64));
1028 }
1029
1030 let oversized = WorkloadRequirements {
1033 max_cpu_per_container: 1000.0,
1034 total_cpu_at_desired: 1000.0,
1035 total_cpu_at_max: 1000.0,
1036 ..req.clone()
1037 };
1038 let message = select_instance_type(Platform::Aws, &oversized)
1039 .expect_err("no AWS machine fits 1000 vCPU per replica");
1040 assert!(
1041 message.starts_with("Arm64: architecture Arm64 is unavailable"),
1042 "{message}"
1043 );
1044 assert!(
1045 message.contains("; X86_64: no aws machine can fit one replica"),
1046 "{message}"
1047 );
1048
1049 let arm = WorkloadRequirements {
1051 architecture: Some(Architecture::Arm64),
1052 ..req
1053 };
1054 assert!(select_instance_type(Platform::Aws, &arm).is_err());
1055 }
1056
1057 #[test]
1058 fn test_configurable_disk_rejects_capacity_above_provider_limit() {
1059 for (platform, requested) in [
1060 (Platform::Aws, 100 * TI),
1061 (Platform::Gcp, 100 * TI),
1062 (Platform::Azure, 8_000 * GI),
1063 ] {
1064 let req = WorkloadRequirements {
1065 total_cpu_at_desired: 2.0,
1066 total_memory_bytes_at_desired: 8 * GI,
1067 total_cpu_at_max: 2.0,
1068 total_memory_bytes_at_max: 8 * GI,
1069 max_cpu_per_container: 2.0,
1070 max_memory_per_container: 8 * GI,
1071 max_ephemeral_storage_bytes: requested,
1072 gpu: None,
1073 architecture: Some(Architecture::X86_64),
1074 nested_virt: true,
1075 };
1076
1077 assert!(select_instance_type(platform, &req).is_err());
1078 }
1079 }
1080
1081 #[test]
1082 fn test_configurable_disk_limits_account_for_controller_rounding() {
1083 for (platform, disk_limit_gib) in [
1084 (Platform::Aws, 64 * 1024),
1085 (Platform::Gcp, 64 * 1024),
1086 (Platform::Azure, 4_095),
1087 ] {
1088 let max = max_configurable_ephemeral_storage_bytes(platform).unwrap();
1089 let materialized_disk_gib = |requested_bytes: u64| {
1090 let requested_gib = requested_bytes.div_ceil(GI);
1091 (requested_gib * 5).div_ceil(4) + 12
1092 };
1093
1094 assert!(materialized_disk_gib(max) <= disk_limit_gib);
1095 assert!(materialized_disk_gib(max + 1) > disk_limit_gib);
1096 }
1097 }
1098
1099 #[test]
1100 fn test_gpu_local_storage_is_not_treated_as_resizable() {
1101 let spec = find_instance_type(Platform::Aws, "g5.xlarge").unwrap();
1102 assert_eq!(spec.family, InstanceFamily::GpuCompute);
1103 assert!(!spec.has_configurable_ephemeral_storage());
1104 }
1105
1106 #[test]
1107 fn test_select_gpu_instance() {
1108 let req = WorkloadRequirements {
1109 total_cpu_at_desired: 8.0,
1110 total_memory_bytes_at_desired: 32 * GI,
1111 total_cpu_at_max: 8.0,
1112 total_memory_bytes_at_max: 32 * GI,
1113 max_cpu_per_container: 4.0,
1114 max_memory_per_container: 16 * GI,
1115 max_ephemeral_storage_bytes: 10 * GI,
1116 gpu: Some(GpuSpec {
1117 gpu_type: "nvidia-a100".to_string(),
1118 count: 1,
1119 }),
1120 architecture: Some(Architecture::X86_64),
1121 nested_virt: false,
1122 };
1123 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1124 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1125 assert_eq!(spec.family, InstanceFamily::GpuCompute);
1126 assert!(spec.gpu.is_some());
1127 }
1128
1129 #[test]
1130 fn test_select_uses_each_cloud_image_target_architecture() {
1131 let req = WorkloadRequirements {
1132 total_cpu_at_desired: 4.0,
1133 total_memory_bytes_at_desired: 16 * GI,
1134 total_cpu_at_max: 4.0,
1135 total_memory_bytes_at_max: 16 * GI,
1136 max_cpu_per_container: 1.0,
1137 max_memory_per_container: 4 * GI,
1138 max_ephemeral_storage_bytes: 10 * GI,
1139 gpu: None,
1140 architecture: None,
1141 nested_virt: false,
1142 };
1143 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1144 let sel = select_instance_type(platform, &req)
1145 .unwrap_or_else(|error| panic!("selection failed for {platform}: {error}"));
1146 let spec = find_instance_type(platform, sel.instance_type)
1147 .expect("selected machine should exist in the catalog");
1148 assert_eq!(
1149 Some(spec.architecture),
1150 default_architecture(platform),
1151 "machine architecture must match the image target for {platform}"
1152 );
1153 }
1154 }
1155
1156 #[test]
1157 fn test_machine_count_reasonable() {
1158 let req = WorkloadRequirements {
1160 total_cpu_at_desired: 20.0,
1161 total_memory_bytes_at_desired: 40 * GI,
1162 total_cpu_at_max: 20.0,
1163 total_memory_bytes_at_max: 40 * GI,
1164 max_cpu_per_container: 1.0,
1165 max_memory_per_container: 2 * GI,
1166 max_ephemeral_storage_bytes: 10 * GI,
1167 gpu: None,
1168 architecture: None,
1169 nested_virt: false,
1170 };
1171 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1172 assert!(sel.min_machines >= 1);
1173 assert!(sel.max_machines <= MAX_MACHINES_PER_CLUSTER);
1174 assert!(sel.max_machines >= sel.min_machines);
1175 }
1176
1177 #[test]
1178 fn test_instance_size_capped_at_8_vcpu() {
1179 let req = WorkloadRequirements {
1181 total_cpu_at_desired: 70.0,
1182 total_memory_bytes_at_desired: 140 * GI,
1183 total_cpu_at_max: 70.0,
1184 total_memory_bytes_at_max: 140 * GI,
1185 max_cpu_per_container: 2.0,
1186 max_memory_per_container: 4 * GI,
1187 max_ephemeral_storage_bytes: 10 * GI,
1188 gpu: None,
1189 architecture: None,
1190 nested_virt: false,
1191 };
1192 let sel = select_instance_type(Platform::Gcp, &req).unwrap();
1193 let spec = find_instance_type(Platform::Gcp, sel.instance_type).unwrap();
1194 assert!(
1195 spec.vcpu <= MAX_STANDARD_VCPU,
1196 "selected {} with {} vCPUs, expected <= {}",
1197 spec.name,
1198 spec.vcpu,
1199 MAX_STANDARD_VCPU
1200 );
1201 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1202 assert!(sel.max_machines > 1);
1204 }
1205
1206 #[test]
1207 fn test_larger_autoscaled_workload_gets_reasonable_instance() {
1208 let req = WorkloadRequirements {
1211 total_cpu_at_desired: 70.0,
1212 total_memory_bytes_at_desired: 140 * GI,
1213 total_cpu_at_max: 70.0, total_memory_bytes_at_max: 140 * GI, max_cpu_per_container: 2.0,
1216 max_memory_per_container: 4 * GI,
1217 max_ephemeral_storage_bytes: 20 * GI,
1218 gpu: None,
1219 architecture: None,
1220 nested_virt: false,
1221 };
1222 let sel = select_instance_type(Platform::Gcp, &req).unwrap();
1223 assert_eq!(sel.instance_type, "n2-standard-8");
1225 assert!(sel.max_machines >= 2);
1226 }
1227
1228 #[test]
1235 fn test_select_aws_picks_m8i_when_nested_virt_required() {
1236 let req = WorkloadRequirements {
1237 total_cpu_at_desired: 4.0,
1238 total_memory_bytes_at_desired: 8 * GI,
1239 total_cpu_at_max: 4.0,
1240 total_memory_bytes_at_max: 8 * GI,
1241 max_cpu_per_container: 4.0,
1242 max_memory_per_container: 8 * GI,
1243 max_ephemeral_storage_bytes: 10 * GI,
1244 gpu: None,
1245 architecture: Some(Architecture::X86_64),
1246 nested_virt: true,
1247 };
1248 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1249 assert!(
1250 sel.instance_type.starts_with("m8i.")
1251 || sel.instance_type.starts_with("c8i.")
1252 || sel.instance_type.starts_with("r8i."),
1253 "expected an m8i/c8i/r8i instance, got {}",
1254 sel.instance_type
1255 );
1256 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1257 assert!(spec.is_nested_virt_capable());
1258 }
1259
1260 #[test]
1261 fn test_select_gcp_picks_n2_when_nested_virt_required() {
1262 let req = WorkloadRequirements {
1263 total_cpu_at_desired: 4.0,
1264 total_memory_bytes_at_desired: 8 * GI,
1265 total_cpu_at_max: 4.0,
1266 total_memory_bytes_at_max: 8 * GI,
1267 max_cpu_per_container: 4.0,
1268 max_memory_per_container: 8 * GI,
1269 max_ephemeral_storage_bytes: 0,
1270 architecture: Some(Architecture::X86_64),
1271 gpu: None,
1272 nested_virt: true,
1273 };
1274
1275 let selection = select_instance_type(Platform::Gcp, &req).unwrap();
1276 assert_eq!(selection.instance_type, "n2-standard-8");
1277 assert!(find_instance_type(Platform::Gcp, selection.instance_type)
1278 .unwrap()
1279 .is_nested_virt_capable());
1280 }
1281
1282 #[test]
1283 fn test_select_azure_picks_dsv5_when_nested_virt_required() {
1284 let req = WorkloadRequirements {
1285 total_cpu_at_desired: 4.0,
1286 total_memory_bytes_at_desired: 8 * GI,
1287 total_cpu_at_max: 4.0,
1288 total_memory_bytes_at_max: 8 * GI,
1289 max_cpu_per_container: 4.0,
1290 max_memory_per_container: 8 * GI,
1291 max_ephemeral_storage_bytes: 0,
1292 architecture: Some(Architecture::X86_64),
1293 gpu: None,
1294 nested_virt: true,
1295 };
1296
1297 let selection = select_instance_type(Platform::Azure, &req).unwrap();
1298 assert_eq!(selection.instance_type, "Standard_D8s_v5");
1299 assert!(find_instance_type(Platform::Azure, selection.instance_type)
1300 .unwrap()
1301 .is_nested_virt_capable());
1302 }
1303
1304 #[test]
1305 fn test_select_aws_defaults_to_image_target_architecture() {
1306 let req = WorkloadRequirements {
1307 total_cpu_at_desired: 4.0,
1308 total_memory_bytes_at_desired: 8 * GI,
1309 total_cpu_at_max: 4.0,
1310 total_memory_bytes_at_max: 8 * GI,
1311 max_cpu_per_container: 4.0,
1312 max_memory_per_container: 8 * GI,
1313 max_ephemeral_storage_bytes: 10 * GI,
1314 gpu: None,
1315 architecture: None,
1316 nested_virt: false,
1317 };
1318 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1319 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1320 assert_eq!(spec.architecture, Architecture::Arm64);
1321 }
1322
1323 #[test]
1324 fn test_cloud_defaults_match_image_target_architectures() {
1325 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1326 let target = BinaryTarget::defaults_for_platform(platform)
1327 .into_iter()
1328 .next()
1329 .expect("managed cloud should have a default image target");
1330 let image_architecture = match target.oci_arch() {
1331 "arm64" => Architecture::Arm64,
1332 "amd64" => Architecture::X86_64,
1333 architecture => {
1334 panic!("unsupported managed-cloud image architecture {architecture}")
1335 }
1336 };
1337
1338 assert_eq!(default_architecture(platform), Some(image_architecture));
1339 }
1340 }
1341
1342 #[test]
1344 fn test_select_aws_uses_graviton_for_explicit_arm64() {
1345 let req = WorkloadRequirements {
1346 total_cpu_at_desired: 4.0,
1347 total_memory_bytes_at_desired: 8 * GI,
1348 total_cpu_at_max: 4.0,
1349 total_memory_bytes_at_max: 8 * GI,
1350 max_cpu_per_container: 4.0,
1351 max_memory_per_container: 8 * GI,
1352 max_ephemeral_storage_bytes: 10 * GI,
1353 gpu: None,
1354 architecture: Some(Architecture::Arm64),
1355 nested_virt: false,
1356 };
1357 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1358 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1359 assert_eq!(spec.architecture, Architecture::Arm64);
1360 }
1361
1362 #[test]
1363 fn test_select_rejects_explicit_architecture_missing_from_cloud_catalog() {
1364 let req = WorkloadRequirements {
1365 total_cpu_at_desired: 1.0,
1366 total_memory_bytes_at_desired: 2 * GI,
1367 total_cpu_at_max: 1.0,
1368 total_memory_bytes_at_max: 2 * GI,
1369 max_cpu_per_container: 1.0,
1370 max_memory_per_container: 2 * GI,
1371 max_ephemeral_storage_bytes: 10 * GI,
1372 gpu: None,
1373 architecture: Some(Architecture::Arm64),
1374 nested_virt: false,
1375 };
1376
1377 let error = select_instance_type(Platform::Gcp, &req)
1378 .expect_err("GCP catalog has no ARM64 machine");
1379
1380 assert!(error.contains("architecture Arm64 is unavailable"));
1381 }
1382
1383 #[test]
1384 fn test_profile_has_required_fields() {
1385 let req = WorkloadRequirements {
1386 total_cpu_at_desired: 4.0,
1387 total_memory_bytes_at_desired: 16 * GI,
1388 total_cpu_at_max: 4.0,
1389 total_memory_bytes_at_max: 16 * GI,
1390 max_cpu_per_container: 1.0,
1391 max_memory_per_container: 4 * GI,
1392 max_ephemeral_storage_bytes: 10 * GI,
1393 gpu: None,
1394 architecture: None,
1395 nested_virt: false,
1396 };
1397 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1398 assert!(!sel.profile.cpu.is_empty());
1399 assert!(sel.profile.memory_bytes > 0);
1400 assert!(sel.profile.ephemeral_storage_bytes > 0);
1401 }
1402
1403 #[test]
1404 fn test_error_for_unsupported_gpu_type() {
1405 let req = WorkloadRequirements {
1406 total_cpu_at_desired: 8.0,
1407 total_memory_bytes_at_desired: 32 * GI,
1408 total_cpu_at_max: 8.0,
1409 total_memory_bytes_at_max: 32 * GI,
1410 max_cpu_per_container: 4.0,
1411 max_memory_per_container: 16 * GI,
1412 max_ephemeral_storage_bytes: 10 * GI,
1413 gpu: Some(GpuSpec {
1414 gpu_type: "amd-mi300".to_string(),
1415 count: 1,
1416 }),
1417 architecture: None,
1418 nested_virt: false,
1419 };
1420 let result = select_instance_type(Platform::Aws, &req);
1421 assert!(result.is_err());
1422 }
1423
1424 type SelectionShape = (
1427 &'static str,
1428 f64,
1429 u64,
1430 u32,
1431 u32,
1432 u64,
1433 Option<(&'static str, u32)>,
1434 );
1435
1436 #[rustfmt::skip]
1437 const SELECTION_SHAPES: &[SelectionShape] = &[
1438 ("tiny", 0.25, 512 * MI, 1, 1, 0, None),
1439 ("burstable", 1.0, 2 * GI, 1, 1, 10 * GI, None),
1440 ("burstable-big-memory", 0.5, 8 * GI, 1, 1, 0, None),
1441 ("small", 1.0, 2 * GI, 2, 4, 10 * GI, None),
1442 ("medium", 4.0, 8 * GI, 1, 1, 10 * GI, None),
1443 ("general", 2.0, 8 * GI, 4, 10, 20 * GI, None),
1444 ("many-replicas", 2.0, 4 * GI, 35, 35, 20 * GI, None),
1445 ("large-replica", 12.0, 48 * GI, 1, 2, 20 * GI, None),
1446 ("huge-replica", 60.0, 240 * GI, 1, 1, 20 * GI, None),
1447 ("memory-heavy", 4.0, 60 * GI, 2, 4, 20 * GI, None),
1448 ("storage-300g", 2.0, 8 * GI, 2, 2, 300 * GI, None),
1449 ("storage-1t", 4.0, 16 * GI, 2, 4, 1000 * GI, None),
1450 ("storage-2t", 4.0, 16 * GI, 2, 4, 2000 * GI, None),
1451 ("storage-8t", 2.0, 8 * GI, 2, 2, 8000 * GI, None),
1452 ("storage-100t", 2.0, 8 * GI, 1, 1, 100 * TI, None),
1453 ("gpu-t4", 2.0, 8 * GI, 1, 1, 10 * GI, Some(("nvidia-t4", 1))),
1454 ("gpu-a10g", 2.0, 8 * GI, 1, 1, 10 * GI, Some(("nvidia-a10g", 1))),
1455 ("gpu-a100", 4.0, 16 * GI, 1, 1, 10 * GI, Some(("nvidia-a100", 1))),
1456 ("gpu-a100x8", 16.0, 128 * GI, 1, 1, 10 * GI, Some(("nvidia-a100", 8))),
1457 ("gpu-h100x8", 16.0, 128 * GI, 1, 1, 10 * GI, Some(("nvidia-h100", 8))),
1458 ];
1459
1460 fn render_selection_matrix() -> String {
1463 let mut table = String::new();
1464 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1465 for architecture in [None, Some(Architecture::Arm64), Some(Architecture::X86_64)] {
1466 for nested_virt in [false, true] {
1467 for &(shape, cpu, memory, desired, max, storage, gpu) in SELECTION_SHAPES {
1468 let requirements = WorkloadRequirements {
1469 total_cpu_at_desired: cpu * desired as f64,
1470 total_memory_bytes_at_desired: memory * desired as u64,
1471 total_cpu_at_max: cpu * max as f64,
1472 total_memory_bytes_at_max: memory * max as u64,
1473 max_cpu_per_container: cpu,
1474 max_memory_per_container: memory,
1475 max_ephemeral_storage_bytes: storage,
1476 gpu: gpu.map(|(gpu_type, count)| GpuSpec {
1477 gpu_type: gpu_type.to_string(),
1478 count,
1479 }),
1480 architecture,
1481 nested_virt,
1482 };
1483 let pick = match select_instance_type(platform, &requirements) {
1484 Ok(selection) => format!(
1485 "{} min={} max={} memory={} storage={}",
1486 selection.instance_type,
1487 selection.min_machines,
1488 selection.max_machines,
1489 selection.profile.memory_bytes,
1490 selection.profile.ephemeral_storage_bytes
1491 ),
1492 Err(error) => format!("error: {error}"),
1493 };
1494 let architecture = architecture
1495 .map_or("default".to_string(), |architecture| {
1496 format!("{architecture:?}")
1497 });
1498 table.push_str(&format!(
1499 "{platform} {architecture} nested={nested_virt} {shape} -> {pick}\n"
1500 ));
1501 }
1502 }
1503 }
1504 }
1505 table
1506 }
1507
1508 #[rustfmt::skip]
1512 const ACCEPTED_PICK_CHANGES: &[(&str, &str)] = &[
1513 (
1514 "aws default nested=false gpu-t4",
1515 "error: Arm64: architecture Arm64 is unavailable for this workload on platform aws; \
1516 X86_64: no instance type for GPU type 'nvidia-t4' x1 on platform aws",
1517 ),
1518 (
1519 "aws default nested=false gpu-a10g",
1520 "g5.xlarge min=1 max=1 memory=17179869184 storage=268435456000",
1521 ),
1522 (
1523 "aws X86_64 nested=false gpu-t4",
1524 "error: no instance type for GPU type 'nvidia-t4' x1 on platform aws",
1525 ),
1526 (
1527 "aws X86_64 nested=false gpu-a10g",
1528 "g5.xlarge min=1 max=1 memory=17179869184 storage=268435456000",
1529 ),
1530 ];
1531
1532 #[test]
1536 fn automatic_picks_match_the_baseline_except_accepted_changes() {
1537 let baseline = include_str!("instance_catalog/auto_selection_baseline.txt");
1538 let rendered = render_selection_matrix();
1539 assert_eq!(baseline.lines().count(), rendered.lines().count());
1540 fn split(line: &str) -> (&str, &str) {
1541 line.split_once(" -> ").unwrap_or((line, ""))
1542 }
1543 for (before, now) in baseline.lines().zip(rendered.lines()) {
1544 let ((case, _), (now_case, pick)) = (split(before), split(now));
1545 assert_eq!(case, now_case);
1546 match ACCEPTED_PICK_CHANGES
1547 .iter()
1548 .find(|(accepted, _)| *accepted == case)
1549 {
1550 Some((_, accepted_pick)) => {
1551 assert_ne!(
1552 before, now,
1553 "{case} no longer changes; drop it from the list"
1554 );
1555 assert_eq!(pick, *accepted_pick, "{case}");
1556 }
1557 None => assert_eq!(before, now, "automatic pick changed"),
1558 }
1559 }
1560 }
1561
1562 type PinnedMachine = (
1563 Platform,
1564 &'static str,
1565 u32,
1566 u64,
1567 u64,
1568 Option<(&'static str, u32)>,
1569 bool,
1570 );
1571
1572 #[rustfmt::skip]
1575 const PINNED_MACHINES: &[PinnedMachine] = &[
1576 (Platform::Aws, "t4g.micro", 2, GI, 20 * GI, None, false),
1577 (Platform::Aws, "t4g.small", 2, 2 * GI, 20 * GI, None, false),
1578 (Platform::Aws, "t4g.medium", 2, 4 * GI, 20 * GI, None, false),
1579 (Platform::Aws, "t4g.large", 2, 8 * GI, 20 * GI, None, false),
1580 (Platform::Aws, "t3.xlarge", 4, 16 * GI, 20 * GI, None, false),
1581 (Platform::Aws, "t4g.xlarge", 4, 16 * GI, 20 * GI, None, false),
1582 (Platform::Aws, "m7g.medium", 1, 4 * GI, 20 * GI, None, false),
1583 (Platform::Aws, "m7i.large", 2, 8 * GI, 20 * GI, None, false),
1584 (Platform::Aws, "m7g.large", 2, 8 * GI, 20 * GI, None, false),
1585 (Platform::Aws, "m8i.large", 2, 8 * GI, 20 * GI, None, true),
1586 (Platform::Aws, "m7i.xlarge", 4, 16 * GI, 20 * GI, None, false),
1587 (Platform::Aws, "m7g.xlarge", 4, 16 * GI, 20 * GI, None, false),
1588 (Platform::Aws, "m8i.xlarge", 4, 16 * GI, 20 * GI, None, true),
1589 (Platform::Aws, "m7i.2xlarge", 8, 32 * GI, 20 * GI, None, false),
1590 (Platform::Aws, "m7g.2xlarge", 8, 32 * GI, 20 * GI, None, false),
1591 (Platform::Aws, "m8i.2xlarge", 8, 32 * GI, 20 * GI, None, true),
1592 (Platform::Aws, "m7i.4xlarge", 16, 64 * GI, 20 * GI, None, false),
1593 (Platform::Aws, "m7g.4xlarge", 16, 64 * GI, 20 * GI, None, false),
1594 (Platform::Aws, "m8i.4xlarge", 16, 64 * GI, 20 * GI, None, true),
1595 (Platform::Aws, "c7g.medium", 1, 2 * GI, 20 * GI, None, false),
1596 (Platform::Aws, "c8g.medium", 1, 2 * GI, 20 * GI, None, false),
1597 (Platform::Aws, "c7g.large", 2, 4 * GI, 20 * GI, None, false),
1598 (Platform::Aws, "c8g.large", 2, 4 * GI, 20 * GI, None, false),
1599 (Platform::Aws, "c8i.large", 2, 4 * GI, 20 * GI, None, true),
1600 (Platform::Aws, "c7g.xlarge", 4, 8 * GI, 20 * GI, None, false),
1601 (Platform::Aws, "c8g.xlarge", 4, 8 * GI, 20 * GI, None, false),
1602 (Platform::Aws, "c8i.xlarge", 4, 8 * GI, 20 * GI, None, true),
1603 (Platform::Aws, "c7g.2xlarge", 8, 16 * GI, 20 * GI, None, false),
1604 (Platform::Aws, "c8g.2xlarge", 8, 16 * GI, 20 * GI, None, false),
1605 (Platform::Aws, "c8i.2xlarge", 8, 16 * GI, 20 * GI, None, true),
1606 (Platform::Aws, "c7g.4xlarge", 16, 32 * GI, 20 * GI, None, false),
1607 (Platform::Aws, "c8g.4xlarge", 16, 32 * GI, 20 * GI, None, false),
1608 (Platform::Aws, "c8i.4xlarge", 16, 32 * GI, 20 * GI, None, true),
1609 (Platform::Aws, "r7g.medium", 1, 8 * GI, 20 * GI, None, false),
1610 (Platform::Aws, "r7g.large", 2, 16 * GI, 20 * GI, None, false),
1611 (Platform::Aws, "r6a.large", 2, 16 * GI, 20 * GI, None, false),
1612 (Platform::Aws, "r7i.large", 2, 16 * GI, 20 * GI, None, false),
1613 (Platform::Aws, "r7g.xlarge", 4, 32 * GI, 20 * GI, None, false),
1614 (Platform::Aws, "r6a.xlarge", 4, 32 * GI, 20 * GI, None, false),
1615 (Platform::Aws, "r7i.xlarge", 4, 32 * GI, 20 * GI, None, false),
1616 (Platform::Aws, "r7g.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1617 (Platform::Aws, "r6a.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1618 (Platform::Aws, "r7i.2xlarge", 8, 64 * GI, 20 * GI, None, false),
1619 (Platform::Aws, "r7g.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1620 (Platform::Aws, "r6a.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1621 (Platform::Aws, "r7i.4xlarge", 16, 128 * GI, 20 * GI, None, false),
1622 (Platform::Aws, "i4i.xlarge", 4, 32 * GI, 937 * GI, None, false),
1623 (Platform::Aws, "i4i.2xlarge", 8, 64 * GI, 1875 * GI, None, false),
1624 (Platform::Aws, "i4i.4xlarge", 16, 128 * GI, 3750 * GI, None, false),
1625 (Platform::Aws, "i4i.8xlarge", 32, 256 * GI, 7500 * GI, None, false),
1626 (Platform::Aws, "g5.xlarge", 4, 16 * GI, 250 * GI, Some(("nvidia-t4", 1)), false),
1627 (Platform::Aws, "g5.2xlarge", 8, 32 * GI, 450 * GI, Some(("nvidia-t4", 1)), false),
1628 (Platform::Aws, "p4d.24xlarge", 96, 1152 * GI, 8000 * GI, Some(("nvidia-a100", 8)), false),
1629 (Platform::Aws, "p5.48xlarge", 192, 2048 * GI, 8000 * GI, Some(("nvidia-h100", 8)), false),
1630 (Platform::Gcp, "e2-micro", 2, GI, 20 * GI, None, false),
1631 (Platform::Gcp, "e2-small", 2, 2 * GI, 20 * GI, None, false),
1632 (Platform::Gcp, "e2-medium", 2, 4 * GI, 20 * GI, None, false),
1633 (Platform::Gcp, "n2-standard-2", 2, 8 * GI, 20 * GI, None, true),
1634 (Platform::Gcp, "n2-standard-4", 4, 16 * GI, 20 * GI, None, true),
1635 (Platform::Gcp, "n2-standard-8", 8, 32 * GI, 20 * GI, None, true),
1636 (Platform::Gcp, "n2-standard-16", 16, 64 * GI, 20 * GI, None, true),
1637 (Platform::Gcp, "c3-standard-4", 4, 8 * GI, 20 * GI, None, false),
1638 (Platform::Gcp, "c3-standard-8", 8, 16 * GI, 20 * GI, None, false),
1639 (Platform::Gcp, "n2-highmem-2", 2, 16 * GI, 20 * GI, None, false),
1640 (Platform::Gcp, "n2-highmem-4", 4, 32 * GI, 20 * GI, None, false),
1641 (Platform::Gcp, "n2-highmem-8", 8, 64 * GI, 20 * GI, None, false),
1642 (Platform::Gcp, "n2-highmem-16", 16, 128 * GI, 20 * GI, None, false),
1643 (Platform::Gcp, "n2-highmem-32", 32, 256 * GI, 20 * GI, None, false),
1644 (Platform::Gcp, "c3d-standard-8", 8, 32 * GI, 480 * GI, None, false),
1645 (Platform::Gcp, "c3d-standard-16", 16, 64 * GI, 960 * GI, None, false),
1646 (Platform::Gcp, "c3d-standard-30", 30, 120 * GI, 1920 * GI, None, false),
1647 (Platform::Gcp, "n1-standard-4-t4", 4, 15 * GI, 100 * GI, Some(("nvidia-t4", 1)), false),
1648 (Platform::Gcp, "a2-highgpu-1g", 12, 85 * GI, 100 * GI, Some(("nvidia-a100", 1)), false),
1649 (Platform::Gcp, "a2-highgpu-8g", 96, 1360 * GI, 100 * GI, Some(("nvidia-a100", 8)), false),
1650 (Platform::Gcp, "a3-highgpu-8g", 208, 1872 * GI, 100 * GI, Some(("nvidia-h100", 8)), false),
1651 (Platform::Azure, "Standard_B1s", 1, GI, 20 * GI, None, false),
1652 (Platform::Azure, "Standard_B2s", 2, 4 * GI, 20 * GI, None, false),
1653 (Platform::Azure, "Standard_B2ms", 2, 8 * GI, 20 * GI, None, false),
1654 (Platform::Azure, "Standard_B4ms", 4, 16 * GI, 20 * GI, None, false),
1655 (Platform::Azure, "Standard_D2s_v4", 2, 8 * GI, 20 * GI, None, false),
1656 (Platform::Azure, "Standard_D2s_v5", 2, 8 * GI, 20 * GI, None, true),
1657 (Platform::Azure, "Standard_D4s_v5", 4, 16 * GI, 20 * GI, None, true),
1658 (Platform::Azure, "Standard_D8s_v5", 8, 32 * GI, 20 * GI, None, true),
1659 (Platform::Azure, "Standard_D16s_v5", 16, 64 * GI, 20 * GI, None, true),
1660 (Platform::Azure, "Standard_F2s_v2", 2, 4 * GI, 20 * GI, None, true),
1661 (Platform::Azure, "Standard_F4s_v2", 4, 8 * GI, 20 * GI, None, true),
1662 (Platform::Azure, "Standard_F8s_v2", 8, 16 * GI, 20 * GI, None, true),
1663 (Platform::Azure, "Standard_F16s_v2", 16, 32 * GI, 20 * GI, None, true),
1664 (Platform::Azure, "Standard_E2s_v5", 2, 16 * GI, 20 * GI, None, true),
1665 (Platform::Azure, "Standard_E4s_v5", 4, 32 * GI, 20 * GI, None, true),
1666 (Platform::Azure, "Standard_E8s_v5", 8, 64 * GI, 20 * GI, None, true),
1667 (Platform::Azure, "Standard_E16s_v5", 16, 128 * GI, 20 * GI, None, true),
1668 (Platform::Azure, "Standard_L8s_v3", 8, 64 * GI, 1788 * GI, None, false),
1669 (Platform::Azure, "Standard_L16s_v3", 16, 128 * GI, 3576 * GI, None, false),
1670 (Platform::Azure, "Standard_L32s_v3", 32, 256 * GI, 7154 * GI, None, false),
1671 (Platform::Azure, "Standard_NC4as_T4_v3", 4, 28 * GI, 176 * GI, Some(("nvidia-t4", 1)), false),
1672 (Platform::Azure, "Standard_NC24ads_A100_v4", 24, 220 * GI, 958 * GI, Some(("nvidia-a100", 1)), false),
1673 (Platform::Azure, "Standard_NC96ads_A100_v4", 96, 880 * GI, 3916 * GI, Some(("nvidia-a100", 4)), false),
1674 (Platform::Azure, "Standard_ND96isr_H100_v5", 96, 1900 * GI, 1000 * GI, Some(("nvidia-h100", 8)), false),
1675 ];
1676
1677 const GPU_TYPE_CORRECTIONS: &[(&str, &str)] =
1679 &[("g5.xlarge", "nvidia-a10g"), ("g5.2xlarge", "nvidia-a10g")];
1680
1681 #[test]
1682 fn every_pinned_machine_is_still_in_the_catalog_with_its_values() {
1683 for &(platform, name, vcpu, memory_bytes, storage_bytes, gpu, nested_virt) in
1684 PINNED_MACHINES
1685 {
1686 let spec = find_instance_type(platform, name)
1687 .unwrap_or_else(|| panic!("{platform} machine {name} left the catalog"));
1688 let gpu = gpu.map(|(gpu_type, count)| {
1689 let corrected = GPU_TYPE_CORRECTIONS
1690 .iter()
1691 .find(|(corrected, _)| *corrected == name);
1692 (corrected.map_or(gpu_type, |(_, gpu_type)| *gpu_type), count)
1693 });
1694 assert_eq!(
1695 (
1696 spec.vcpu,
1697 spec.memory_bytes,
1698 spec.ephemeral_storage_bytes,
1699 spec.gpu.map(|gpu| (gpu.gpu_type, gpu.count)),
1700 spec.nested_virt,
1701 ),
1702 (vcpu, memory_bytes, storage_bytes, gpu, nested_virt),
1703 "{platform} {name}"
1704 );
1705 }
1706 }
1707
1708 #[test]
1709 fn automatic_selection_considers_exactly_the_pinned_machines() {
1710 let pinned: std::collections::HashSet<_> = PINNED_MACHINES
1711 .iter()
1712 .map(|&(platform, name, ..)| (platform, name))
1713 .collect();
1714 let auto_select: std::collections::HashSet<_> = CATALOG
1715 .iter()
1716 .filter(|spec| spec.auto_select)
1717 .map(|spec| (spec.platform, spec.name))
1718 .collect();
1719 assert_eq!(auto_select, pinned);
1720 }
1721
1722 #[test]
1723 fn catalog_lists_no_bare_metal_machine() {
1724 for spec in CATALOG {
1725 assert!(!spec.name.contains("metal"), "{}", spec.name);
1726 }
1727 }
1728
1729 #[test]
1731 fn test_catalog_instance_types_sorted_by_vcpu_within_series() {
1732 let mut finished = std::collections::HashSet::new();
1733 for window in CATALOG.windows(2) {
1734 let (previous, next) = (&window[0], &window[1]);
1735 if (previous.platform, previous.series) == (next.platform, next.series) {
1736 assert!(
1737 previous.vcpu <= next.vcpu,
1738 "{} ({}) before {} ({})",
1739 previous.name,
1740 previous.vcpu,
1741 next.name,
1742 next.vcpu
1743 );
1744 } else {
1745 finished.insert((previous.platform, previous.series));
1746 assert!(
1747 !finished.contains(&(next.platform, next.series)),
1748 "series {} is split",
1749 next.series
1750 );
1751 }
1752 }
1753 }
1754}