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 = select_instance_type_for_architecture(
360 platform,
361 &with_architecture(default),
362 );
363 if selection.is_ok() {
364 return selection;
365 }
366 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 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 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 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 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 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 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 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 candidates
564 .iter()
565 .filter(|spec| spec.vcpu <= vcpu_cap)
566 .max_by_key(|spec| spec.vcpu)
567 })
568 .or_else(|| {
569 candidates.iter().min_by_key(|spec| spec.vcpu)
571 })
572 .ok_or_else(|| format!("no instance types available for platform {platform}"))?;
573
574 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
586pub fn select_family(requirements: &WorkloadRequirements) -> InstanceFamily {
592 if requirements.gpu.is_some() {
594 return InstanceFamily::GpuCompute;
595 }
596
597 if requirements.max_ephemeral_storage_bytes > STORAGE_OPTIMIZED_THRESHOLD {
599 return InstanceFamily::StorageOptimized;
600 }
601
602 if requirements.total_cpu_at_max < 2.0 {
604 return InstanceFamily::Burstable;
605 }
606
607 InstanceFamily::GeneralPurpose
609}
610
611fn 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 cpu_machines
623 .max(mem_machines)
624 .max(1)
625 .min(MAX_MACHINES_PER_CLUSTER)
626}
627
628fn 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#[cfg(test)]
685mod tests {
686 use super::*;
687 use crate::BinaryTarget;
688
689 #[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 #[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 #[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 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 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 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 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 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 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 assert!(sel.max_machines > 1);
1203 }
1204
1205 #[test]
1206 fn test_larger_autoscaled_workload_gets_reasonable_instance() {
1207 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, total_memory_bytes_at_max: 140 * GI, 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 assert_eq!(sel.instance_type, "n2-standard-8");
1224 assert!(sel.max_machines >= 2);
1225 }
1226
1227 #[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 #[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 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 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 #[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 #[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 #[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 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 #[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}