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 if let Some(num) = s.strip_suffix("Ti") {
50 let v: f64 = num
51 .parse()
52 .map_err(|_| format!("invalid memory value: '{s}'"))?;
53 return Ok((v * 1024.0 * 1024.0 * 1024.0 * 1024.0) as u64);
54 }
55 if let Some(num) = s.strip_suffix("Gi") {
56 let v: f64 = num
57 .parse()
58 .map_err(|_| format!("invalid memory value: '{s}'"))?;
59 return Ok((v * 1024.0 * 1024.0 * 1024.0) as u64);
60 }
61 if let Some(num) = s.strip_suffix("Mi") {
62 let v: f64 = num
63 .parse()
64 .map_err(|_| format!("invalid memory value: '{s}'"))?;
65 return Ok((v * 1024.0 * 1024.0) as u64);
66 }
67 if let Some(num) = s.strip_suffix("Ki") {
68 let v: f64 = num
69 .parse()
70 .map_err(|_| format!("invalid memory value: '{s}'"))?;
71 return Ok((v * 1024.0) as u64);
72 }
73
74 if let Some(num) = s.strip_suffix('T') {
76 let v: f64 = num
77 .parse()
78 .map_err(|_| format!("invalid memory value: '{s}'"))?;
79 return Ok((v * 1_000_000_000_000.0) as u64);
80 }
81 if let Some(num) = s.strip_suffix('G') {
82 let v: f64 = num
83 .parse()
84 .map_err(|_| format!("invalid memory value: '{s}'"))?;
85 return Ok((v * 1_000_000_000.0) as u64);
86 }
87 if let Some(num) = s.strip_suffix('M') {
88 let v: f64 = num
89 .parse()
90 .map_err(|_| format!("invalid memory value: '{s}'"))?;
91 return Ok((v * 1_000_000.0) as u64);
92 }
93 if let Some(num) = s.strip_suffix('k') {
94 let v: f64 = num
95 .parse()
96 .map_err(|_| format!("invalid memory value: '{s}'"))?;
97 return Ok((v * 1000.0) as u64);
98 }
99
100 s.parse()
102 .map_err(|_| format!("invalid memory value: '{s}'"))
103}
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub enum InstanceFamily {
112 Burstable,
113 GeneralPurpose,
114 ComputeOptimized,
115 MemoryOptimized,
116 StorageOptimized,
117 GpuCompute,
118}
119
120#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
122#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
123#[serde(rename_all = "snake_case")]
124pub enum Architecture {
125 Arm64,
126 X86_64,
127}
128
129pub fn default_architecture(platform: Platform) -> Option<Architecture> {
131 match platform {
132 Platform::Aws => Some(Architecture::Arm64),
133 Platform::Gcp | Platform::Azure => Some(Architecture::X86_64),
134 Platform::Kubernetes | Platform::Machines | Platform::Local | Platform::Test => None,
135 }
136}
137
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140pub struct CatalogGpu {
141 pub gpu_type: &'static str,
142 pub count: u32,
143}
144
145#[derive(Debug, Clone)]
149pub struct InstanceTypeSpec {
150 pub name: &'static str,
151 pub platform: Platform,
152 pub family: InstanceFamily,
153 pub architecture: Architecture,
154 pub vcpu: u32,
156 pub memory_bytes: u64,
158 pub ephemeral_storage_bytes: u64,
160 pub gpu: Option<CatalogGpu>,
162}
163
164impl InstanceTypeSpec {
165 pub fn has_configurable_ephemeral_storage(&self) -> bool {
168 matches!(
169 self.family,
170 InstanceFamily::Burstable
171 | InstanceFamily::GeneralPurpose
172 | InstanceFamily::ComputeOptimized
173 | InstanceFamily::MemoryOptimized
174 )
175 }
176
177 pub fn is_nested_virt_capable(&self) -> bool {
183 match self.platform {
184 Platform::Aws => {
185 let name = self.name;
186 name.starts_with("m8i.")
187 || name.starts_with("c8i.")
188 || name.starts_with("r8i.")
189 || name.starts_with("m8i-flex.")
190 || name.starts_with("c8i-flex.")
191 || name.starts_with("r8i-flex.")
192 }
193 Platform::Gcp => self.name.starts_with("n2-standard-"),
194 Platform::Azure => {
195 let name = self.name;
196 (name.starts_with("Standard_D") && name.ends_with("s_v5"))
197 || (name.starts_with("Standard_E") && name.ends_with("s_v5"))
198 || (name.starts_with("Standard_F") && name.ends_with("s_v2"))
199 }
200 _ => false,
201 }
202 }
203
204 pub fn to_machine_profile(&self) -> MachineProfile {
206 MachineProfile {
207 cpu: format!("{}.0", self.vcpu),
208 memory_bytes: self.memory_bytes,
209 ephemeral_storage_bytes: self.ephemeral_storage_bytes,
210 architecture: Some(self.architecture),
211 gpu: self.gpu.map(|g| GpuSpec {
212 gpu_type: g.gpu_type.to_string(),
213 count: g.count,
214 }),
215 }
216 }
217
218 pub fn to_machine_profile_for_storage(&self, requested_bytes: u64) -> MachineProfile {
222 let mut profile = self.to_machine_profile();
223 if self.has_configurable_ephemeral_storage() {
224 profile.ephemeral_storage_bytes = profile.ephemeral_storage_bytes.max(requested_bytes);
225 }
226 profile
227 }
228}
229
230const KI: u64 = 1024;
232const MI: u64 = KI * 1024;
233const GI: u64 = MI * 1024;
234const TI: u64 = GI * 1024;
235
236pub fn max_configurable_ephemeral_storage_bytes(platform: Platform) -> Option<u64> {
242 let disk_limit_gib = match platform {
243 Platform::Aws | Platform::Gcp => 64 * 1024,
245 Platform::Azure => 4_095,
247 Platform::Kubernetes | Platform::Machines | Platform::Local | Platform::Test => {
248 return None;
249 }
250 };
251
252 let max_requested_gib = (disk_limit_gib - 12) * 4 / 5;
257 Some(max_requested_gib * GI)
258}
259
260static CATALOG: &[InstanceTypeSpec] = &[
268 InstanceTypeSpec {
274 name: "t4g.micro",
275 platform: Platform::Aws,
276 family: InstanceFamily::Burstable,
277 architecture: Architecture::Arm64,
278 vcpu: 2,
279 memory_bytes: 1 * GI,
280 ephemeral_storage_bytes: 20 * GI,
281 gpu: None,
282 },
283 InstanceTypeSpec {
284 name: "t4g.small",
285 platform: Platform::Aws,
286 family: InstanceFamily::Burstable,
287 architecture: Architecture::Arm64,
288 vcpu: 2,
289 memory_bytes: 2 * GI,
290 ephemeral_storage_bytes: 20 * GI,
291 gpu: None,
292 },
293 InstanceTypeSpec {
294 name: "t4g.medium",
295 platform: Platform::Aws,
296 family: InstanceFamily::Burstable,
297 architecture: Architecture::Arm64,
298 vcpu: 2,
299 memory_bytes: 4 * GI,
300 ephemeral_storage_bytes: 20 * GI,
301 gpu: None,
302 },
303 InstanceTypeSpec {
304 name: "t4g.large",
305 platform: Platform::Aws,
306 family: InstanceFamily::Burstable,
307 architecture: Architecture::Arm64,
308 vcpu: 2,
309 memory_bytes: 8 * GI,
310 ephemeral_storage_bytes: 20 * GI,
311 gpu: None,
312 },
313 InstanceTypeSpec {
314 name: "t3.xlarge",
315 platform: Platform::Aws,
316 family: InstanceFamily::Burstable,
317 architecture: Architecture::X86_64,
318 vcpu: 4,
319 memory_bytes: 16 * GI,
320 ephemeral_storage_bytes: 20 * GI,
321 gpu: None,
322 },
323 InstanceTypeSpec {
324 name: "t4g.xlarge",
325 platform: Platform::Aws,
326 family: InstanceFamily::Burstable,
327 architecture: Architecture::Arm64,
328 vcpu: 4,
329 memory_bytes: 16 * GI,
330 ephemeral_storage_bytes: 20 * GI,
331 gpu: None,
332 },
333 InstanceTypeSpec {
335 name: "m7g.medium",
336 platform: Platform::Aws,
337 family: InstanceFamily::GeneralPurpose,
338 architecture: Architecture::Arm64,
339 vcpu: 1,
340 memory_bytes: 4 * GI,
341 ephemeral_storage_bytes: 20 * GI,
342 gpu: None,
343 },
344 InstanceTypeSpec {
345 name: "m7i.large",
346 platform: Platform::Aws,
347 family: InstanceFamily::GeneralPurpose,
348 architecture: Architecture::X86_64,
349 vcpu: 2,
350 memory_bytes: 8 * GI,
351 ephemeral_storage_bytes: 20 * GI,
352 gpu: None,
353 },
354 InstanceTypeSpec {
355 name: "m7g.large",
356 platform: Platform::Aws,
357 family: InstanceFamily::GeneralPurpose,
358 architecture: Architecture::Arm64,
359 vcpu: 2,
360 memory_bytes: 8 * GI,
361 ephemeral_storage_bytes: 20 * GI,
362 gpu: None,
363 },
364 InstanceTypeSpec {
372 name: "m8i.large",
373 platform: Platform::Aws,
374 family: InstanceFamily::GeneralPurpose,
375 architecture: Architecture::X86_64,
376 vcpu: 2,
377 memory_bytes: 8 * GI,
378 ephemeral_storage_bytes: 20 * GI,
379 gpu: None,
380 },
381 InstanceTypeSpec {
382 name: "m7i.xlarge",
383 platform: Platform::Aws,
384 family: InstanceFamily::GeneralPurpose,
385 architecture: Architecture::X86_64,
386 vcpu: 4,
387 memory_bytes: 16 * GI,
388 ephemeral_storage_bytes: 20 * GI,
389 gpu: None,
390 },
391 InstanceTypeSpec {
392 name: "m7g.xlarge",
393 platform: Platform::Aws,
394 family: InstanceFamily::GeneralPurpose,
395 architecture: Architecture::Arm64,
396 vcpu: 4,
397 memory_bytes: 16 * GI,
398 ephemeral_storage_bytes: 20 * GI,
399 gpu: None,
400 },
401 InstanceTypeSpec {
402 name: "m8i.xlarge",
403 platform: Platform::Aws,
404 family: InstanceFamily::GeneralPurpose,
405 architecture: Architecture::X86_64,
406 vcpu: 4,
407 memory_bytes: 16 * GI,
408 ephemeral_storage_bytes: 20 * GI,
409 gpu: None,
410 },
411 InstanceTypeSpec {
412 name: "m7i.2xlarge",
413 platform: Platform::Aws,
414 family: InstanceFamily::GeneralPurpose,
415 architecture: Architecture::X86_64,
416 vcpu: 8,
417 memory_bytes: 32 * GI,
418 ephemeral_storage_bytes: 20 * GI,
419 gpu: None,
420 },
421 InstanceTypeSpec {
422 name: "m7g.2xlarge",
423 platform: Platform::Aws,
424 family: InstanceFamily::GeneralPurpose,
425 architecture: Architecture::Arm64,
426 vcpu: 8,
427 memory_bytes: 32 * GI,
428 ephemeral_storage_bytes: 20 * GI,
429 gpu: None,
430 },
431 InstanceTypeSpec {
432 name: "m8i.2xlarge",
433 platform: Platform::Aws,
434 family: InstanceFamily::GeneralPurpose,
435 architecture: Architecture::X86_64,
436 vcpu: 8,
437 memory_bytes: 32 * GI,
438 ephemeral_storage_bytes: 20 * GI,
439 gpu: None,
440 },
441 InstanceTypeSpec {
442 name: "m7i.4xlarge",
443 platform: Platform::Aws,
444 family: InstanceFamily::GeneralPurpose,
445 architecture: Architecture::X86_64,
446 vcpu: 16,
447 memory_bytes: 64 * GI,
448 ephemeral_storage_bytes: 20 * GI,
449 gpu: None,
450 },
451 InstanceTypeSpec {
452 name: "m7g.4xlarge",
453 platform: Platform::Aws,
454 family: InstanceFamily::GeneralPurpose,
455 architecture: Architecture::Arm64,
456 vcpu: 16,
457 memory_bytes: 64 * GI,
458 ephemeral_storage_bytes: 20 * GI,
459 gpu: None,
460 },
461 InstanceTypeSpec {
462 name: "m8i.4xlarge",
463 platform: Platform::Aws,
464 family: InstanceFamily::GeneralPurpose,
465 architecture: Architecture::X86_64,
466 vcpu: 16,
467 memory_bytes: 64 * GI,
468 ephemeral_storage_bytes: 20 * GI,
469 gpu: None,
470 },
471 InstanceTypeSpec {
473 name: "c7g.medium",
474 platform: Platform::Aws,
475 family: InstanceFamily::ComputeOptimized,
476 architecture: Architecture::Arm64,
477 vcpu: 1,
478 memory_bytes: 2 * GI,
479 ephemeral_storage_bytes: 20 * GI,
480 gpu: None,
481 },
482 InstanceTypeSpec {
483 name: "c8g.medium",
484 platform: Platform::Aws,
485 family: InstanceFamily::ComputeOptimized,
486 architecture: Architecture::Arm64,
487 vcpu: 1,
488 memory_bytes: 2 * GI,
489 ephemeral_storage_bytes: 20 * GI,
490 gpu: None,
491 },
492 InstanceTypeSpec {
493 name: "c7g.large",
494 platform: Platform::Aws,
495 family: InstanceFamily::ComputeOptimized,
496 architecture: Architecture::Arm64,
497 vcpu: 2,
498 memory_bytes: 4 * GI,
499 ephemeral_storage_bytes: 20 * GI,
500 gpu: None,
501 },
502 InstanceTypeSpec {
503 name: "c8g.large",
504 platform: Platform::Aws,
505 family: InstanceFamily::ComputeOptimized,
506 architecture: Architecture::Arm64,
507 vcpu: 2,
508 memory_bytes: 4 * GI,
509 ephemeral_storage_bytes: 20 * GI,
510 gpu: None,
511 },
512 InstanceTypeSpec {
513 name: "c8i.large",
514 platform: Platform::Aws,
515 family: InstanceFamily::ComputeOptimized,
516 architecture: Architecture::X86_64,
517 vcpu: 2,
518 memory_bytes: 4 * GI,
519 ephemeral_storage_bytes: 20 * GI,
520 gpu: None,
521 },
522 InstanceTypeSpec {
523 name: "c7g.xlarge",
524 platform: Platform::Aws,
525 family: InstanceFamily::ComputeOptimized,
526 architecture: Architecture::Arm64,
527 vcpu: 4,
528 memory_bytes: 8 * GI,
529 ephemeral_storage_bytes: 20 * GI,
530 gpu: None,
531 },
532 InstanceTypeSpec {
533 name: "c8g.xlarge",
534 platform: Platform::Aws,
535 family: InstanceFamily::ComputeOptimized,
536 architecture: Architecture::Arm64,
537 vcpu: 4,
538 memory_bytes: 8 * GI,
539 ephemeral_storage_bytes: 20 * GI,
540 gpu: None,
541 },
542 InstanceTypeSpec {
543 name: "c8i.xlarge",
544 platform: Platform::Aws,
545 family: InstanceFamily::ComputeOptimized,
546 architecture: Architecture::X86_64,
547 vcpu: 4,
548 memory_bytes: 8 * GI,
549 ephemeral_storage_bytes: 20 * GI,
550 gpu: None,
551 },
552 InstanceTypeSpec {
553 name: "c7g.2xlarge",
554 platform: Platform::Aws,
555 family: InstanceFamily::ComputeOptimized,
556 architecture: Architecture::Arm64,
557 vcpu: 8,
558 memory_bytes: 16 * GI,
559 ephemeral_storage_bytes: 20 * GI,
560 gpu: None,
561 },
562 InstanceTypeSpec {
563 name: "c8g.2xlarge",
564 platform: Platform::Aws,
565 family: InstanceFamily::ComputeOptimized,
566 architecture: Architecture::Arm64,
567 vcpu: 8,
568 memory_bytes: 16 * GI,
569 ephemeral_storage_bytes: 20 * GI,
570 gpu: None,
571 },
572 InstanceTypeSpec {
573 name: "c8i.2xlarge",
574 platform: Platform::Aws,
575 family: InstanceFamily::ComputeOptimized,
576 architecture: Architecture::X86_64,
577 vcpu: 8,
578 memory_bytes: 16 * GI,
579 ephemeral_storage_bytes: 20 * GI,
580 gpu: None,
581 },
582 InstanceTypeSpec {
583 name: "c7g.4xlarge",
584 platform: Platform::Aws,
585 family: InstanceFamily::ComputeOptimized,
586 architecture: Architecture::Arm64,
587 vcpu: 16,
588 memory_bytes: 32 * GI,
589 ephemeral_storage_bytes: 20 * GI,
590 gpu: None,
591 },
592 InstanceTypeSpec {
593 name: "c8g.4xlarge",
594 platform: Platform::Aws,
595 family: InstanceFamily::ComputeOptimized,
596 architecture: Architecture::Arm64,
597 vcpu: 16,
598 memory_bytes: 32 * GI,
599 ephemeral_storage_bytes: 20 * GI,
600 gpu: None,
601 },
602 InstanceTypeSpec {
603 name: "c8i.4xlarge",
604 platform: Platform::Aws,
605 family: InstanceFamily::ComputeOptimized,
606 architecture: Architecture::X86_64,
607 vcpu: 16,
608 memory_bytes: 32 * GI,
609 ephemeral_storage_bytes: 20 * GI,
610 gpu: None,
611 },
612 InstanceTypeSpec {
614 name: "r7g.medium",
615 platform: Platform::Aws,
616 family: InstanceFamily::MemoryOptimized,
617 architecture: Architecture::Arm64,
618 vcpu: 1,
619 memory_bytes: 8 * GI,
620 ephemeral_storage_bytes: 20 * GI,
621 gpu: None,
622 },
623 InstanceTypeSpec {
624 name: "r7g.large",
625 platform: Platform::Aws,
626 family: InstanceFamily::MemoryOptimized,
627 architecture: Architecture::Arm64,
628 vcpu: 2,
629 memory_bytes: 16 * GI,
630 ephemeral_storage_bytes: 20 * GI,
631 gpu: None,
632 },
633 InstanceTypeSpec {
634 name: "r7g.xlarge",
635 platform: Platform::Aws,
636 family: InstanceFamily::MemoryOptimized,
637 architecture: Architecture::Arm64,
638 vcpu: 4,
639 memory_bytes: 32 * GI,
640 ephemeral_storage_bytes: 20 * GI,
641 gpu: None,
642 },
643 InstanceTypeSpec {
644 name: "r7g.2xlarge",
645 platform: Platform::Aws,
646 family: InstanceFamily::MemoryOptimized,
647 architecture: Architecture::Arm64,
648 vcpu: 8,
649 memory_bytes: 64 * GI,
650 ephemeral_storage_bytes: 20 * GI,
651 gpu: None,
652 },
653 InstanceTypeSpec {
654 name: "r7g.4xlarge",
655 platform: Platform::Aws,
656 family: InstanceFamily::MemoryOptimized,
657 architecture: Architecture::Arm64,
658 vcpu: 16,
659 memory_bytes: 128 * GI,
660 ephemeral_storage_bytes: 20 * GI,
661 gpu: None,
662 },
663 InstanceTypeSpec {
665 name: "i4i.xlarge",
666 platform: Platform::Aws,
667 family: InstanceFamily::StorageOptimized,
668 architecture: Architecture::X86_64,
669 vcpu: 4,
670 memory_bytes: 32 * GI,
671 ephemeral_storage_bytes: 937 * GI,
672 gpu: None,
673 },
674 InstanceTypeSpec {
675 name: "i4i.2xlarge",
676 platform: Platform::Aws,
677 family: InstanceFamily::StorageOptimized,
678 architecture: Architecture::X86_64,
679 vcpu: 8,
680 memory_bytes: 64 * GI,
681 ephemeral_storage_bytes: 1875 * GI,
682 gpu: None,
683 },
684 InstanceTypeSpec {
685 name: "i4i.4xlarge",
686 platform: Platform::Aws,
687 family: InstanceFamily::StorageOptimized,
688 architecture: Architecture::X86_64,
689 vcpu: 16,
690 memory_bytes: 128 * GI,
691 ephemeral_storage_bytes: 3750 * GI,
692 gpu: None,
693 },
694 InstanceTypeSpec {
695 name: "i4i.8xlarge",
696 platform: Platform::Aws,
697 family: InstanceFamily::StorageOptimized,
698 architecture: Architecture::X86_64,
699 vcpu: 32,
700 memory_bytes: 256 * GI,
701 ephemeral_storage_bytes: 7500 * GI,
702 gpu: None,
703 },
704 InstanceTypeSpec {
706 name: "g5.xlarge",
707 platform: Platform::Aws,
708 family: InstanceFamily::GpuCompute,
709 architecture: Architecture::X86_64,
710 vcpu: 4,
711 memory_bytes: 16 * GI,
712 ephemeral_storage_bytes: 250 * GI,
713 gpu: Some(CatalogGpu {
714 gpu_type: "nvidia-t4",
715 count: 1,
716 }),
717 },
718 InstanceTypeSpec {
719 name: "g5.2xlarge",
720 platform: Platform::Aws,
721 family: InstanceFamily::GpuCompute,
722 architecture: Architecture::X86_64,
723 vcpu: 8,
724 memory_bytes: 32 * GI,
725 ephemeral_storage_bytes: 450 * GI,
726 gpu: Some(CatalogGpu {
727 gpu_type: "nvidia-t4",
728 count: 1,
729 }),
730 },
731 InstanceTypeSpec {
733 name: "p4d.24xlarge",
734 platform: Platform::Aws,
735 family: InstanceFamily::GpuCompute,
736 architecture: Architecture::X86_64,
737 vcpu: 96,
738 memory_bytes: 1152 * GI,
739 ephemeral_storage_bytes: 8000 * GI,
740 gpu: Some(CatalogGpu {
741 gpu_type: "nvidia-a100",
742 count: 8,
743 }),
744 },
745 InstanceTypeSpec {
747 name: "p5.48xlarge",
748 platform: Platform::Aws,
749 family: InstanceFamily::GpuCompute,
750 architecture: Architecture::X86_64,
751 vcpu: 192,
752 memory_bytes: 2048 * GI,
753 ephemeral_storage_bytes: 8000 * GI,
754 gpu: Some(CatalogGpu {
755 gpu_type: "nvidia-h100",
756 count: 8,
757 }),
758 },
759 InstanceTypeSpec {
765 name: "e2-micro",
766 platform: Platform::Gcp,
767 family: InstanceFamily::Burstable,
768 architecture: Architecture::X86_64,
769 vcpu: 2,
770 memory_bytes: 1 * GI,
771 ephemeral_storage_bytes: 20 * GI,
772 gpu: None,
773 },
774 InstanceTypeSpec {
775 name: "e2-small",
776 platform: Platform::Gcp,
777 family: InstanceFamily::Burstable,
778 architecture: Architecture::X86_64,
779 vcpu: 2,
780 memory_bytes: 2 * GI,
781 ephemeral_storage_bytes: 20 * GI,
782 gpu: None,
783 },
784 InstanceTypeSpec {
785 name: "e2-medium",
786 platform: Platform::Gcp,
787 family: InstanceFamily::Burstable,
788 architecture: Architecture::X86_64,
789 vcpu: 2,
790 memory_bytes: 4 * GI,
791 ephemeral_storage_bytes: 20 * GI,
792 gpu: None,
793 },
794 InstanceTypeSpec {
796 name: "n2-standard-2",
797 platform: Platform::Gcp,
798 family: InstanceFamily::GeneralPurpose,
799 architecture: Architecture::X86_64,
800 vcpu: 2,
801 memory_bytes: 8 * GI,
802 ephemeral_storage_bytes: 20 * GI,
803 gpu: None,
804 },
805 InstanceTypeSpec {
806 name: "n2-standard-4",
807 platform: Platform::Gcp,
808 family: InstanceFamily::GeneralPurpose,
809 architecture: Architecture::X86_64,
810 vcpu: 4,
811 memory_bytes: 16 * GI,
812 ephemeral_storage_bytes: 20 * GI,
813 gpu: None,
814 },
815 InstanceTypeSpec {
816 name: "n2-standard-8",
817 platform: Platform::Gcp,
818 family: InstanceFamily::GeneralPurpose,
819 architecture: Architecture::X86_64,
820 vcpu: 8,
821 memory_bytes: 32 * GI,
822 ephemeral_storage_bytes: 20 * GI,
823 gpu: None,
824 },
825 InstanceTypeSpec {
826 name: "n2-standard-16",
827 platform: Platform::Gcp,
828 family: InstanceFamily::GeneralPurpose,
829 architecture: Architecture::X86_64,
830 vcpu: 16,
831 memory_bytes: 64 * GI,
832 ephemeral_storage_bytes: 20 * GI,
833 gpu: None,
834 },
835 InstanceTypeSpec {
837 name: "c3-standard-4",
838 platform: Platform::Gcp,
839 family: InstanceFamily::ComputeOptimized,
840 architecture: Architecture::X86_64,
841 vcpu: 4,
842 memory_bytes: 8 * GI,
843 ephemeral_storage_bytes: 20 * GI,
844 gpu: None,
845 },
846 InstanceTypeSpec {
847 name: "c3-standard-8",
848 platform: Platform::Gcp,
849 family: InstanceFamily::ComputeOptimized,
850 architecture: Architecture::X86_64,
851 vcpu: 8,
852 memory_bytes: 16 * GI,
853 ephemeral_storage_bytes: 20 * GI,
854 gpu: None,
855 },
856 InstanceTypeSpec {
858 name: "n2-highmem-2",
859 platform: Platform::Gcp,
860 family: InstanceFamily::MemoryOptimized,
861 architecture: Architecture::X86_64,
862 vcpu: 2,
863 memory_bytes: 16 * GI,
864 ephemeral_storage_bytes: 20 * GI,
865 gpu: None,
866 },
867 InstanceTypeSpec {
868 name: "n2-highmem-4",
869 platform: Platform::Gcp,
870 family: InstanceFamily::MemoryOptimized,
871 architecture: Architecture::X86_64,
872 vcpu: 4,
873 memory_bytes: 32 * GI,
874 ephemeral_storage_bytes: 20 * GI,
875 gpu: None,
876 },
877 InstanceTypeSpec {
878 name: "n2-highmem-8",
879 platform: Platform::Gcp,
880 family: InstanceFamily::MemoryOptimized,
881 architecture: Architecture::X86_64,
882 vcpu: 8,
883 memory_bytes: 64 * GI,
884 ephemeral_storage_bytes: 20 * GI,
885 gpu: None,
886 },
887 InstanceTypeSpec {
888 name: "n2-highmem-16",
889 platform: Platform::Gcp,
890 family: InstanceFamily::MemoryOptimized,
891 architecture: Architecture::X86_64,
892 vcpu: 16,
893 memory_bytes: 128 * GI,
894 ephemeral_storage_bytes: 20 * GI,
895 gpu: None,
896 },
897 InstanceTypeSpec {
898 name: "n2-highmem-32",
899 platform: Platform::Gcp,
900 family: InstanceFamily::MemoryOptimized,
901 architecture: Architecture::X86_64,
902 vcpu: 32,
903 memory_bytes: 256 * GI,
904 ephemeral_storage_bytes: 20 * GI,
905 gpu: None,
906 },
907 InstanceTypeSpec {
909 name: "c3d-standard-8",
910 platform: Platform::Gcp,
911 family: InstanceFamily::StorageOptimized,
912 architecture: Architecture::X86_64,
913 vcpu: 8,
914 memory_bytes: 32 * GI,
915 ephemeral_storage_bytes: 480 * GI,
916 gpu: None,
917 },
918 InstanceTypeSpec {
919 name: "c3d-standard-16",
920 platform: Platform::Gcp,
921 family: InstanceFamily::StorageOptimized,
922 architecture: Architecture::X86_64,
923 vcpu: 16,
924 memory_bytes: 64 * GI,
925 ephemeral_storage_bytes: 960 * GI,
926 gpu: None,
927 },
928 InstanceTypeSpec {
929 name: "c3d-standard-30",
930 platform: Platform::Gcp,
931 family: InstanceFamily::StorageOptimized,
932 architecture: Architecture::X86_64,
933 vcpu: 30,
934 memory_bytes: 120 * GI,
935 ephemeral_storage_bytes: 1920 * GI,
936 gpu: None,
937 },
938 InstanceTypeSpec {
940 name: "n1-standard-4-t4",
941 platform: Platform::Gcp,
942 family: InstanceFamily::GpuCompute,
943 architecture: Architecture::X86_64,
944 vcpu: 4,
945 memory_bytes: 15 * GI,
946 ephemeral_storage_bytes: 100 * GI,
947 gpu: Some(CatalogGpu {
948 gpu_type: "nvidia-t4",
949 count: 1,
950 }),
951 },
952 InstanceTypeSpec {
954 name: "a2-highgpu-1g",
955 platform: Platform::Gcp,
956 family: InstanceFamily::GpuCompute,
957 architecture: Architecture::X86_64,
958 vcpu: 12,
959 memory_bytes: 85 * GI,
960 ephemeral_storage_bytes: 100 * GI,
961 gpu: Some(CatalogGpu {
962 gpu_type: "nvidia-a100",
963 count: 1,
964 }),
965 },
966 InstanceTypeSpec {
967 name: "a2-highgpu-8g",
968 platform: Platform::Gcp,
969 family: InstanceFamily::GpuCompute,
970 architecture: Architecture::X86_64,
971 vcpu: 96,
972 memory_bytes: 1360 * GI,
973 ephemeral_storage_bytes: 100 * GI,
974 gpu: Some(CatalogGpu {
975 gpu_type: "nvidia-a100",
976 count: 8,
977 }),
978 },
979 InstanceTypeSpec {
981 name: "a3-highgpu-8g",
982 platform: Platform::Gcp,
983 family: InstanceFamily::GpuCompute,
984 architecture: Architecture::X86_64,
985 vcpu: 208,
986 memory_bytes: 1872 * GI,
987 ephemeral_storage_bytes: 100 * GI,
988 gpu: Some(CatalogGpu {
989 gpu_type: "nvidia-h100",
990 count: 8,
991 }),
992 },
993 InstanceTypeSpec {
999 name: "Standard_B1s",
1000 platform: Platform::Azure,
1001 family: InstanceFamily::Burstable,
1002 architecture: Architecture::X86_64,
1003 vcpu: 1,
1004 memory_bytes: 1 * GI,
1005 ephemeral_storage_bytes: 20 * GI,
1006 gpu: None,
1007 },
1008 InstanceTypeSpec {
1009 name: "Standard_B2s",
1010 platform: Platform::Azure,
1011 family: InstanceFamily::Burstable,
1012 architecture: Architecture::X86_64,
1013 vcpu: 2,
1014 memory_bytes: 4 * GI,
1015 ephemeral_storage_bytes: 20 * GI,
1016 gpu: None,
1017 },
1018 InstanceTypeSpec {
1019 name: "Standard_B2ms",
1020 platform: Platform::Azure,
1021 family: InstanceFamily::Burstable,
1022 architecture: Architecture::X86_64,
1023 vcpu: 2,
1024 memory_bytes: 8 * GI,
1025 ephemeral_storage_bytes: 20 * GI,
1026 gpu: None,
1027 },
1028 InstanceTypeSpec {
1029 name: "Standard_B4ms",
1030 platform: Platform::Azure,
1031 family: InstanceFamily::Burstable,
1032 architecture: Architecture::X86_64,
1033 vcpu: 4,
1034 memory_bytes: 16 * GI,
1035 ephemeral_storage_bytes: 20 * GI,
1036 gpu: None,
1037 },
1038 InstanceTypeSpec {
1040 name: "Standard_D2s_v5",
1041 platform: Platform::Azure,
1042 family: InstanceFamily::GeneralPurpose,
1043 architecture: Architecture::X86_64,
1044 vcpu: 2,
1045 memory_bytes: 8 * GI,
1046 ephemeral_storage_bytes: 20 * GI,
1047 gpu: None,
1048 },
1049 InstanceTypeSpec {
1050 name: "Standard_D4s_v5",
1051 platform: Platform::Azure,
1052 family: InstanceFamily::GeneralPurpose,
1053 architecture: Architecture::X86_64,
1054 vcpu: 4,
1055 memory_bytes: 16 * GI,
1056 ephemeral_storage_bytes: 20 * GI,
1057 gpu: None,
1058 },
1059 InstanceTypeSpec {
1060 name: "Standard_D8s_v5",
1061 platform: Platform::Azure,
1062 family: InstanceFamily::GeneralPurpose,
1063 architecture: Architecture::X86_64,
1064 vcpu: 8,
1065 memory_bytes: 32 * GI,
1066 ephemeral_storage_bytes: 20 * GI,
1067 gpu: None,
1068 },
1069 InstanceTypeSpec {
1070 name: "Standard_D16s_v5",
1071 platform: Platform::Azure,
1072 family: InstanceFamily::GeneralPurpose,
1073 architecture: Architecture::X86_64,
1074 vcpu: 16,
1075 memory_bytes: 64 * GI,
1076 ephemeral_storage_bytes: 20 * GI,
1077 gpu: None,
1078 },
1079 InstanceTypeSpec {
1081 name: "Standard_F2s_v2",
1082 platform: Platform::Azure,
1083 family: InstanceFamily::ComputeOptimized,
1084 architecture: Architecture::X86_64,
1085 vcpu: 2,
1086 memory_bytes: 4 * GI,
1087 ephemeral_storage_bytes: 20 * GI,
1088 gpu: None,
1089 },
1090 InstanceTypeSpec {
1091 name: "Standard_F4s_v2",
1092 platform: Platform::Azure,
1093 family: InstanceFamily::ComputeOptimized,
1094 architecture: Architecture::X86_64,
1095 vcpu: 4,
1096 memory_bytes: 8 * GI,
1097 ephemeral_storage_bytes: 20 * GI,
1098 gpu: None,
1099 },
1100 InstanceTypeSpec {
1101 name: "Standard_F8s_v2",
1102 platform: Platform::Azure,
1103 family: InstanceFamily::ComputeOptimized,
1104 architecture: Architecture::X86_64,
1105 vcpu: 8,
1106 memory_bytes: 16 * GI,
1107 ephemeral_storage_bytes: 20 * GI,
1108 gpu: None,
1109 },
1110 InstanceTypeSpec {
1111 name: "Standard_F16s_v2",
1112 platform: Platform::Azure,
1113 family: InstanceFamily::ComputeOptimized,
1114 architecture: Architecture::X86_64,
1115 vcpu: 16,
1116 memory_bytes: 32 * GI,
1117 ephemeral_storage_bytes: 20 * GI,
1118 gpu: None,
1119 },
1120 InstanceTypeSpec {
1122 name: "Standard_E2s_v5",
1123 platform: Platform::Azure,
1124 family: InstanceFamily::MemoryOptimized,
1125 architecture: Architecture::X86_64,
1126 vcpu: 2,
1127 memory_bytes: 16 * GI,
1128 ephemeral_storage_bytes: 20 * GI,
1129 gpu: None,
1130 },
1131 InstanceTypeSpec {
1132 name: "Standard_E4s_v5",
1133 platform: Platform::Azure,
1134 family: InstanceFamily::MemoryOptimized,
1135 architecture: Architecture::X86_64,
1136 vcpu: 4,
1137 memory_bytes: 32 * GI,
1138 ephemeral_storage_bytes: 20 * GI,
1139 gpu: None,
1140 },
1141 InstanceTypeSpec {
1142 name: "Standard_E8s_v5",
1143 platform: Platform::Azure,
1144 family: InstanceFamily::MemoryOptimized,
1145 architecture: Architecture::X86_64,
1146 vcpu: 8,
1147 memory_bytes: 64 * GI,
1148 ephemeral_storage_bytes: 20 * GI,
1149 gpu: None,
1150 },
1151 InstanceTypeSpec {
1152 name: "Standard_E16s_v5",
1153 platform: Platform::Azure,
1154 family: InstanceFamily::MemoryOptimized,
1155 architecture: Architecture::X86_64,
1156 vcpu: 16,
1157 memory_bytes: 128 * GI,
1158 ephemeral_storage_bytes: 20 * GI,
1159 gpu: None,
1160 },
1161 InstanceTypeSpec {
1163 name: "Standard_L8s_v3",
1164 platform: Platform::Azure,
1165 family: InstanceFamily::StorageOptimized,
1166 architecture: Architecture::X86_64,
1167 vcpu: 8,
1168 memory_bytes: 64 * GI,
1169 ephemeral_storage_bytes: 1788 * GI,
1170 gpu: None,
1171 },
1172 InstanceTypeSpec {
1173 name: "Standard_L16s_v3",
1174 platform: Platform::Azure,
1175 family: InstanceFamily::StorageOptimized,
1176 architecture: Architecture::X86_64,
1177 vcpu: 16,
1178 memory_bytes: 128 * GI,
1179 ephemeral_storage_bytes: 3576 * GI,
1180 gpu: None,
1181 },
1182 InstanceTypeSpec {
1183 name: "Standard_L32s_v3",
1184 platform: Platform::Azure,
1185 family: InstanceFamily::StorageOptimized,
1186 architecture: Architecture::X86_64,
1187 vcpu: 32,
1188 memory_bytes: 256 * GI,
1189 ephemeral_storage_bytes: 7154 * GI,
1190 gpu: None,
1191 },
1192 InstanceTypeSpec {
1194 name: "Standard_NC4as_T4_v3",
1195 platform: Platform::Azure,
1196 family: InstanceFamily::GpuCompute,
1197 architecture: Architecture::X86_64,
1198 vcpu: 4,
1199 memory_bytes: 28 * GI,
1200 ephemeral_storage_bytes: 176 * GI,
1201 gpu: Some(CatalogGpu {
1202 gpu_type: "nvidia-t4",
1203 count: 1,
1204 }),
1205 },
1206 InstanceTypeSpec {
1208 name: "Standard_NC24ads_A100_v4",
1209 platform: Platform::Azure,
1210 family: InstanceFamily::GpuCompute,
1211 architecture: Architecture::X86_64,
1212 vcpu: 24,
1213 memory_bytes: 220 * GI,
1214 ephemeral_storage_bytes: 958 * GI,
1215 gpu: Some(CatalogGpu {
1216 gpu_type: "nvidia-a100",
1217 count: 1,
1218 }),
1219 },
1220 InstanceTypeSpec {
1221 name: "Standard_NC96ads_A100_v4",
1222 platform: Platform::Azure,
1223 family: InstanceFamily::GpuCompute,
1224 architecture: Architecture::X86_64,
1225 vcpu: 96,
1226 memory_bytes: 880 * GI,
1227 ephemeral_storage_bytes: 3916 * GI,
1228 gpu: Some(CatalogGpu {
1229 gpu_type: "nvidia-a100",
1230 count: 4,
1231 }),
1232 },
1233 InstanceTypeSpec {
1235 name: "Standard_ND96isr_H100_v5",
1236 platform: Platform::Azure,
1237 family: InstanceFamily::GpuCompute,
1238 architecture: Architecture::X86_64,
1239 vcpu: 96,
1240 memory_bytes: 1900 * GI,
1241 ephemeral_storage_bytes: 1000 * GI,
1242 gpu: Some(CatalogGpu {
1243 gpu_type: "nvidia-h100",
1244 count: 8,
1245 }),
1246 },
1247];
1248
1249pub fn catalog_for_platform(platform: Platform) -> Vec<&'static InstanceTypeSpec> {
1255 CATALOG
1256 .iter()
1257 .filter(|spec| spec.platform == platform)
1258 .collect()
1259}
1260
1261pub fn find_instance_type(platform: Platform, name: &str) -> Option<&'static InstanceTypeSpec> {
1263 CATALOG
1264 .iter()
1265 .find(|spec| spec.platform == platform && spec.name == name)
1266}
1267
1268pub fn is_same_architecture_aws_machine(old: &str, new: &str) -> bool {
1271 is_same_architecture_machine(Platform::Aws, old, new)
1272}
1273
1274pub fn is_same_architecture_machine(platform: Platform, old: &str, new: &str) -> bool {
1276 matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure)
1277 && find_instance_type(platform, old)
1278 .zip(find_instance_type(platform, new))
1279 .is_some_and(|(old, new)| old.architecture == new.architecture)
1280}
1281
1282#[derive(Debug, Clone)]
1288pub struct WorkloadRequirements {
1289 pub total_cpu_at_desired: f64,
1291 pub total_memory_bytes_at_desired: u64,
1293 pub total_cpu_at_max: f64,
1295 pub total_memory_bytes_at_max: u64,
1297 pub max_cpu_per_container: f64,
1299 pub max_memory_per_container: u64,
1301 pub max_ephemeral_storage_bytes: u64,
1303 pub gpu: Option<GpuSpec>,
1305 pub architecture: Option<Architecture>,
1307 pub nested_virt: bool,
1311}
1312
1313#[derive(Debug, Clone)]
1315pub struct InstanceSelection {
1316 pub instance_type: &'static str,
1318 pub profile: MachineProfile,
1320 pub min_machines: u32,
1322 pub max_machines: u32,
1324}
1325
1326const STORAGE_OPTIMIZED_THRESHOLD: u64 = 200 * GI;
1328
1329const MAX_MACHINES_PER_CLUSTER: u32 = 10;
1331
1332const MAX_STANDARD_VCPU: u32 = 8;
1335
1336const SYSTEM_RESERVE_CPU: f64 = 0.5;
1338
1339const WORKLOAD_HEADROOM_FACTOR: f64 = 1.15;
1341
1342pub fn select_instance_type(
1353 platform: Platform,
1354 requirements: &WorkloadRequirements,
1355) -> Result<InstanceSelection, String> {
1356 let family = select_family(requirements);
1357 let family = if requirements.nested_virt && family == InstanceFamily::Burstable {
1358 InstanceFamily::GeneralPurpose
1359 } else {
1360 family
1361 };
1362 let selection = select_instance_type_in_family(platform, requirements, family);
1363 if selection.is_err() && family == InstanceFamily::Burstable {
1364 return select_instance_type_in_family(
1365 platform,
1366 requirements,
1367 InstanceFamily::GeneralPurpose,
1368 );
1369 }
1370 selection
1371}
1372
1373fn select_instance_type_in_family(
1374 platform: Platform,
1375 requirements: &WorkloadRequirements,
1376 family: InstanceFamily,
1377) -> Result<InstanceSelection, String> {
1378 let architecture = requirements
1379 .architecture
1380 .or_else(|| default_architecture(platform))
1381 .ok_or_else(|| format!("platform {platform} has no default compute architecture"))?;
1382
1383 let candidates: Vec<&InstanceTypeSpec> = CATALOG
1384 .iter()
1385 .filter(|spec| spec.platform == platform && spec.family == family)
1386 .filter(|spec| {
1387 if requirements.nested_virt {
1388 spec.is_nested_virt_capable()
1389 } else {
1390 platform != Platform::Aws || !spec.is_nested_virt_capable()
1391 }
1392 })
1393 .filter(|spec| spec.architecture == architecture)
1394 .collect();
1395
1396 if candidates.is_empty() && family != InstanceFamily::StorageOptimized {
1401 let family_has_other_architecture = CATALOG.iter().any(|spec| {
1402 spec.platform == platform
1403 && spec.family == family
1404 && if requirements.nested_virt {
1405 spec.is_nested_virt_capable()
1406 } else {
1407 platform != Platform::Aws || !spec.is_nested_virt_capable()
1408 }
1409 });
1410 if family_has_other_architecture {
1411 return Err(format!(
1412 "architecture {architecture:?} is unavailable for this workload on platform {platform}"
1413 ));
1414 }
1415 return Err(if requirements.nested_virt {
1416 format!(
1417 "no nested-virt-capable {family:?} instance types in catalog for platform {platform}"
1418 )
1419 } else {
1420 format!("no {family:?} instance types in catalog for platform {platform}")
1421 });
1422 }
1423
1424 let candidates = if let Some(ref gpu) = requirements.gpu {
1426 let filtered: Vec<&InstanceTypeSpec> = candidates
1427 .into_iter()
1428 .filter(|spec| {
1429 spec.gpu.as_ref().map_or(false, |g| {
1430 g.gpu_type == gpu.gpu_type && g.count >= gpu.count
1431 })
1432 })
1433 .collect();
1434 if filtered.is_empty() {
1435 return Err(format!(
1436 "no instance type for GPU type '{}' x{} on platform {platform}",
1437 gpu.gpu_type, gpu.count
1438 ));
1439 }
1440 filtered
1441 } else {
1442 candidates
1443 };
1444
1445 let candidates = if family == InstanceFamily::StorageOptimized {
1447 let filtered: Vec<&InstanceTypeSpec> = candidates
1448 .into_iter()
1449 .filter(|spec| spec.ephemeral_storage_bytes >= requirements.max_ephemeral_storage_bytes)
1450 .collect();
1451 if filtered.is_empty() {
1452 CATALOG
1456 .iter()
1457 .filter(|spec| {
1458 spec.platform == platform
1459 && spec.family == InstanceFamily::GeneralPurpose
1460 && spec.has_configurable_ephemeral_storage()
1461 && max_configurable_ephemeral_storage_bytes(platform)
1462 .is_some_and(|max| requirements.max_ephemeral_storage_bytes <= max)
1463 })
1464 .filter(|spec| {
1465 if requirements.nested_virt {
1466 spec.is_nested_virt_capable()
1467 } else {
1468 platform != Platform::Aws || !spec.is_nested_virt_capable()
1469 }
1470 })
1471 .filter(|spec| spec.architecture == architecture)
1472 .collect()
1473 } else {
1474 filtered
1475 }
1476 } else {
1477 candidates
1478 };
1479
1480 if candidates.is_empty() {
1481 return Err(format!(
1482 "architecture {architecture:?} is unavailable for this workload on platform {platform}"
1483 ));
1484 }
1485
1486 let candidates: Vec<_> = candidates
1489 .into_iter()
1490 .filter(|spec| {
1491 allocatable_cpu(spec) >= requirements.max_cpu_per_container * WORKLOAD_HEADROOM_FACTOR
1492 && allocatable_memory_bytes(spec) as f64
1493 >= requirements.max_memory_per_container as f64 * WORKLOAD_HEADROOM_FACTOR
1494 })
1495 .collect();
1496 let smallest_replica_machine = candidates.iter().map(|spec| spec.vcpu).min()
1497 .ok_or_else(|| format!("no {platform} machine can fit one replica with system reserve and workload headroom"))?;
1498
1499 let effective_family = candidates[0].family;
1504 let vcpu_cap = if effective_family == InstanceFamily::GpuCompute
1505 || effective_family == InstanceFamily::StorageOptimized
1506 {
1507 u32::MAX
1508 } else {
1509 MAX_STANDARD_VCPU.max(smallest_replica_machine)
1510 };
1511
1512 let desired_target_machines = desired_target_machines(requirements);
1513 let target_cpu = requirements
1514 .max_cpu_per_container
1515 .max(requirements.total_cpu_at_desired / desired_target_machines as f64)
1516 * WORKLOAD_HEADROOM_FACTOR;
1517 let target_memory = (requirements.max_memory_per_container as f64)
1518 .max(requirements.total_memory_bytes_at_desired as f64 / desired_target_machines as f64)
1519 * WORKLOAD_HEADROOM_FACTOR;
1520
1521 let selected = candidates
1526 .iter()
1527 .filter(|spec| {
1528 spec.vcpu <= vcpu_cap
1529 && allocatable_cpu(spec) >= target_cpu
1530 && allocatable_memory_bytes(spec) as f64 >= target_memory
1531 })
1532 .min_by_key(|spec| spec.vcpu)
1533 .or_else(|| {
1534 candidates
1536 .iter()
1537 .filter(|spec| spec.vcpu <= vcpu_cap)
1538 .max_by_key(|spec| spec.vcpu)
1539 })
1540 .or_else(|| {
1541 candidates.iter().min_by_key(|spec| spec.vcpu)
1543 })
1544 .ok_or_else(|| format!("no instance types available for platform {platform}"))?;
1545
1546 let max_machines = compute_max_machines(requirements, selected);
1548 let min_machines = compute_min_machines(requirements, selected, max_machines);
1549
1550 Ok(InstanceSelection {
1551 instance_type: selected.name,
1552 profile: selected.to_machine_profile_for_storage(requirements.max_ephemeral_storage_bytes),
1553 min_machines,
1554 max_machines,
1555 })
1556}
1557
1558pub fn select_family(requirements: &WorkloadRequirements) -> InstanceFamily {
1564 if requirements.gpu.is_some() {
1566 return InstanceFamily::GpuCompute;
1567 }
1568
1569 if requirements.max_ephemeral_storage_bytes > STORAGE_OPTIMIZED_THRESHOLD {
1571 return InstanceFamily::StorageOptimized;
1572 }
1573
1574 if requirements.total_cpu_at_max < 2.0 {
1576 return InstanceFamily::Burstable;
1577 }
1578
1579 InstanceFamily::GeneralPurpose
1581}
1582
1583fn compute_max_machines(requirements: &WorkloadRequirements, instance: &InstanceTypeSpec) -> u32 {
1585 let cpu_with_headroom = requirements.total_cpu_at_max * WORKLOAD_HEADROOM_FACTOR;
1586 let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
1587
1588 let mem_with_headroom =
1589 requirements.total_memory_bytes_at_max as f64 * WORKLOAD_HEADROOM_FACTOR;
1590 let mem_machines =
1591 (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
1592
1593 cpu_machines
1595 .max(mem_machines)
1596 .max(1)
1597 .min(MAX_MACHINES_PER_CLUSTER)
1598}
1599
1600fn compute_min_machines(
1602 requirements: &WorkloadRequirements,
1603 instance: &InstanceTypeSpec,
1604 max_machines: u32,
1605) -> u32 {
1606 let cpu_with_headroom = requirements.total_cpu_at_desired * WORKLOAD_HEADROOM_FACTOR;
1607 let cpu_machines = (cpu_with_headroom / allocatable_cpu(instance)).ceil() as u32;
1608
1609 let mem_with_headroom =
1610 requirements.total_memory_bytes_at_desired as f64 * WORKLOAD_HEADROOM_FACTOR;
1611 let mem_machines =
1612 (mem_with_headroom / allocatable_memory_bytes(instance) as f64).ceil() as u32;
1613
1614 cpu_machines
1615 .max(mem_machines)
1616 .max(1)
1617 .min(2)
1618 .min(max_machines)
1619}
1620
1621fn desired_target_machines(requirements: &WorkloadRequirements) -> u32 {
1622 if requirements.total_cpu_at_desired >= 2.0
1623 || requirements.total_memory_bytes_at_desired >= 4 * GI
1624 {
1625 2
1626 } else {
1627 1
1628 }
1629}
1630
1631pub(crate) fn allocatable_cpu(instance: &InstanceTypeSpec) -> f64 {
1632 (instance.vcpu as f64 - SYSTEM_RESERVE_CPU).max(0.25)
1633}
1634
1635pub(crate) fn allocatable_memory_bytes(instance: &InstanceTypeSpec) -> u64 {
1636 instance
1637 .memory_bytes
1638 .saturating_sub(system_reserve_memory_bytes(instance.memory_bytes))
1639 .max(256 * MI)
1640}
1641
1642fn system_reserve_memory_bytes(memory_bytes: u64) -> u64 {
1643 if memory_bytes < 4 * GI {
1644 256 * MI
1645 } else if memory_bytes < 16 * GI {
1646 512 * MI
1647 } else {
1648 GI
1649 }
1650}
1651
1652#[cfg(test)]
1657mod tests {
1658 use super::*;
1659 use crate::BinaryTarget;
1660
1661 #[test]
1664 fn test_parse_cpu_plain() {
1665 assert_eq!(parse_cpu("1").unwrap(), 1.0);
1666 assert_eq!(parse_cpu("0.5").unwrap(), 0.5);
1667 assert_eq!(parse_cpu("2.0").unwrap(), 2.0);
1668 assert_eq!(parse_cpu("16").unwrap(), 16.0);
1669 }
1670
1671 #[test]
1672 fn test_parse_cpu_millicore() {
1673 assert_eq!(parse_cpu("500m").unwrap(), 0.5);
1674 assert_eq!(parse_cpu("250m").unwrap(), 0.25);
1675 assert_eq!(parse_cpu("1000m").unwrap(), 1.0);
1676 assert_eq!(parse_cpu("100m").unwrap(), 0.1);
1677 }
1678
1679 #[test]
1680 fn test_parse_cpu_invalid() {
1681 assert!(parse_cpu("").is_err());
1682 assert!(parse_cpu("abc").is_err());
1683 assert!(parse_cpu("m").is_err());
1684 }
1685
1686 #[test]
1687 fn test_parse_memory_binary_suffixes() {
1688 assert_eq!(parse_memory_bytes("1Ki").unwrap(), 1024);
1689 assert_eq!(parse_memory_bytes("1Mi").unwrap(), 1024 * 1024);
1690 assert_eq!(parse_memory_bytes("1Gi").unwrap(), 1024 * 1024 * 1024);
1691 assert_eq!(parse_memory_bytes("4Gi").unwrap(), 4 * 1024 * 1024 * 1024);
1692 assert_eq!(parse_memory_bytes("512Mi").unwrap(), 512 * 1024 * 1024);
1693 assert_eq!(
1694 parse_memory_bytes("1Ti").unwrap(),
1695 1024u64 * 1024 * 1024 * 1024
1696 );
1697 }
1698
1699 #[test]
1700 fn test_parse_memory_decimal_suffixes() {
1701 assert_eq!(parse_memory_bytes("1k").unwrap(), 1000);
1702 assert_eq!(parse_memory_bytes("1M").unwrap(), 1_000_000);
1703 assert_eq!(parse_memory_bytes("1G").unwrap(), 1_000_000_000);
1704 assert_eq!(parse_memory_bytes("1T").unwrap(), 1_000_000_000_000);
1705 }
1706
1707 #[test]
1708 fn test_parse_memory_plain_bytes() {
1709 assert_eq!(parse_memory_bytes("1024").unwrap(), 1024);
1710 assert_eq!(parse_memory_bytes("0").unwrap(), 0);
1711 }
1712
1713 #[test]
1714 fn test_parse_memory_invalid() {
1715 assert!(parse_memory_bytes("").is_err());
1716 assert!(parse_memory_bytes("abc").is_err());
1717 assert!(parse_memory_bytes("Gi").is_err());
1718 }
1719
1720 #[test]
1721 fn test_parse_memory_fractional() {
1722 assert_eq!(parse_memory_bytes("0.5Gi").unwrap(), GI / 2);
1723 assert_eq!(parse_memory_bytes("1.5Gi").unwrap(), GI + GI / 2);
1724 }
1725
1726 #[test]
1729 fn test_catalog_has_entries_for_all_cloud_platforms() {
1730 assert!(!catalog_for_platform(Platform::Aws).is_empty());
1731 assert!(!catalog_for_platform(Platform::Gcp).is_empty());
1732 assert!(!catalog_for_platform(Platform::Azure).is_empty());
1733 }
1734
1735 #[test]
1736 fn test_catalog_no_entries_for_non_cloud_platforms() {
1737 assert!(catalog_for_platform(Platform::Local).is_empty());
1738 assert!(catalog_for_platform(Platform::Kubernetes).is_empty());
1739 }
1740
1741 #[test]
1742 fn test_find_known_instance_type() {
1743 let spec =
1744 find_instance_type(Platform::Aws, "m7g.2xlarge").expect("should find m7g.2xlarge");
1745 assert_eq!(spec.vcpu, 8);
1746 assert_eq!(spec.memory_bytes, 32 * GI);
1747 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1748 }
1749
1750 #[test]
1751 fn test_find_aws_c8i_nested_virt_instance_type() {
1752 let spec = find_instance_type(Platform::Aws, "c8i.large").expect("should find c8i.large");
1753 assert_eq!(spec.vcpu, 2);
1754 assert_eq!(spec.memory_bytes, 4 * GI);
1755 assert_eq!(spec.family, InstanceFamily::ComputeOptimized);
1756 assert_eq!(spec.architecture, Architecture::X86_64);
1757 assert!(spec.is_nested_virt_capable());
1758 }
1759
1760 #[test]
1761 fn test_find_unknown_instance_type() {
1762 assert!(find_instance_type(Platform::Aws, "nonexistent.xlarge").is_none());
1763 }
1764
1765 #[test]
1766 fn test_find_wrong_platform() {
1767 assert!(find_instance_type(Platform::Gcp, "m7g.2xlarge").is_none());
1768 }
1769
1770 #[test]
1771 fn test_to_machine_profile() {
1772 let spec = find_instance_type(Platform::Aws, "m7g.2xlarge").unwrap();
1773 let profile = spec.to_machine_profile();
1774 assert_eq!(profile.cpu, "8.0");
1775 assert_eq!(profile.memory_bytes, 32 * GI);
1776 assert_eq!(profile.ephemeral_storage_bytes, 20 * GI);
1777 assert!(profile.gpu.is_none());
1778 }
1779
1780 #[test]
1781 fn test_to_machine_profile_with_gpu() {
1782 let spec = find_instance_type(Platform::Aws, "p4d.24xlarge").unwrap();
1783 let profile = spec.to_machine_profile();
1784 let gpu = profile.gpu.as_ref().expect("should have GPU");
1785 assert_eq!(gpu.gpu_type, "nvidia-a100");
1786 assert_eq!(gpu.count, 8);
1787 }
1788
1789 #[test]
1792 fn small_cpu_with_large_memory_falls_back_from_burstable() {
1793 let requirements = WorkloadRequirements {
1794 total_cpu_at_desired: 0.5,
1795 total_memory_bytes_at_desired: 8 * GI,
1796 total_cpu_at_max: 0.5,
1797 total_memory_bytes_at_max: 8 * GI,
1798 max_cpu_per_container: 0.5,
1799 max_memory_per_container: 8 * GI,
1800 max_ephemeral_storage_bytes: 0,
1801 gpu: None,
1802 architecture: Some(Architecture::X86_64),
1803 nested_virt: false,
1804 };
1805 let selected = select_instance_type(Platform::Gcp, &requirements).unwrap();
1806 let machine = find_instance_type(Platform::Gcp, selected.instance_type).unwrap();
1807 assert!(
1808 allocatable_memory_bytes(machine) as f64 >= 8.0 * GI as f64 * WORKLOAD_HEADROOM_FACTOR
1809 );
1810 }
1811
1812 #[test]
1813 fn test_select_burstable_for_small_workload() {
1814 let req = WorkloadRequirements {
1815 total_cpu_at_desired: 1.0,
1816 total_memory_bytes_at_desired: 2 * GI,
1817 total_cpu_at_max: 1.0,
1818 total_memory_bytes_at_max: 2 * GI,
1819 max_cpu_per_container: 0.5,
1820 max_memory_per_container: 1 * GI,
1821 max_ephemeral_storage_bytes: 10 * GI,
1822 gpu: None,
1823 architecture: None,
1824 nested_virt: false,
1825 };
1826 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1827 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1828 assert_eq!(spec.family, InstanceFamily::Burstable);
1829 }
1830
1831 #[test]
1832 fn test_selects_smallest_burstable_machine_with_real_headroom() {
1833 let req = WorkloadRequirements {
1834 total_cpu_at_desired: 1.0,
1835 total_memory_bytes_at_desired: 2 * GI,
1836 total_cpu_at_max: 1.0,
1837 total_memory_bytes_at_max: 2 * GI,
1838 max_cpu_per_container: 1.0,
1839 max_memory_per_container: 2 * GI,
1840 max_ephemeral_storage_bytes: 10 * GI,
1841 gpu: None,
1842 architecture: None,
1843 nested_virt: false,
1844 };
1845
1846 let selection = select_instance_type(Platform::Aws, &req).unwrap();
1847
1848 assert_eq!(selection.instance_type, "t4g.medium");
1849 assert_eq!(selection.min_machines, 1);
1850 assert_eq!(selection.max_machines, 1);
1851 }
1852
1853 #[test]
1854 fn test_select_general_purpose_for_standard_workload() {
1855 let req = WorkloadRequirements {
1857 total_cpu_at_desired: 20.0,
1858 total_memory_bytes_at_desired: 80 * GI,
1859 total_cpu_at_max: 20.0,
1860 total_memory_bytes_at_max: 80 * GI,
1861 max_cpu_per_container: 2.0,
1862 max_memory_per_container: 8 * GI,
1863 max_ephemeral_storage_bytes: 10 * GI,
1864 gpu: None,
1865 architecture: None,
1866 nested_virt: false,
1867 };
1868 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1869 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1870 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1871 }
1872
1873 #[test]
1874 fn test_select_general_purpose_even_for_cpu_heavy() {
1875 let req = WorkloadRequirements {
1877 total_cpu_at_desired: 20.0,
1878 total_memory_bytes_at_desired: 20 * GI,
1879 total_cpu_at_max: 20.0,
1880 total_memory_bytes_at_max: 20 * GI,
1881 max_cpu_per_container: 2.0,
1882 max_memory_per_container: 2 * GI,
1883 max_ephemeral_storage_bytes: 10 * GI,
1884 gpu: None,
1885 architecture: None,
1886 nested_virt: false,
1887 };
1888 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1889 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1890 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1891 }
1892
1893 #[test]
1894 fn test_select_storage_optimized_for_large_ephemeral() {
1895 let req = WorkloadRequirements {
1896 total_cpu_at_desired: 8.0,
1897 total_memory_bytes_at_desired: 32 * GI,
1898 total_cpu_at_max: 8.0,
1899 total_memory_bytes_at_max: 32 * GI,
1900 max_cpu_per_container: 2.0,
1901 max_memory_per_container: 8 * GI,
1902 max_ephemeral_storage_bytes: 500 * GI,
1903 gpu: None,
1904 architecture: Some(Architecture::X86_64),
1905 nested_virt: false,
1906 };
1907 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1908 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1909 assert_eq!(spec.family, InstanceFamily::StorageOptimized);
1910 }
1911
1912 #[test]
1913 fn test_select_configurable_storage_above_fixed_local_catalog() {
1914 let req = WorkloadRequirements {
1915 total_cpu_at_desired: 8.0,
1916 total_memory_bytes_at_desired: 32 * GI,
1917 total_cpu_at_max: 8.0,
1918 total_memory_bytes_at_max: 32 * GI,
1919 max_cpu_per_container: 2.0,
1920 max_memory_per_container: 8 * GI,
1921 max_ephemeral_storage_bytes: 8_000 * GI,
1922 gpu: None,
1923 architecture: Some(Architecture::X86_64),
1924 nested_virt: false,
1925 };
1926 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1927 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1928 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1929 assert_eq!(sel.profile.ephemeral_storage_bytes, 8_000 * GI);
1930 }
1931
1932 #[test]
1933 fn test_configurable_storage_fallback_retains_standard_vcpu_cap() {
1934 let req = WorkloadRequirements {
1935 total_cpu_at_desired: 70.0,
1936 total_memory_bytes_at_desired: 140 * GI,
1937 total_cpu_at_max: 70.0,
1938 total_memory_bytes_at_max: 140 * GI,
1939 max_cpu_per_container: 2.0,
1940 max_memory_per_container: 4 * GI,
1941 max_ephemeral_storage_bytes: 8_000 * GI,
1942 gpu: None,
1943 architecture: Some(Architecture::X86_64),
1944 nested_virt: false,
1945 };
1946
1947 let sel = select_instance_type(Platform::Aws, &req).unwrap();
1948 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
1949 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1950 assert!(spec.vcpu <= MAX_STANDARD_VCPU);
1951 assert!(sel.max_machines > 1);
1952 }
1953
1954 #[test]
1955 fn test_nested_virtualization_storage_falls_back_to_configurable_disk() {
1956 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
1957 let req = WorkloadRequirements {
1958 total_cpu_at_desired: 2.0,
1959 total_memory_bytes_at_desired: 8 * GI,
1960 total_cpu_at_max: 2.0,
1961 total_memory_bytes_at_max: 8 * GI,
1962 max_cpu_per_container: 2.0,
1963 max_memory_per_container: 8 * GI,
1964 max_ephemeral_storage_bytes: 500 * GI,
1965 gpu: None,
1966 architecture: Some(Architecture::X86_64),
1967 nested_virt: true,
1968 };
1969
1970 let sel = select_instance_type(platform, &req)
1971 .unwrap_or_else(|error| panic!("{platform} should support fallback: {error}"));
1972 let spec = find_instance_type(platform, sel.instance_type).unwrap();
1973 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
1974 assert!(spec.is_nested_virt_capable());
1975 assert_eq!(sel.profile.ephemeral_storage_bytes, 500 * GI);
1976 }
1977 }
1978
1979 #[test]
1980 fn test_configurable_disk_rejects_capacity_above_provider_limit() {
1981 for (platform, requested) in [
1982 (Platform::Aws, 100 * TI),
1983 (Platform::Gcp, 100 * TI),
1984 (Platform::Azure, 8_000 * GI),
1985 ] {
1986 let req = WorkloadRequirements {
1987 total_cpu_at_desired: 2.0,
1988 total_memory_bytes_at_desired: 8 * GI,
1989 total_cpu_at_max: 2.0,
1990 total_memory_bytes_at_max: 8 * GI,
1991 max_cpu_per_container: 2.0,
1992 max_memory_per_container: 8 * GI,
1993 max_ephemeral_storage_bytes: requested,
1994 gpu: None,
1995 architecture: Some(Architecture::X86_64),
1996 nested_virt: true,
1997 };
1998
1999 assert!(select_instance_type(platform, &req).is_err());
2000 }
2001 }
2002
2003 #[test]
2004 fn test_configurable_disk_limits_account_for_controller_rounding() {
2005 for (platform, disk_limit_gib) in [
2006 (Platform::Aws, 64 * 1024),
2007 (Platform::Gcp, 64 * 1024),
2008 (Platform::Azure, 4_095),
2009 ] {
2010 let max = max_configurable_ephemeral_storage_bytes(platform).unwrap();
2011 let materialized_disk_gib = |requested_bytes: u64| {
2012 let requested_gib = requested_bytes.div_ceil(GI);
2013 (requested_gib * 5).div_ceil(4) + 12
2014 };
2015
2016 assert!(materialized_disk_gib(max) <= disk_limit_gib);
2017 assert!(materialized_disk_gib(max + 1) > disk_limit_gib);
2018 }
2019 }
2020
2021 #[test]
2022 fn test_gpu_local_storage_is_not_treated_as_resizable() {
2023 let spec = find_instance_type(Platform::Aws, "g5.xlarge").unwrap();
2024 assert_eq!(spec.family, InstanceFamily::GpuCompute);
2025 assert!(!spec.has_configurable_ephemeral_storage());
2026 }
2027
2028 #[test]
2029 fn test_select_gpu_instance() {
2030 let req = WorkloadRequirements {
2031 total_cpu_at_desired: 8.0,
2032 total_memory_bytes_at_desired: 32 * GI,
2033 total_cpu_at_max: 8.0,
2034 total_memory_bytes_at_max: 32 * GI,
2035 max_cpu_per_container: 4.0,
2036 max_memory_per_container: 16 * GI,
2037 max_ephemeral_storage_bytes: 10 * GI,
2038 gpu: Some(GpuSpec {
2039 gpu_type: "nvidia-a100".to_string(),
2040 count: 1,
2041 }),
2042 architecture: Some(Architecture::X86_64),
2043 nested_virt: false,
2044 };
2045 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2046 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
2047 assert_eq!(spec.family, InstanceFamily::GpuCompute);
2048 assert!(spec.gpu.is_some());
2049 }
2050
2051 #[test]
2052 fn test_select_uses_each_cloud_image_target_architecture() {
2053 let req = WorkloadRequirements {
2054 total_cpu_at_desired: 4.0,
2055 total_memory_bytes_at_desired: 16 * GI,
2056 total_cpu_at_max: 4.0,
2057 total_memory_bytes_at_max: 16 * GI,
2058 max_cpu_per_container: 1.0,
2059 max_memory_per_container: 4 * GI,
2060 max_ephemeral_storage_bytes: 10 * GI,
2061 gpu: None,
2062 architecture: None,
2063 nested_virt: false,
2064 };
2065 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
2066 let sel = select_instance_type(platform, &req)
2067 .unwrap_or_else(|error| panic!("selection failed for {platform}: {error}"));
2068 let spec = find_instance_type(platform, sel.instance_type)
2069 .expect("selected machine should exist in the catalog");
2070 assert_eq!(
2071 Some(spec.architecture),
2072 default_architecture(platform),
2073 "machine architecture must match the image target for {platform}"
2074 );
2075 }
2076 }
2077
2078 #[test]
2079 fn test_machine_count_reasonable() {
2080 let req = WorkloadRequirements {
2082 total_cpu_at_desired: 20.0,
2083 total_memory_bytes_at_desired: 40 * GI,
2084 total_cpu_at_max: 20.0,
2085 total_memory_bytes_at_max: 40 * GI,
2086 max_cpu_per_container: 1.0,
2087 max_memory_per_container: 2 * GI,
2088 max_ephemeral_storage_bytes: 10 * GI,
2089 gpu: None,
2090 architecture: None,
2091 nested_virt: false,
2092 };
2093 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2094 assert!(sel.min_machines >= 1);
2095 assert!(sel.max_machines <= MAX_MACHINES_PER_CLUSTER);
2096 assert!(sel.max_machines >= sel.min_machines);
2097 }
2098
2099 #[test]
2100 fn test_instance_size_capped_at_8_vcpu() {
2101 let req = WorkloadRequirements {
2103 total_cpu_at_desired: 70.0,
2104 total_memory_bytes_at_desired: 140 * GI,
2105 total_cpu_at_max: 70.0,
2106 total_memory_bytes_at_max: 140 * GI,
2107 max_cpu_per_container: 2.0,
2108 max_memory_per_container: 4 * GI,
2109 max_ephemeral_storage_bytes: 10 * GI,
2110 gpu: None,
2111 architecture: None,
2112 nested_virt: false,
2113 };
2114 let sel = select_instance_type(Platform::Gcp, &req).unwrap();
2115 let spec = find_instance_type(Platform::Gcp, sel.instance_type).unwrap();
2116 assert!(
2117 spec.vcpu <= MAX_STANDARD_VCPU,
2118 "selected {} with {} vCPUs, expected <= {}",
2119 spec.name,
2120 spec.vcpu,
2121 MAX_STANDARD_VCPU
2122 );
2123 assert_eq!(spec.family, InstanceFamily::GeneralPurpose);
2124 assert!(sel.max_machines > 1);
2126 }
2127
2128 #[test]
2129 fn test_larger_autoscaled_workload_gets_reasonable_instance() {
2130 let req = WorkloadRequirements {
2133 total_cpu_at_desired: 70.0,
2134 total_memory_bytes_at_desired: 140 * GI,
2135 total_cpu_at_max: 70.0, total_memory_bytes_at_max: 140 * GI, max_cpu_per_container: 2.0,
2138 max_memory_per_container: 4 * GI,
2139 max_ephemeral_storage_bytes: 20 * GI,
2140 gpu: None,
2141 architecture: None,
2142 nested_virt: false,
2143 };
2144 let sel = select_instance_type(Platform::Gcp, &req).unwrap();
2145 assert_eq!(sel.instance_type, "n2-standard-8");
2147 assert!(sel.max_machines >= 2);
2148 }
2149
2150 #[test]
2157 fn test_select_aws_picks_m8i_when_nested_virt_required() {
2158 let req = WorkloadRequirements {
2159 total_cpu_at_desired: 4.0,
2160 total_memory_bytes_at_desired: 8 * GI,
2161 total_cpu_at_max: 4.0,
2162 total_memory_bytes_at_max: 8 * GI,
2163 max_cpu_per_container: 4.0,
2164 max_memory_per_container: 8 * GI,
2165 max_ephemeral_storage_bytes: 10 * GI,
2166 gpu: None,
2167 architecture: Some(Architecture::X86_64),
2168 nested_virt: true,
2169 };
2170 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2171 assert!(
2172 sel.instance_type.starts_with("m8i.")
2173 || sel.instance_type.starts_with("c8i.")
2174 || sel.instance_type.starts_with("r8i."),
2175 "expected an m8i/c8i/r8i instance, got {}",
2176 sel.instance_type
2177 );
2178 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
2179 assert!(spec.is_nested_virt_capable());
2180 }
2181
2182 #[test]
2183 fn test_select_gcp_picks_n2_when_nested_virt_required() {
2184 let req = WorkloadRequirements {
2185 total_cpu_at_desired: 4.0,
2186 total_memory_bytes_at_desired: 8 * GI,
2187 total_cpu_at_max: 4.0,
2188 total_memory_bytes_at_max: 8 * GI,
2189 max_cpu_per_container: 4.0,
2190 max_memory_per_container: 8 * GI,
2191 max_ephemeral_storage_bytes: 0,
2192 architecture: Some(Architecture::X86_64),
2193 gpu: None,
2194 nested_virt: true,
2195 };
2196
2197 let selection = select_instance_type(Platform::Gcp, &req).unwrap();
2198 assert_eq!(selection.instance_type, "n2-standard-8");
2199 assert!(find_instance_type(Platform::Gcp, selection.instance_type)
2200 .unwrap()
2201 .is_nested_virt_capable());
2202 }
2203
2204 #[test]
2205 fn test_select_azure_picks_dsv5_when_nested_virt_required() {
2206 let req = WorkloadRequirements {
2207 total_cpu_at_desired: 4.0,
2208 total_memory_bytes_at_desired: 8 * GI,
2209 total_cpu_at_max: 4.0,
2210 total_memory_bytes_at_max: 8 * GI,
2211 max_cpu_per_container: 4.0,
2212 max_memory_per_container: 8 * GI,
2213 max_ephemeral_storage_bytes: 0,
2214 architecture: Some(Architecture::X86_64),
2215 gpu: None,
2216 nested_virt: true,
2217 };
2218
2219 let selection = select_instance_type(Platform::Azure, &req).unwrap();
2220 assert_eq!(selection.instance_type, "Standard_D8s_v5");
2221 assert!(find_instance_type(Platform::Azure, selection.instance_type)
2222 .unwrap()
2223 .is_nested_virt_capable());
2224 }
2225
2226 #[test]
2227 fn test_select_aws_defaults_to_image_target_architecture() {
2228 let req = WorkloadRequirements {
2229 total_cpu_at_desired: 4.0,
2230 total_memory_bytes_at_desired: 8 * GI,
2231 total_cpu_at_max: 4.0,
2232 total_memory_bytes_at_max: 8 * GI,
2233 max_cpu_per_container: 4.0,
2234 max_memory_per_container: 8 * GI,
2235 max_ephemeral_storage_bytes: 10 * GI,
2236 gpu: None,
2237 architecture: None,
2238 nested_virt: false,
2239 };
2240 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2241 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
2242 assert_eq!(spec.architecture, Architecture::Arm64);
2243 }
2244
2245 #[test]
2246 fn test_cloud_defaults_match_image_target_architectures() {
2247 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
2248 let target = BinaryTarget::defaults_for_platform(platform)
2249 .into_iter()
2250 .next()
2251 .expect("managed cloud should have a default image target");
2252 let image_architecture = match target.oci_arch() {
2253 "arm64" => Architecture::Arm64,
2254 "amd64" => Architecture::X86_64,
2255 architecture => {
2256 panic!("unsupported managed-cloud image architecture {architecture}")
2257 }
2258 };
2259
2260 assert_eq!(default_architecture(platform), Some(image_architecture));
2261 }
2262 }
2263
2264 #[test]
2266 fn test_select_aws_uses_graviton_for_explicit_arm64() {
2267 let req = WorkloadRequirements {
2268 total_cpu_at_desired: 4.0,
2269 total_memory_bytes_at_desired: 8 * GI,
2270 total_cpu_at_max: 4.0,
2271 total_memory_bytes_at_max: 8 * GI,
2272 max_cpu_per_container: 4.0,
2273 max_memory_per_container: 8 * GI,
2274 max_ephemeral_storage_bytes: 10 * GI,
2275 gpu: None,
2276 architecture: Some(Architecture::Arm64),
2277 nested_virt: false,
2278 };
2279 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2280 let spec = find_instance_type(Platform::Aws, sel.instance_type).unwrap();
2281 assert_eq!(spec.architecture, Architecture::Arm64);
2282 }
2283
2284 #[test]
2285 fn test_select_rejects_explicit_architecture_missing_from_cloud_catalog() {
2286 let req = WorkloadRequirements {
2287 total_cpu_at_desired: 1.0,
2288 total_memory_bytes_at_desired: 2 * GI,
2289 total_cpu_at_max: 1.0,
2290 total_memory_bytes_at_max: 2 * GI,
2291 max_cpu_per_container: 1.0,
2292 max_memory_per_container: 2 * GI,
2293 max_ephemeral_storage_bytes: 10 * GI,
2294 gpu: None,
2295 architecture: Some(Architecture::Arm64),
2296 nested_virt: false,
2297 };
2298
2299 let error = select_instance_type(Platform::Gcp, &req)
2300 .expect_err("GCP catalog has no ARM64 machine");
2301
2302 assert!(error.contains("architecture Arm64 is unavailable"));
2303 }
2304
2305 #[test]
2306 fn test_profile_has_required_fields() {
2307 let req = WorkloadRequirements {
2308 total_cpu_at_desired: 4.0,
2309 total_memory_bytes_at_desired: 16 * GI,
2310 total_cpu_at_max: 4.0,
2311 total_memory_bytes_at_max: 16 * GI,
2312 max_cpu_per_container: 1.0,
2313 max_memory_per_container: 4 * GI,
2314 max_ephemeral_storage_bytes: 10 * GI,
2315 gpu: None,
2316 architecture: None,
2317 nested_virt: false,
2318 };
2319 let sel = select_instance_type(Platform::Aws, &req).unwrap();
2320 assert!(!sel.profile.cpu.is_empty());
2321 assert!(sel.profile.memory_bytes > 0);
2322 assert!(sel.profile.ephemeral_storage_bytes > 0);
2323 }
2324
2325 #[test]
2326 fn test_error_for_unsupported_gpu_type() {
2327 let req = WorkloadRequirements {
2328 total_cpu_at_desired: 8.0,
2329 total_memory_bytes_at_desired: 32 * GI,
2330 total_cpu_at_max: 8.0,
2331 total_memory_bytes_at_max: 32 * GI,
2332 max_cpu_per_container: 4.0,
2333 max_memory_per_container: 16 * GI,
2334 max_ephemeral_storage_bytes: 10 * GI,
2335 gpu: Some(GpuSpec {
2336 gpu_type: "amd-mi300".to_string(),
2337 count: 1,
2338 }),
2339 architecture: None,
2340 nested_virt: false,
2341 };
2342 let result = select_instance_type(Platform::Aws, &req);
2343 assert!(result.is_err());
2344 }
2345
2346 #[test]
2347 fn test_catalog_instance_types_sorted_by_vcpu_within_family() {
2348 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
2351 let entries = catalog_for_platform(platform);
2352 let mut by_family: std::collections::HashMap<_, Vec<_>> =
2353 std::collections::HashMap::new();
2354 for entry in entries {
2355 by_family
2356 .entry(format!("{:?}", entry.family))
2357 .or_default()
2358 .push(entry);
2359 }
2360 for (family, instances) in &by_family {
2361 for window in instances.windows(2) {
2362 assert!(
2363 window[0].vcpu <= window[1].vcpu,
2364 "catalog not sorted by vcpu for {platform}/{family}: {} ({}) > {} ({})",
2365 window[0].name,
2366 window[0].vcpu,
2367 window[1].name,
2368 window[1].vcpu
2369 );
2370 }
2371 }
2372 }
2373 }
2374}