1use crate::{
8 instance_catalog::{self, Architecture, WorkloadRequirements},
9 CapacityGroup, CapacityGroupScalePolicy, ComputeChoiceRange, ComputePoolSelection, Container,
10 Daemon, ErrorData, FailureDomainSelection, GpuSpec, MachineProfile, Platform, ResourceSpec,
11 Stack,
12};
13use alien_error::{AlienError, Result};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16
17#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
19#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
20#[serde(rename_all = "camelCase")]
21pub struct ComputePlan {
22 pub pools: Vec<ComputePoolPlan>,
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
29#[serde(rename_all = "camelCase")]
30pub struct ComputePoolPlan {
31 pub pool_id: String,
33 pub workloads: Vec<String>,
35 pub requirements: MachineProfile,
37 pub scale: CapacityGroupScalePolicy,
39 pub selected: ComputePoolSelection,
41 pub recommended: ComputePoolSelection,
43 pub machines: Vec<ComputeMachineOption>,
45 #[serde(default, skip_serializing_if = "Vec::is_empty")]
47 pub errors: Vec<String>,
48}
49
50#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
53#[serde(rename_all = "camelCase")]
54pub struct ComputeMachineOption {
55 pub machine: String,
57 pub profile: MachineProfile,
59 pub recommended: bool,
61}
62
63pub fn plan_compute(
65 stack: &Stack,
66 platform: Platform,
67 selected_settings: Option<&crate::ComputeSettings>,
68) -> Result<ComputePlan, ErrorData> {
69 let mut groups = collect_workload_groups(stack)?;
70 merge_explicit_compute_groups(stack, &mut groups)?;
71
72 let mut pool_ids: Vec<String> = groups.keys().cloned().collect();
73 pool_ids.sort();
74
75 let mut pools = Vec::new();
76 for pool_id in pool_ids {
77 let group = groups.remove(&pool_id).expect("pool id came from map keys");
78 let requirements = group.requirements;
79 let selected = selected_settings.and_then(|settings| settings.pools.get(&pool_id));
80 let recommended = recommended_selection(
81 platform,
82 &requirements,
83 &group.scale,
84 group.requires_failure_domain,
85 )?;
86 let mut selected_choice = selected.cloned().unwrap_or_else(|| recommended.clone());
87 if selected_choice.failure_domains().is_none() {
88 if let Some(default_failure_domains) = recommended.failure_domains().cloned() {
89 match &mut selected_choice {
90 ComputePoolSelection::Fixed {
91 failure_domains, ..
92 }
93 | ComputePoolSelection::Autoscale {
94 failure_domains, ..
95 } => *failure_domains = Some(default_failure_domains),
96 }
97 }
98 }
99 let errors = validate_compute_pool_selection(
100 platform,
101 &pool_id,
102 &selected_choice,
103 &requirements,
104 &group.scale,
105 );
106 let machines = machine_options(platform, &requirements, selected_choice.machine())?;
107
108 pools.push(ComputePoolPlan {
109 pool_id,
110 workloads: group.workloads,
111 requirements: requirements_to_profile(&requirements),
112 scale: group.scale,
113 selected: selected_choice,
114 recommended,
115 machines,
116 errors,
117 });
118 }
119
120 Ok(ComputePlan { pools })
121}
122
123#[derive(Debug, Clone)]
124struct PlannedGroup {
125 workloads: Vec<String>,
126 requirements: WorkloadRequirements,
127 scale: CapacityGroupScalePolicy,
128 requires_failure_domain: bool,
129}
130
131fn collect_workload_groups(stack: &Stack) -> Result<HashMap<String, PlannedGroup>, ErrorData> {
132 let mut groups: HashMap<String, Vec<Workload>> = HashMap::new();
133
134 for entry in stack.resources.values() {
135 if let Some(container) = entry.config.downcast_ref::<Container>() {
136 groups
137 .entry(
138 container
139 .pool
140 .clone()
141 .unwrap_or_else(|| needed_container_pool(container).to_string()),
142 )
143 .or_default()
144 .push(Workload::from_container(container)?);
145 }
146 if let Some(daemon) = entry.config.downcast_ref::<Daemon>() {
147 if daemon.cluster.is_some() {
148 groups
149 .entry(daemon.pool.clone().unwrap_or_else(|| "general".to_string()))
150 .or_default()
151 .push(Workload::from_daemon(daemon)?);
152 }
153 }
154 }
155
156 let mut planned = HashMap::new();
157 for (pool_id, workloads) in groups {
158 let requirements = aggregate_workloads(&workloads);
159 let requires_failure_domain = workloads
160 .iter()
161 .any(|workload| workload.requires_failure_domain);
162 let min_size = default_min_machines(&requirements);
163 let max_size = default_max_machines(&requirements);
164 planned.insert(
165 pool_id,
166 PlannedGroup {
167 workloads: workloads.into_iter().map(|w| w.id).collect(),
168 scale: CapacityGroupScalePolicy::from_selected_bounds(min_size, max_size),
169 requirements,
170 requires_failure_domain,
171 },
172 );
173 }
174 if planned.is_empty() {
175 let requirements = default_requirements();
176 planned.insert(
177 "general".to_string(),
178 PlannedGroup {
179 workloads: Vec::new(),
180 scale: CapacityGroupScalePolicy::from_selected_bounds(1, 1),
181 requirements,
182 requires_failure_domain: false,
183 },
184 );
185 }
186 Ok(planned)
187}
188
189fn merge_explicit_compute_groups(
190 stack: &Stack,
191 groups: &mut HashMap<String, PlannedGroup>,
192) -> Result<(), ErrorData> {
193 for entry in stack.resources.values() {
194 let Some(cluster) = entry.config.downcast_ref::<crate::ComputeCluster>() else {
195 continue;
196 };
197 for group in &cluster.capacity_groups {
198 let explicit_requirements = profile_to_requirements(
199 group.profile.as_ref(),
200 group.nested_virtualization.unwrap_or(false),
201 );
202 let scale = group.scale_policy.clone().unwrap_or_else(|| {
203 CapacityGroupScalePolicy::from_selected_bounds(group.min_size, group.max_size)
204 });
205 groups
206 .entry(group.group_id.clone())
207 .and_modify(|planned| {
208 if group.instance_type.is_none() {
209 merge_requirements(&mut planned.requirements, &explicit_requirements);
210 } else {
211 planned.requirements.architecture = planned
218 .requirements
219 .architecture
220 .or(explicit_requirements.architecture);
221 planned.requirements.nested_virt |= explicit_requirements.nested_virt;
222 }
223 planned.scale = merge_scale_policy(&planned.scale, &scale);
224 })
225 .or_insert_with(|| PlannedGroup {
226 workloads: Vec::new(),
227 scale,
228 requirements: explicit_requirements,
229 requires_failure_domain: false,
230 });
231 }
232 }
233 Ok(())
234}
235
236fn recommended_selection(
237 platform: Platform,
238 requirements: &WorkloadRequirements,
239 scale: &CapacityGroupScalePolicy,
240 requires_failure_domain: bool,
241) -> Result<ComputePoolSelection, ErrorData> {
242 let machine = match platform {
243 Platform::Aws | Platform::Gcp | Platform::Azure => Some(
244 instance_catalog::select_instance_type(platform, requirements)
245 .map_err(|message| {
246 AlienError::new(ErrorData::GenericError {
247 message: format!("Failed to select {platform} machine: {message}"),
248 })
249 })?
250 .instance_type
251 .to_string(),
252 ),
253 Platform::Local | Platform::Kubernetes | Platform::Machines | Platform::Test => None,
254 };
255
256 let failure_domains = (requires_failure_domain
257 && matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure))
258 .then_some(FailureDomainSelection {
259 spread: 1,
260 selected_failure_domains: Vec::new(),
261 });
262
263 match scale {
264 CapacityGroupScalePolicy::Fixed { machines } => Ok(ComputePoolSelection::Fixed {
265 machines: machines.default.max(1),
266 machine,
267 failure_domains,
268 }),
269 CapacityGroupScalePolicy::Autoscale { min, max } => Ok(ComputePoolSelection::Autoscale {
270 min: min.default,
271 max: max.default.max(min.default),
272 machine,
273 failure_domains,
274 }),
275 }
276}
277
278pub fn validate_compute_pool_selection(
281 platform: Platform,
282 pool_id: &str,
283 selection: &ComputePoolSelection,
284 requirements: &WorkloadRequirements,
285 scale: &CapacityGroupScalePolicy,
286) -> Vec<String> {
287 let mut errors = Vec::new();
288 if let Err(message) = selection.validate() {
289 errors.push(message);
290 }
291 if let Err(message) = validate_selection_against_scale(selection, scale) {
292 errors.push(format!("Pool '{pool_id}' {message}"));
293 }
294 if matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
295 match selection.machine() {
296 Some(machine) => match instance_catalog::find_instance_type(platform, machine) {
297 Some(spec) => {
298 let architecture = requirements.architecture.unwrap_or(spec.architecture);
299 if !instance_satisfies(spec, requirements, architecture) {
300 errors.push(format!(
301 "{} machine '{}' does not satisfy pool '{}' requirements",
302 platform, machine, pool_id
303 ));
304 }
305 }
306 None => errors.push(format!(
307 "Unknown {} machine '{}' for pool '{}'",
308 platform, machine, pool_id
309 )),
310 },
311 None => errors.push(format!(
312 "Pool '{}' requires a provider machine on {}",
313 pool_id, platform
314 )),
315 }
316 }
317 errors
318}
319
320pub fn capacity_group_requirements(group: &CapacityGroup) -> WorkloadRequirements {
322 profile_to_requirements(
323 group.profile.as_ref(),
324 group.nested_virtualization.unwrap_or(false),
325 )
326}
327
328fn machine_options(
329 platform: Platform,
330 requirements: &WorkloadRequirements,
331 selected_machine: Option<&str>,
332) -> Result<Vec<ComputeMachineOption>, ErrorData> {
333 if !matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
334 return Ok(Vec::new());
335 }
336 let recommended =
337 instance_catalog::select_instance_type(platform, requirements).map_err(|message| {
338 AlienError::new(ErrorData::GenericError {
339 message: format!("Failed to select {platform} machine: {message}"),
340 })
341 })?;
342 let resolved_architecture = requirements
343 .architecture
344 .or_else(|| {
345 selected_machine.and_then(|machine| {
346 instance_catalog::find_instance_type(platform, machine)
347 .map(|spec| spec.architecture)
348 })
349 })
350 .or(recommended.profile.architecture)
351 .ok_or_else(|| {
352 AlienError::new(ErrorData::GenericError {
353 message: format!("Selected {platform} machine has no CPU architecture"),
354 })
355 })?;
356 let recommended = recommended.instance_type.to_string();
357
358 let mut options: Vec<ComputeMachineOption> = instance_catalog::catalog_for_platform(platform)
359 .into_iter()
360 .filter(|spec| instance_satisfies(spec, requirements, resolved_architecture))
361 .map(|spec| ComputeMachineOption {
362 machine: spec.name.to_string(),
363 profile: spec.to_machine_profile(),
364 recommended: spec.name == recommended || Some(spec.name) == selected_machine,
365 })
366 .collect();
367 options.sort_by(|a, b| a.machine.cmp(&b.machine));
368 Ok(options)
369}
370
371fn instance_satisfies(
372 spec: &instance_catalog::InstanceTypeSpec,
373 requirements: &WorkloadRequirements,
374 resolved_architecture: Architecture,
375) -> bool {
376 if spec.architecture != resolved_architecture {
377 return false;
378 }
379 if requirements.nested_virt && !spec.is_nested_virt_capable() {
380 return false;
381 }
382 if spec.vcpu < requirements.max_cpu_per_container.ceil() as u32 {
383 return false;
384 }
385 if spec.memory_bytes < requirements.max_memory_per_container {
386 return false;
387 }
388 if spec.ephemeral_storage_bytes < requirements.max_ephemeral_storage_bytes {
389 return false;
390 }
391 match (&requirements.gpu, spec.gpu) {
392 (Some(required), Some(actual)) => {
393 (required.gpu_type == "any" || required.gpu_type == actual.gpu_type)
394 && actual.count >= required.count
395 }
396 (Some(_), None) => false,
397 (None, _) => true,
398 }
399}
400
401#[derive(Debug, Clone)]
402struct Workload {
403 id: String,
404 cpu: f64,
405 memory_bytes: u64,
406 desired_replicas: f64,
407 max_replicas: f64,
408 ephemeral_storage_bytes: u64,
409 gpu: Option<GpuSpec>,
410 requires_failure_domain: bool,
411}
412
413impl Workload {
414 fn from_container(container: &Container) -> Result<Self, ErrorData> {
415 let cpu = parse_cpu(&container.id, &container.cpu)?;
416 let memory_bytes = parse_memory(&container.id, &container.memory)?;
417 let desired_replicas = container
418 .autoscaling
419 .as_ref()
420 .map(|a| a.desired)
421 .or(container.replicas)
422 .unwrap_or(1) as f64;
423 let max_replicas = container
424 .autoscaling
425 .as_ref()
426 .map(|a| a.max)
427 .or(container.replicas)
428 .unwrap_or(1) as f64;
429 let ephemeral_storage_bytes = container
430 .ephemeral_storage
431 .as_deref()
432 .map(instance_catalog::parse_memory_bytes)
433 .transpose()
434 .map_err(|message| {
435 AlienError::new(ErrorData::GenericError {
436 message: format!(
437 "Failed to parse ephemeral storage for '{}': {message}",
438 container.id
439 ),
440 })
441 })?
442 .unwrap_or(0);
443
444 Ok(Self {
445 id: container.id.clone(),
446 cpu,
447 memory_bytes,
448 desired_replicas,
449 max_replicas,
450 ephemeral_storage_bytes,
451 gpu: container.gpu.as_ref().map(|gpu| GpuSpec {
452 gpu_type: gpu.gpu_type.clone(),
453 count: gpu.count,
454 }),
455 requires_failure_domain: container.stateful && container.persistent_storage.is_some(),
456 })
457 }
458
459 fn from_daemon(daemon: &Daemon) -> Result<Self, ErrorData> {
460 Ok(Self {
461 id: daemon.id.clone(),
462 cpu: parse_cpu(&daemon.id, &daemon.cpu)?,
463 memory_bytes: parse_memory(&daemon.id, &daemon.memory)?,
464 desired_replicas: 1.0,
465 max_replicas: 1.0,
466 ephemeral_storage_bytes: 0,
467 gpu: None,
468 requires_failure_domain: false,
469 })
470 }
471}
472
473fn parse_cpu(resource_id: &str, spec: &ResourceSpec) -> Result<f64, ErrorData> {
474 instance_catalog::parse_cpu(&spec.desired).map_err(|message| {
475 AlienError::new(ErrorData::GenericError {
476 message: format!(
477 "Failed to parse CPU requirement '{}' for '{}': {message}",
478 spec.desired, resource_id
479 ),
480 })
481 })
482}
483
484fn parse_memory(resource_id: &str, spec: &ResourceSpec) -> Result<u64, ErrorData> {
485 instance_catalog::parse_memory_bytes(&spec.desired).map_err(|message| {
486 AlienError::new(ErrorData::GenericError {
487 message: format!(
488 "Failed to parse memory requirement '{}' for '{}': {message}",
489 spec.desired, resource_id
490 ),
491 })
492 })
493}
494
495fn aggregate_workloads(workloads: &[Workload]) -> WorkloadRequirements {
496 let mut requirements = default_requirements();
497 requirements.total_cpu_at_desired = 0.0;
498 requirements.total_memory_bytes_at_desired = 0;
499 requirements.total_cpu_at_max = 0.0;
500 requirements.total_memory_bytes_at_max = 0;
501 requirements.max_cpu_per_container = 0.0;
502 requirements.max_memory_per_container = 0;
503 requirements.max_ephemeral_storage_bytes = 0;
504 requirements.gpu = None;
505
506 for workload in workloads {
507 requirements.total_cpu_at_desired += workload.cpu * workload.desired_replicas;
508 requirements.total_cpu_at_max += workload.cpu * workload.max_replicas;
509 requirements.total_memory_bytes_at_desired +=
510 (workload.memory_bytes as f64 * workload.desired_replicas) as u64;
511 requirements.total_memory_bytes_at_max +=
512 (workload.memory_bytes as f64 * workload.max_replicas) as u64;
513 requirements.max_cpu_per_container = requirements.max_cpu_per_container.max(workload.cpu);
514 requirements.max_memory_per_container = requirements
515 .max_memory_per_container
516 .max(workload.memory_bytes);
517 requirements.max_ephemeral_storage_bytes = requirements
518 .max_ephemeral_storage_bytes
519 .max(workload.ephemeral_storage_bytes);
520 if requirements.gpu.is_none() {
521 requirements.gpu = workload.gpu.clone();
522 }
523 }
524 requirements
525}
526
527fn default_requirements() -> WorkloadRequirements {
528 WorkloadRequirements {
529 total_cpu_at_desired: 1.0,
530 total_memory_bytes_at_desired: 2 * 1024 * 1024 * 1024,
531 total_cpu_at_max: 1.0,
532 total_memory_bytes_at_max: 2 * 1024 * 1024 * 1024,
533 max_cpu_per_container: 1.0,
534 max_memory_per_container: 2 * 1024 * 1024 * 1024,
535 max_ephemeral_storage_bytes: 0,
536 gpu: None,
537 architecture: None,
538 nested_virt: false,
539 }
540}
541
542fn profile_to_requirements(
543 profile: Option<&MachineProfile>,
544 nested_virt: bool,
545) -> WorkloadRequirements {
546 let Some(profile) = profile else {
547 return WorkloadRequirements {
548 nested_virt,
549 ..default_requirements()
550 };
551 };
552 let cpu = instance_catalog::parse_cpu(&profile.cpu).unwrap_or(1.0);
553 WorkloadRequirements {
554 total_cpu_at_desired: cpu,
555 total_memory_bytes_at_desired: profile.memory_bytes,
556 total_cpu_at_max: cpu,
557 total_memory_bytes_at_max: profile.memory_bytes,
558 max_cpu_per_container: cpu,
559 max_memory_per_container: profile.memory_bytes,
560 max_ephemeral_storage_bytes: profile.ephemeral_storage_bytes,
561 gpu: profile.gpu.clone(),
562 architecture: profile.architecture,
563 nested_virt,
564 }
565}
566
567fn merge_requirements(existing: &mut WorkloadRequirements, declared: &WorkloadRequirements) {
568 existing.total_cpu_at_desired = existing
569 .total_cpu_at_desired
570 .max(declared.total_cpu_at_desired);
571 existing.total_memory_bytes_at_desired = existing
572 .total_memory_bytes_at_desired
573 .max(declared.total_memory_bytes_at_desired);
574 existing.total_cpu_at_max = existing.total_cpu_at_max.max(declared.total_cpu_at_max);
575 existing.total_memory_bytes_at_max = existing
576 .total_memory_bytes_at_max
577 .max(declared.total_memory_bytes_at_max);
578 existing.max_cpu_per_container = existing
579 .max_cpu_per_container
580 .max(declared.max_cpu_per_container);
581 existing.max_memory_per_container = existing
582 .max_memory_per_container
583 .max(declared.max_memory_per_container);
584 existing.max_ephemeral_storage_bytes = existing
585 .max_ephemeral_storage_bytes
586 .max(declared.max_ephemeral_storage_bytes);
587 if existing.gpu.is_none() {
588 existing.gpu = declared.gpu.clone();
589 }
590 existing.nested_virt |= declared.nested_virt;
591 if existing.architecture.is_none() {
592 existing.architecture = declared.architecture;
593 }
594}
595
596fn requirements_to_profile(requirements: &WorkloadRequirements) -> MachineProfile {
597 MachineProfile {
598 cpu: requirements.max_cpu_per_container.to_string(),
599 memory_bytes: requirements.max_memory_per_container,
600 ephemeral_storage_bytes: requirements.max_ephemeral_storage_bytes,
601 architecture: requirements.architecture,
602 gpu: requirements.gpu.clone(),
603 }
604}
605
606fn merge_scale_policy(
607 existing: &CapacityGroupScalePolicy,
608 declared: &CapacityGroupScalePolicy,
609) -> CapacityGroupScalePolicy {
610 match (existing, declared) {
611 (
612 CapacityGroupScalePolicy::Fixed {
613 machines: existing_machines,
614 },
615 CapacityGroupScalePolicy::Fixed {
616 machines: declared_machines,
617 },
618 ) => CapacityGroupScalePolicy::Fixed {
619 machines: merge_choice_range(existing_machines, declared_machines),
620 },
621 (_, declared) => declared.clone(),
622 }
623}
624
625fn merge_choice_range(
626 existing: &ComputeChoiceRange,
627 declared: &ComputeChoiceRange,
628) -> ComputeChoiceRange {
629 ComputeChoiceRange {
630 min: existing.min.max(declared.min),
631 max: existing.max.max(declared.max),
632 default: declared.default,
633 }
634}
635
636fn validate_selection_against_scale(
637 selection: &ComputePoolSelection,
638 scale: &CapacityGroupScalePolicy,
639) -> std::result::Result<(), String> {
640 match (selection, scale) {
641 (
642 ComputePoolSelection::Fixed { machines, .. },
643 CapacityGroupScalePolicy::Fixed { machines: allowed },
644 ) => {
645 if allowed.contains(*machines) {
646 Ok(())
647 } else {
648 Err(format!(
649 "fixed machine count {machines} is outside the allowed range {}-{}",
650 allowed.min, allowed.max
651 ))
652 }
653 }
654 (
655 ComputePoolSelection::Autoscale { min, max, .. },
656 CapacityGroupScalePolicy::Autoscale {
657 min: allowed_min,
658 max: allowed_max,
659 },
660 ) => {
661 if !allowed_min.contains(*min) {
662 return Err(format!(
663 "autoscale minimum {min} is outside the allowed range {}-{}",
664 allowed_min.min, allowed_min.max
665 ));
666 }
667 if !allowed_max.contains(*max) {
668 return Err(format!(
669 "autoscale maximum {max} is outside the allowed range {}-{}",
670 allowed_max.min, allowed_max.max
671 ));
672 }
673 Ok(())
674 }
675 (ComputePoolSelection::Fixed { .. }, CapacityGroupScalePolicy::Autoscale { .. }) => {
676 Err("must use autoscale mode".to_string())
677 }
678 (ComputePoolSelection::Autoscale { .. }, CapacityGroupScalePolicy::Fixed { .. }) => {
679 Err("must use fixed mode".to_string())
680 }
681 }
682}
683
684fn needed_container_pool(container: &Container) -> &'static str {
685 if container.stateful && container.persistent_storage.is_some() {
686 return "stateful";
687 }
688 if container.gpu.is_some() {
689 return "gpu";
690 }
691 if let Some(storage) = &container.ephemeral_storage {
692 if instance_catalog::parse_memory_bytes(storage).unwrap_or(0) > 200 * 1024 * 1024 * 1024 {
693 return "storage";
694 }
695 }
696 "general"
697}
698
699fn default_min_machines(requirements: &WorkloadRequirements) -> u32 {
700 if requirements.total_cpu_at_desired > 0.0 || requirements.total_memory_bytes_at_desired > 0 {
701 1
702 } else {
703 0
704 }
705}
706
707fn default_max_machines(requirements: &WorkloadRequirements) -> u32 {
708 let min = default_min_machines(requirements);
709 let by_cpu =
710 (requirements.total_cpu_at_max / requirements.max_cpu_per_container.max(1.0)).ceil() as u32;
711 let by_mem = requirements
712 .total_memory_bytes_at_max
713 .div_ceil(requirements.max_memory_per_container.max(1)) as u32;
714 min.max(by_cpu).max(by_mem).max(1)
715}
716
717#[cfg(test)]
718mod tests {
719 use super::*;
720 use crate::{
721 instance_catalog::Architecture, CapacityGroup, CapacityGroupScalePolicy,
722 ComputeChoiceRange, ComputeCluster, ComputeSettings, ContainerCode, DaemonCode, Resource,
723 ResourceEntry, ResourceLifecycle, Stack,
724 };
725
726 fn stack_with_container() -> Stack {
727 let container = Container::new("api".to_string())
728 .code(ContainerCode::Image {
729 image: "api:latest".to_string(),
730 })
731 .cpu(ResourceSpec {
732 min: "1".to_string(),
733 desired: "2".to_string(),
734 })
735 .memory(ResourceSpec {
736 min: "2Gi".to_string(),
737 desired: "4Gi".to_string(),
738 })
739 .permissions("api".to_string())
740 .build();
741 Stack {
742 id: "test".to_string(),
743 resources: [(
744 "api".to_string(),
745 ResourceEntry {
746 config: Resource::new(container),
747 lifecycle: ResourceLifecycle::Live,
748 dependencies: Vec::new(),
749 remote_access: false,
750 },
751 )]
752 .into_iter()
753 .collect(),
754 permissions: crate::permissions::PermissionsConfig::default(),
755 supported_platforms: None,
756 inputs: vec![],
757 }
758 }
759
760 #[test]
761 fn cloud_plan_recommends_provider_machine_without_mutating_selection() {
762 let stack = stack_with_container();
763
764 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
765
766 let pool = plan.pools.first().expect("general pool should exist");
767 assert_eq!(pool.pool_id, "general");
768 assert_eq!(pool.workloads, vec!["api"]);
769 assert!(pool.selected.machine().is_some());
770 assert!(pool.machines.iter().any(|machine| machine.recommended));
771 }
772
773 #[test]
774 fn cloud_plan_only_offers_machines_matching_the_default_image_target() {
775 let stack = stack_with_container();
776
777 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
778 let plan = plan_compute(&stack, platform, None).expect("plan should build");
779 let pool = plan.pools.first().expect("general pool should exist");
780 let expected = instance_catalog::default_architecture(platform)
781 .expect("managed cloud should have a default architecture");
782
783 assert!(pool.machines.iter().all(|machine| {
784 instance_catalog::find_instance_type(platform, &machine.machine)
785 .is_some_and(|spec| spec.architecture == expected)
786 }));
787 }
788 }
789
790 #[test]
791 fn selected_machine_is_preserved_as_static_deployment_choice() {
792 let stack = stack_with_container();
793 let settings = ComputeSettings {
794 pools: [(
795 "general".to_string(),
796 ComputePoolSelection::Fixed {
797 machines: 1,
798 machine: Some("m7g.xlarge".to_string()),
799 failure_domains: None,
800 },
801 )]
802 .into_iter()
803 .collect(),
804 };
805
806 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
807
808 let pool = plan.pools.first().expect("general pool should exist");
809 assert_eq!(pool.selected.machine(), Some("m7g.xlarge"));
810 assert!(pool.errors.is_empty());
811 }
812
813 #[test]
814 fn selected_machine_defines_architecture_when_workloads_do_not() {
815 let stack = stack_with_container();
816 let settings = ComputeSettings {
817 pools: [(
818 "general".to_string(),
819 ComputePoolSelection::Fixed {
820 machines: 1,
821 machine: Some("m7i.xlarge".to_string()),
822 failure_domains: None,
823 },
824 )]
825 .into_iter()
826 .collect(),
827 };
828
829 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
830
831 let pool = plan.pools.first().expect("general pool should exist");
832 assert!(pool.errors.is_empty());
833 assert!(pool.machines.iter().all(|machine| {
834 instance_catalog::find_instance_type(Platform::Aws, &machine.machine)
835 .is_some_and(|spec| spec.architecture == Architecture::X86_64)
836 }));
837 }
838
839 #[test]
840 fn explicit_capacity_group_requirements_are_merged_with_workloads() {
841 let mut stack = stack_with_container();
842 let cluster = ComputeCluster::new("compute".to_string())
843 .capacity_group(CapacityGroup {
844 group_id: "general".to_string(),
845 instance_type: None,
846 profile: Some(MachineProfile {
847 cpu: "4".to_string(),
848 memory_bytes: 16 * 1024 * 1024 * 1024,
849 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
850 architecture: Some(Architecture::X86_64),
851 gpu: None,
852 }),
853 min_size: 2,
854 max_size: 5,
855 scale_policy: None,
856 nested_virtualization: Some(true),
857 })
858 .build();
859 stack.resources.insert(
860 "compute".to_string(),
861 ResourceEntry {
862 config: Resource::new(cluster),
863 lifecycle: ResourceLifecycle::Frozen,
864 dependencies: Vec::new(),
865 remote_access: false,
866 },
867 );
868
869 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
870
871 let pool = plan.pools.first().expect("general pool should exist");
872 let machine = pool
873 .selected
874 .machine()
875 .expect("AWS selection should include a machine");
876 let spec = instance_catalog::find_instance_type(Platform::Aws, machine)
877 .expect("selected machine should exist in the catalog");
878 assert!(spec.is_nested_virt_capable());
879 assert_eq!(pool.selected.min_size(), 2);
880 assert_eq!(pool.selected.max_size(), 5);
881 assert!(pool.errors.is_empty());
882 }
883
884 #[test]
885 fn materialized_machine_profile_is_not_counted_as_workload_demand() {
886 let mut stack = stack_with_container();
887 let selected = instance_catalog::find_instance_type(Platform::Aws, "m7g.xlarge")
888 .expect("machine should exist in catalog");
889 let cluster = ComputeCluster::new("compute".to_string())
890 .capacity_group(CapacityGroup {
891 group_id: "general".to_string(),
892 instance_type: Some(selected.name.to_string()),
893 profile: Some(selected.to_machine_profile()),
894 min_size: 1,
895 max_size: 1,
896 scale_policy: None,
897 nested_virtualization: None,
898 })
899 .build();
900 stack.resources.insert(
901 "compute".to_string(),
902 ResourceEntry {
903 config: Resource::new(cluster),
904 lifecycle: ResourceLifecycle::Frozen,
905 dependencies: Vec::new(),
906 remote_access: false,
907 },
908 );
909
910 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
911 let pool = plan.pools.first().expect("general pool should exist");
912
913 assert_eq!(pool.requirements.cpu, "2");
914 assert_eq!(pool.requirements.memory_bytes, 4 * 1024 * 1024 * 1024);
915 assert_eq!(pool.recommended.machine(), Some("m7g.xlarge"));
916 assert!(pool
917 .machines
918 .iter()
919 .any(|option| option.machine == "m7g.xlarge"));
920 }
921
922 #[test]
923 fn nested_x86_fixed_range_pool_preserves_bounds_and_rejects_graviton() {
924 let daemon = Daemon::new("vm-runtime-loader".to_string())
925 .code(DaemonCode::Image {
926 image: "example.com/vm-runtime:latest".to_string(),
927 })
928 .cluster("vm-runtime".to_string())
929 .cpu(ResourceSpec {
930 min: "2".to_string(),
931 desired: "2".to_string(),
932 })
933 .memory(ResourceSpec {
934 min: "4Gi".to_string(),
935 desired: "4Gi".to_string(),
936 })
937 .permissions("loader".to_string())
938 .build();
939 let cluster = ComputeCluster::new("vm-runtime".to_string())
940 .capacity_group(CapacityGroup {
941 group_id: "general".to_string(),
942 instance_type: None,
943 profile: Some(MachineProfile {
944 cpu: "4".to_string(),
945 memory_bytes: 16 * 1024 * 1024 * 1024,
946 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
947 architecture: Some(Architecture::X86_64),
948 gpu: None,
949 }),
950 min_size: 2,
951 max_size: 2,
952 scale_policy: Some(CapacityGroupScalePolicy::Fixed {
953 machines: ComputeChoiceRange {
954 min: 1,
955 max: 5,
956 default: 2,
957 },
958 }),
959 nested_virtualization: Some(true),
960 })
961 .build();
962 let stack = Stack {
963 id: "vm-runtime".to_string(),
964 resources: [
965 (
966 "vm-runtime-loader".to_string(),
967 ResourceEntry {
968 config: Resource::new(daemon),
969 lifecycle: ResourceLifecycle::Live,
970 dependencies: Vec::new(),
971 remote_access: false,
972 },
973 ),
974 (
975 "vm-runtime".to_string(),
976 ResourceEntry {
977 config: Resource::new(cluster),
978 lifecycle: ResourceLifecycle::Frozen,
979 dependencies: Vec::new(),
980 remote_access: false,
981 },
982 ),
983 ]
984 .into_iter()
985 .collect(),
986 permissions: crate::permissions::PermissionsConfig::default(),
987 supported_platforms: None,
988 inputs: vec![],
989 };
990
991 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
992 let pool = plan.pools.first().expect("general pool should exist");
993 assert_eq!(pool.recommended.machine(), Some("m8i.2xlarge"));
994 assert_eq!(pool.recommended.min_size(), 2);
995 assert_eq!(pool.recommended.max_size(), 2);
996 assert_eq!(
997 pool.scale,
998 CapacityGroupScalePolicy::Fixed {
999 machines: ComputeChoiceRange {
1000 min: 1,
1001 max: 5,
1002 default: 2,
1003 },
1004 }
1005 );
1006 assert!(!pool
1007 .machines
1008 .iter()
1009 .any(|option| option.machine == "m7g.2xlarge"));
1010
1011 let invalid_settings = ComputeSettings {
1012 pools: [(
1013 "general".to_string(),
1014 ComputePoolSelection::Fixed {
1015 machines: 2,
1016 machine: Some("m7g.2xlarge".to_string()),
1017 failure_domains: None,
1018 },
1019 )]
1020 .into_iter()
1021 .collect(),
1022 };
1023 let invalid_plan = plan_compute(&stack, Platform::Aws, Some(&invalid_settings))
1024 .expect("plan should build");
1025 assert!(!invalid_plan.pools[0].errors.is_empty());
1026 }
1027
1028 #[test]
1029 fn local_plan_has_no_provider_machine_choices() {
1030 let stack = stack_with_container();
1031
1032 let plan = plan_compute(&stack, Platform::Local, None).expect("plan should build");
1033
1034 let pool = plan.pools.first().expect("general pool should exist");
1035 assert_eq!(pool.selected.machine(), None);
1036 assert!(pool.machines.is_empty());
1037 assert!(pool.errors.is_empty());
1038 }
1039}