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 Ok(planned)
175}
176
177fn merge_explicit_compute_groups(
178 stack: &Stack,
179 groups: &mut HashMap<String, PlannedGroup>,
180) -> Result<(), ErrorData> {
181 for entry in stack.resources.values() {
182 let Some(cluster) = entry.config.downcast_ref::<crate::ComputeCluster>() else {
183 continue;
184 };
185 for group in &cluster.capacity_groups {
186 let explicit_requirements = profile_to_requirements(
187 group.profile.as_ref(),
188 group.nested_virtualization.unwrap_or(false),
189 );
190 let scale = group.scale_policy.clone().unwrap_or_else(|| {
191 CapacityGroupScalePolicy::from_selected_bounds(group.min_size, group.max_size)
192 });
193 groups
194 .entry(group.group_id.clone())
195 .and_modify(|planned| {
196 if group.instance_type.is_none() {
197 merge_requirements(&mut planned.requirements, &explicit_requirements);
198 } else {
199 planned.requirements.architecture = planned
206 .requirements
207 .architecture
208 .or(explicit_requirements.architecture);
209 planned.requirements.nested_virt |= explicit_requirements.nested_virt;
210 }
211 planned.scale = merge_scale_policy(&planned.scale, &scale);
212 })
213 .or_insert_with(|| PlannedGroup {
214 workloads: Vec::new(),
215 scale,
216 requirements: explicit_requirements,
217 requires_failure_domain: false,
218 });
219 }
220 }
221 Ok(())
222}
223
224fn recommended_selection(
225 platform: Platform,
226 requirements: &WorkloadRequirements,
227 scale: &CapacityGroupScalePolicy,
228 requires_failure_domain: bool,
229) -> Result<ComputePoolSelection, ErrorData> {
230 let machine = match platform {
231 Platform::Aws | Platform::Gcp | Platform::Azure => Some(
232 instance_catalog::select_instance_type(platform, requirements)
233 .map_err(|message| {
234 AlienError::new(ErrorData::GenericError {
235 message: format!("Failed to select {platform} machine: {message}"),
236 })
237 })?
238 .instance_type
239 .to_string(),
240 ),
241 Platform::Local | Platform::Kubernetes | Platform::Machines | Platform::Test => None,
242 };
243
244 let failure_domains = (requires_failure_domain
245 && matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure))
246 .then_some(FailureDomainSelection {
247 spread: 1,
248 selected_failure_domains: Vec::new(),
249 });
250
251 match scale {
252 CapacityGroupScalePolicy::Fixed { machines } => Ok(ComputePoolSelection::Fixed {
253 machines: machines.default.max(1),
254 machine,
255 failure_domains,
256 }),
257 CapacityGroupScalePolicy::Autoscale { min, max } => Ok(ComputePoolSelection::Autoscale {
258 min: min.default,
259 max: max.default.max(min.default),
260 machine,
261 failure_domains,
262 }),
263 }
264}
265
266pub fn validate_compute_pool_selection(
269 platform: Platform,
270 pool_id: &str,
271 selection: &ComputePoolSelection,
272 requirements: &WorkloadRequirements,
273 scale: &CapacityGroupScalePolicy,
274) -> Vec<String> {
275 let mut errors = Vec::new();
276 if let Err(message) = selection.validate() {
277 errors.push(message);
278 }
279 if let Err(message) = validate_selection_against_scale(selection, scale) {
280 errors.push(format!("Pool '{pool_id}' {message}"));
281 }
282 if matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
283 match selection.machine() {
284 Some(machine) => match instance_catalog::find_instance_type(platform, machine) {
285 Some(spec) => {
286 let architecture = requirements.architecture.unwrap_or(spec.architecture);
287 if !instance_satisfies(spec, requirements, architecture) {
288 errors.push(format!(
289 "{} machine '{}' does not satisfy pool '{}' requirements",
290 platform, machine, pool_id
291 ));
292 }
293 }
294 None => errors.push(format!(
295 "Unknown {} machine '{}' for pool '{}'",
296 platform, machine, pool_id
297 )),
298 },
299 None => errors.push(format!(
300 "Pool '{}' requires a provider machine on {}",
301 pool_id, platform
302 )),
303 }
304 }
305 errors
306}
307
308pub fn capacity_group_requirements(group: &CapacityGroup) -> WorkloadRequirements {
310 profile_to_requirements(
311 group.profile.as_ref(),
312 group.nested_virtualization.unwrap_or(false),
313 )
314}
315
316fn machine_options(
317 platform: Platform,
318 requirements: &WorkloadRequirements,
319 selected_machine: Option<&str>,
320) -> Result<Vec<ComputeMachineOption>, ErrorData> {
321 if !matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
322 return Ok(Vec::new());
323 }
324 let recommended =
325 instance_catalog::select_instance_type(platform, requirements).map_err(|message| {
326 AlienError::new(ErrorData::GenericError {
327 message: format!("Failed to select {platform} machine: {message}"),
328 })
329 })?;
330 let resolved_architecture = requirements
331 .architecture
332 .or_else(|| {
333 selected_machine.and_then(|machine| {
334 instance_catalog::find_instance_type(platform, machine)
335 .map(|spec| spec.architecture)
336 })
337 })
338 .or(recommended.profile.architecture)
339 .ok_or_else(|| {
340 AlienError::new(ErrorData::GenericError {
341 message: format!("Selected {platform} machine has no CPU architecture"),
342 })
343 })?;
344 let recommended = recommended.instance_type.to_string();
345
346 let mut options: Vec<ComputeMachineOption> = instance_catalog::catalog_for_platform(platform)
347 .into_iter()
348 .filter(|spec| instance_satisfies(spec, requirements, resolved_architecture))
349 .map(|spec| ComputeMachineOption {
350 machine: spec.name.to_string(),
351 profile: spec.to_machine_profile_for_storage(requirements.max_ephemeral_storage_bytes),
352 recommended: spec.name == recommended || Some(spec.name) == selected_machine,
353 })
354 .collect();
355 options.sort_by(|a, b| a.machine.cmp(&b.machine));
356 Ok(options)
357}
358
359fn instance_satisfies(
360 spec: &instance_catalog::InstanceTypeSpec,
361 requirements: &WorkloadRequirements,
362 resolved_architecture: Architecture,
363) -> bool {
364 if spec.architecture != resolved_architecture {
365 return false;
366 }
367 if requirements.nested_virt && !spec.is_nested_virt_capable() {
368 return false;
369 }
370 if spec.vcpu < requirements.max_cpu_per_container.ceil() as u32 {
371 return false;
372 }
373 if spec.memory_bytes < requirements.max_memory_per_container {
374 return false;
375 }
376 if spec.has_configurable_ephemeral_storage() {
377 if crate::instance_catalog::max_configurable_ephemeral_storage_bytes(spec.platform)
378 .is_none_or(|max| requirements.max_ephemeral_storage_bytes > max)
379 {
380 return false;
381 }
382 } else if spec.ephemeral_storage_bytes < requirements.max_ephemeral_storage_bytes {
383 return false;
384 }
385 match (&requirements.gpu, spec.gpu) {
386 (Some(required), Some(actual)) => {
387 (required.gpu_type == "any" || required.gpu_type == actual.gpu_type)
388 && actual.count >= required.count
389 }
390 (Some(_), None) => false,
391 (None, _) => true,
392 }
393}
394
395#[derive(Debug, Clone)]
396struct Workload {
397 id: String,
398 cpu: f64,
399 memory_bytes: u64,
400 desired_replicas: f64,
401 max_replicas: f64,
402 ephemeral_storage_bytes: u64,
403 gpu: Option<GpuSpec>,
404 requires_failure_domain: bool,
405}
406
407impl Workload {
408 fn from_container(container: &Container) -> Result<Self, ErrorData> {
409 let cpu = parse_cpu(&container.id, &container.cpu)?;
410 let memory_bytes = parse_memory(&container.id, &container.memory)?;
411 let desired_replicas = container
412 .autoscaling
413 .as_ref()
414 .map(|a| a.desired)
415 .or(container.replicas)
416 .unwrap_or(1) as f64;
417 let max_replicas = container
418 .autoscaling
419 .as_ref()
420 .map(|a| a.max)
421 .or(container.replicas)
422 .unwrap_or(1) as f64;
423 let ephemeral_storage_bytes = container
424 .ephemeral_storage
425 .as_deref()
426 .map(instance_catalog::parse_memory_bytes)
427 .transpose()
428 .map_err(|message| {
429 AlienError::new(ErrorData::GenericError {
430 message: format!(
431 "Failed to parse ephemeral storage for '{}': {message}",
432 container.id
433 ),
434 })
435 })?
436 .unwrap_or(0);
437
438 Ok(Self {
439 id: container.id.clone(),
440 cpu,
441 memory_bytes,
442 desired_replicas,
443 max_replicas,
444 ephemeral_storage_bytes,
445 gpu: container.gpu.as_ref().map(|gpu| GpuSpec {
446 gpu_type: gpu.gpu_type.clone(),
447 count: gpu.count,
448 }),
449 requires_failure_domain: container.stateful && container.persistent_storage.is_some(),
450 })
451 }
452
453 fn from_daemon(daemon: &Daemon) -> Result<Self, ErrorData> {
454 Ok(Self {
455 id: daemon.id.clone(),
456 cpu: parse_cpu(&daemon.id, &daemon.cpu)?,
457 memory_bytes: parse_memory(&daemon.id, &daemon.memory)?,
458 desired_replicas: 1.0,
459 max_replicas: 1.0,
460 ephemeral_storage_bytes: 0,
461 gpu: None,
462 requires_failure_domain: false,
463 })
464 }
465}
466
467fn parse_cpu(resource_id: &str, spec: &ResourceSpec) -> Result<f64, ErrorData> {
468 instance_catalog::parse_cpu(&spec.desired).map_err(|message| {
469 AlienError::new(ErrorData::GenericError {
470 message: format!(
471 "Failed to parse CPU requirement '{}' for '{}': {message}",
472 spec.desired, resource_id
473 ),
474 })
475 })
476}
477
478fn parse_memory(resource_id: &str, spec: &ResourceSpec) -> Result<u64, ErrorData> {
479 instance_catalog::parse_memory_bytes(&spec.desired).map_err(|message| {
480 AlienError::new(ErrorData::GenericError {
481 message: format!(
482 "Failed to parse memory requirement '{}' for '{}': {message}",
483 spec.desired, resource_id
484 ),
485 })
486 })
487}
488
489fn aggregate_workloads(workloads: &[Workload]) -> WorkloadRequirements {
490 let mut requirements = default_requirements();
491 requirements.total_cpu_at_desired = 0.0;
492 requirements.total_memory_bytes_at_desired = 0;
493 requirements.total_cpu_at_max = 0.0;
494 requirements.total_memory_bytes_at_max = 0;
495 requirements.max_cpu_per_container = 0.0;
496 requirements.max_memory_per_container = 0;
497 requirements.max_ephemeral_storage_bytes = 0;
498 requirements.gpu = None;
499
500 for workload in workloads {
501 requirements.total_cpu_at_desired += workload.cpu * workload.desired_replicas;
502 requirements.total_cpu_at_max += workload.cpu * workload.max_replicas;
503 requirements.total_memory_bytes_at_desired +=
504 (workload.memory_bytes as f64 * workload.desired_replicas) as u64;
505 requirements.total_memory_bytes_at_max +=
506 (workload.memory_bytes as f64 * workload.max_replicas) as u64;
507 requirements.max_cpu_per_container = requirements.max_cpu_per_container.max(workload.cpu);
508 requirements.max_memory_per_container = requirements
509 .max_memory_per_container
510 .max(workload.memory_bytes);
511 requirements.max_ephemeral_storage_bytes = requirements
512 .max_ephemeral_storage_bytes
513 .max(workload.ephemeral_storage_bytes);
514 if requirements.gpu.is_none() {
515 requirements.gpu = workload.gpu.clone();
516 }
517 }
518 requirements
519}
520
521fn default_requirements() -> WorkloadRequirements {
522 WorkloadRequirements {
523 total_cpu_at_desired: 1.0,
524 total_memory_bytes_at_desired: 2 * 1024 * 1024 * 1024,
525 total_cpu_at_max: 1.0,
526 total_memory_bytes_at_max: 2 * 1024 * 1024 * 1024,
527 max_cpu_per_container: 1.0,
528 max_memory_per_container: 2 * 1024 * 1024 * 1024,
529 max_ephemeral_storage_bytes: 0,
530 gpu: None,
531 architecture: None,
532 nested_virt: false,
533 }
534}
535
536fn profile_to_requirements(
537 profile: Option<&MachineProfile>,
538 nested_virt: bool,
539) -> WorkloadRequirements {
540 let Some(profile) = profile else {
541 return WorkloadRequirements {
542 nested_virt,
543 ..default_requirements()
544 };
545 };
546 let cpu = instance_catalog::parse_cpu(&profile.cpu).unwrap_or(1.0);
547 WorkloadRequirements {
548 total_cpu_at_desired: cpu,
549 total_memory_bytes_at_desired: profile.memory_bytes,
550 total_cpu_at_max: cpu,
551 total_memory_bytes_at_max: profile.memory_bytes,
552 max_cpu_per_container: cpu,
553 max_memory_per_container: profile.memory_bytes,
554 max_ephemeral_storage_bytes: profile.ephemeral_storage_bytes,
555 gpu: profile.gpu.clone(),
556 architecture: profile.architecture,
557 nested_virt,
558 }
559}
560
561fn merge_requirements(existing: &mut WorkloadRequirements, declared: &WorkloadRequirements) {
562 existing.total_cpu_at_desired = existing
563 .total_cpu_at_desired
564 .max(declared.total_cpu_at_desired);
565 existing.total_memory_bytes_at_desired = existing
566 .total_memory_bytes_at_desired
567 .max(declared.total_memory_bytes_at_desired);
568 existing.total_cpu_at_max = existing.total_cpu_at_max.max(declared.total_cpu_at_max);
569 existing.total_memory_bytes_at_max = existing
570 .total_memory_bytes_at_max
571 .max(declared.total_memory_bytes_at_max);
572 existing.max_cpu_per_container = existing
573 .max_cpu_per_container
574 .max(declared.max_cpu_per_container);
575 existing.max_memory_per_container = existing
576 .max_memory_per_container
577 .max(declared.max_memory_per_container);
578 existing.max_ephemeral_storage_bytes = existing
579 .max_ephemeral_storage_bytes
580 .max(declared.max_ephemeral_storage_bytes);
581 if existing.gpu.is_none() {
582 existing.gpu = declared.gpu.clone();
583 }
584 existing.nested_virt |= declared.nested_virt;
585 if existing.architecture.is_none() {
586 existing.architecture = declared.architecture;
587 }
588}
589
590fn requirements_to_profile(requirements: &WorkloadRequirements) -> MachineProfile {
591 MachineProfile {
592 cpu: requirements.max_cpu_per_container.to_string(),
593 memory_bytes: requirements.max_memory_per_container,
594 ephemeral_storage_bytes: requirements.max_ephemeral_storage_bytes,
595 architecture: requirements.architecture,
596 gpu: requirements.gpu.clone(),
597 }
598}
599
600fn merge_scale_policy(
601 existing: &CapacityGroupScalePolicy,
602 declared: &CapacityGroupScalePolicy,
603) -> CapacityGroupScalePolicy {
604 match (existing, declared) {
605 (
606 CapacityGroupScalePolicy::Fixed {
607 machines: existing_machines,
608 },
609 CapacityGroupScalePolicy::Fixed {
610 machines: declared_machines,
611 },
612 ) => CapacityGroupScalePolicy::Fixed {
613 machines: merge_choice_range(existing_machines, declared_machines),
614 },
615 (_, declared) => declared.clone(),
616 }
617}
618
619fn merge_choice_range(
620 existing: &ComputeChoiceRange,
621 declared: &ComputeChoiceRange,
622) -> ComputeChoiceRange {
623 ComputeChoiceRange {
624 min: existing.min.max(declared.min),
625 max: existing.max.max(declared.max),
626 default: declared.default,
627 }
628}
629
630fn validate_selection_against_scale(
631 selection: &ComputePoolSelection,
632 scale: &CapacityGroupScalePolicy,
633) -> std::result::Result<(), String> {
634 match (selection, scale) {
635 (
636 ComputePoolSelection::Fixed { machines, .. },
637 CapacityGroupScalePolicy::Fixed { machines: allowed },
638 ) => {
639 if allowed.contains(*machines) {
640 Ok(())
641 } else {
642 Err(format!(
643 "fixed machine count {machines} is outside the allowed range {}-{}",
644 allowed.min, allowed.max
645 ))
646 }
647 }
648 (
649 ComputePoolSelection::Autoscale { min, max, .. },
650 CapacityGroupScalePolicy::Autoscale {
651 min: allowed_min,
652 max: allowed_max,
653 },
654 ) => {
655 if !allowed_min.contains(*min) {
656 return Err(format!(
657 "autoscale minimum {min} is outside the allowed range {}-{}",
658 allowed_min.min, allowed_min.max
659 ));
660 }
661 if !allowed_max.contains(*max) {
662 return Err(format!(
663 "autoscale maximum {max} is outside the allowed range {}-{}",
664 allowed_max.min, allowed_max.max
665 ));
666 }
667 Ok(())
668 }
669 (ComputePoolSelection::Fixed { .. }, CapacityGroupScalePolicy::Autoscale { .. }) => {
670 Err("must use autoscale mode".to_string())
671 }
672 (ComputePoolSelection::Autoscale { .. }, CapacityGroupScalePolicy::Fixed { .. }) => {
673 Err("must use fixed mode".to_string())
674 }
675 }
676}
677
678fn needed_container_pool(container: &Container) -> &'static str {
679 if container.stateful && container.persistent_storage.is_some() {
680 return "stateful";
681 }
682 if container.gpu.is_some() {
683 return "gpu";
684 }
685 if let Some(storage) = &container.ephemeral_storage {
686 if instance_catalog::parse_memory_bytes(storage).unwrap_or(0) > 200 * 1024 * 1024 * 1024 {
687 return "storage";
688 }
689 }
690 "general"
691}
692
693fn default_min_machines(requirements: &WorkloadRequirements) -> u32 {
694 if requirements.total_cpu_at_desired > 0.0 || requirements.total_memory_bytes_at_desired > 0 {
695 1
696 } else {
697 0
698 }
699}
700
701fn default_max_machines(requirements: &WorkloadRequirements) -> u32 {
702 let min = default_min_machines(requirements);
703 let by_cpu =
704 (requirements.total_cpu_at_max / requirements.max_cpu_per_container.max(1.0)).ceil() as u32;
705 let by_mem = requirements
706 .total_memory_bytes_at_max
707 .div_ceil(requirements.max_memory_per_container.max(1)) as u32;
708 min.max(by_cpu).max(by_mem).max(1)
709}
710
711#[cfg(test)]
712mod tests {
713 use super::*;
714 use crate::{
715 instance_catalog::Architecture, CapacityGroup, CapacityGroupScalePolicy,
716 ComputeChoiceRange, ComputeCluster, ComputeSettings, ContainerCode, DaemonCode, Resource,
717 ResourceEntry, ResourceLifecycle, Stack, Storage, Worker, WorkerCode,
718 };
719
720 fn stack_with_container() -> Stack {
721 let container = Container::new("api".to_string())
722 .code(ContainerCode::Image {
723 image: "api:latest".to_string(),
724 })
725 .cpu(ResourceSpec {
726 min: "1".to_string(),
727 desired: "2".to_string(),
728 })
729 .memory(ResourceSpec {
730 min: "2Gi".to_string(),
731 desired: "4Gi".to_string(),
732 })
733 .permissions("api".to_string())
734 .build();
735 Stack {
736 id: "test".to_string(),
737 resources: [(
738 "api".to_string(),
739 ResourceEntry {
740 config: Resource::new(container),
741 lifecycle: ResourceLifecycle::Live,
742 dependencies: Vec::new(),
743 remote_access: false,
744 enabled_when: None,
745 },
746 )]
747 .into_iter()
748 .collect(),
749 permissions: crate::permissions::PermissionsConfig::default(),
750 supported_platforms: None,
751 inputs: vec![],
752 dynamic_container_repositories: Vec::new(),
753 dynamic_container_image_resources: Vec::new(),
754 }
755 }
756
757 #[test]
758 fn serverless_workers_and_storage_do_not_allocate_machine_pools() {
759 let mut stack = stack_with_container();
760 stack.resources.clear();
761 for (id, resource) in [
762 (
763 "api",
764 Resource::new(
765 Worker::new("api".to_string())
766 .code(WorkerCode::Image {
767 image: "example/worker:v1".to_string(),
768 })
769 .permissions("execution".to_string())
770 .build(),
771 ),
772 ),
773 (
774 "data",
775 Resource::new(Storage::new("data".to_string()).build()),
776 ),
777 ] {
778 stack.resources.insert(
779 id.to_string(),
780 ResourceEntry {
781 config: resource,
782 lifecycle: ResourceLifecycle::Frozen,
783 dependencies: Vec::new(),
784 remote_access: false,
785 enabled_when: None,
786 },
787 );
788 }
789
790 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
791 let plan = plan_compute(&stack, platform, None).expect("plan should build");
792 assert!(plan.pools.is_empty(), "serverless plan allocated {plan:?}");
793 }
794 }
795
796 #[test]
797 fn empty_stack_does_not_allocate_machine_pools() {
798 let mut stack = stack_with_container();
799 stack.resources.clear();
800 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
801 assert!(plan.pools.is_empty());
802 }
803
804 #[test]
805 fn explicit_cluster_without_workloads_only_plans_declared_pools() {
806 let mut stack = stack_with_container();
807 stack.resources.clear();
808 let cluster = ComputeCluster::new("compute".to_string())
809 .capacity_group(CapacityGroup {
810 group_id: "batch".to_string(),
811 instance_type: Some("t4g.medium".to_string()),
812 profile: None,
813 min_size: 1,
814 max_size: 1,
815 scale_policy: None,
816 nested_virtualization: None,
817 })
818 .build();
819 stack.resources.insert(
820 "compute".to_string(),
821 ResourceEntry {
822 config: Resource::new(cluster),
823 lifecycle: ResourceLifecycle::Frozen,
824 dependencies: Vec::new(),
825 remote_access: false,
826 enabled_when: None,
827 },
828 );
829
830 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
831 assert_eq!(plan.pools.len(), 1);
832 assert_eq!(plan.pools[0].pool_id, "batch");
833 assert!(plan.pools[0].errors.is_empty());
834 }
835
836 #[test]
837 fn cloud_plan_recommends_provider_machine_without_mutating_selection() {
838 let stack = stack_with_container();
839
840 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
841
842 let pool = plan.pools.first().expect("general pool should exist");
843 assert_eq!(pool.pool_id, "general");
844 assert_eq!(pool.workloads, vec!["api"]);
845 assert!(pool.selected.machine().is_some());
846 assert!(pool.machines.iter().any(|machine| machine.recommended));
847 }
848
849 #[test]
850 fn cloud_plan_only_offers_machines_matching_the_default_image_target() {
851 let stack = stack_with_container();
852
853 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
854 let plan = plan_compute(&stack, platform, None).expect("plan should build");
855 let pool = plan.pools.first().expect("general pool should exist");
856 let expected = instance_catalog::default_architecture(platform)
857 .expect("managed cloud should have a default architecture");
858
859 assert!(pool.machines.iter().all(|machine| {
860 instance_catalog::find_instance_type(platform, &machine.machine)
861 .is_some_and(|spec| spec.architecture == expected)
862 }));
863 }
864 }
865
866 #[test]
867 fn selected_machine_is_preserved_as_static_deployment_choice() {
868 let stack = stack_with_container();
869 let settings = ComputeSettings {
870 pools: [(
871 "general".to_string(),
872 ComputePoolSelection::Fixed {
873 machines: 1,
874 machine: Some("m7g.xlarge".to_string()),
875 failure_domains: None,
876 },
877 )]
878 .into_iter()
879 .collect(),
880 };
881
882 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
883
884 let pool = plan.pools.first().expect("general pool should exist");
885 assert_eq!(pool.selected.machine(), Some("m7g.xlarge"));
886 assert!(pool.errors.is_empty());
887 }
888
889 #[test]
890 fn configurable_cloud_disk_satisfies_storage_above_fixed_local_catalog() {
891 let mut stack = stack_with_container();
892 let container = stack
893 .resources
894 .get_mut("api")
895 .and_then(|entry| entry.config.downcast_mut::<Container>())
896 .expect("test stack should contain a container");
897 container.ephemeral_storage = Some("8000Gi".to_string());
898 let settings = ComputeSettings {
899 pools: [(
900 "storage".to_string(),
901 ComputePoolSelection::Fixed {
902 machines: 1,
903 machine: Some("m8i.2xlarge".to_string()),
904 failure_domains: None,
905 },
906 )]
907 .into_iter()
908 .collect(),
909 };
910
911 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
912 let pool = plan.pools.first().expect("general pool should exist");
913 assert!(pool.errors.is_empty());
914 let selected = pool
915 .machines
916 .iter()
917 .find(|option| option.machine == "m8i.2xlarge")
918 .expect("selected machine should remain an option");
919 assert_eq!(
920 selected.profile.ephemeral_storage_bytes,
921 8_000 * 1024 * 1024 * 1024
922 );
923 }
924
925 #[test]
926 fn selected_machine_defines_architecture_when_workloads_do_not() {
927 let stack = stack_with_container();
928 let settings = ComputeSettings {
929 pools: [(
930 "general".to_string(),
931 ComputePoolSelection::Fixed {
932 machines: 1,
933 machine: Some("m7i.xlarge".to_string()),
934 failure_domains: None,
935 },
936 )]
937 .into_iter()
938 .collect(),
939 };
940
941 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
942
943 let pool = plan.pools.first().expect("general pool should exist");
944 assert!(pool.errors.is_empty());
945 assert!(pool.machines.iter().all(|machine| {
946 instance_catalog::find_instance_type(Platform::Aws, &machine.machine)
947 .is_some_and(|spec| spec.architecture == Architecture::X86_64)
948 }));
949 }
950
951 #[test]
952 fn explicit_capacity_group_requirements_are_merged_with_workloads() {
953 let mut stack = stack_with_container();
954 let cluster = ComputeCluster::new("compute".to_string())
955 .capacity_group(CapacityGroup {
956 group_id: "general".to_string(),
957 instance_type: None,
958 profile: Some(MachineProfile {
959 cpu: "4".to_string(),
960 memory_bytes: 16 * 1024 * 1024 * 1024,
961 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
962 architecture: Some(Architecture::X86_64),
963 gpu: None,
964 }),
965 min_size: 2,
966 max_size: 5,
967 scale_policy: None,
968 nested_virtualization: Some(true),
969 })
970 .build();
971 stack.resources.insert(
972 "compute".to_string(),
973 ResourceEntry {
974 config: Resource::new(cluster),
975 lifecycle: ResourceLifecycle::Frozen,
976 dependencies: Vec::new(),
977 remote_access: false,
978 enabled_when: None,
979 },
980 );
981
982 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
983
984 let pool = plan.pools.first().expect("general pool should exist");
985 let machine = pool
986 .selected
987 .machine()
988 .expect("AWS selection should include a machine");
989 let spec = instance_catalog::find_instance_type(Platform::Aws, machine)
990 .expect("selected machine should exist in the catalog");
991 assert!(spec.is_nested_virt_capable());
992 assert_eq!(pool.selected.min_size(), 2);
993 assert_eq!(pool.selected.max_size(), 5);
994 assert!(pool.errors.is_empty());
995 }
996
997 #[test]
998 fn materialized_machine_profile_is_not_counted_as_workload_demand() {
999 let mut stack = stack_with_container();
1000 let selected = instance_catalog::find_instance_type(Platform::Aws, "m7g.xlarge")
1001 .expect("machine should exist in catalog");
1002 let cluster = ComputeCluster::new("compute".to_string())
1003 .capacity_group(CapacityGroup {
1004 group_id: "general".to_string(),
1005 instance_type: Some(selected.name.to_string()),
1006 profile: Some(selected.to_machine_profile()),
1007 min_size: 1,
1008 max_size: 1,
1009 scale_policy: None,
1010 nested_virtualization: None,
1011 })
1012 .build();
1013 stack.resources.insert(
1014 "compute".to_string(),
1015 ResourceEntry {
1016 config: Resource::new(cluster),
1017 lifecycle: ResourceLifecycle::Frozen,
1018 dependencies: Vec::new(),
1019 remote_access: false,
1020 enabled_when: None,
1021 },
1022 );
1023
1024 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1025 let pool = plan.pools.first().expect("general pool should exist");
1026
1027 assert_eq!(pool.requirements.cpu, "2");
1028 assert_eq!(pool.requirements.memory_bytes, 4 * 1024 * 1024 * 1024);
1029 assert_eq!(pool.recommended.machine(), Some("m7g.xlarge"));
1030 assert!(pool
1031 .machines
1032 .iter()
1033 .any(|option| option.machine == "m7g.xlarge"));
1034 }
1035
1036 #[test]
1037 fn nested_x86_fixed_range_pool_preserves_bounds_and_rejects_graviton() {
1038 let daemon = Daemon::new("vm-runtime-loader".to_string())
1039 .code(DaemonCode::Image {
1040 image: "example.com/vm-runtime:latest".to_string(),
1041 })
1042 .cluster("vm-runtime".to_string())
1043 .cpu(ResourceSpec {
1044 min: "2".to_string(),
1045 desired: "2".to_string(),
1046 })
1047 .memory(ResourceSpec {
1048 min: "4Gi".to_string(),
1049 desired: "4Gi".to_string(),
1050 })
1051 .permissions("loader".to_string())
1052 .build();
1053 let cluster = ComputeCluster::new("vm-runtime".to_string())
1054 .capacity_group(CapacityGroup {
1055 group_id: "general".to_string(),
1056 instance_type: None,
1057 profile: Some(MachineProfile {
1058 cpu: "4".to_string(),
1059 memory_bytes: 16 * 1024 * 1024 * 1024,
1060 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
1061 architecture: Some(Architecture::X86_64),
1062 gpu: None,
1063 }),
1064 min_size: 2,
1065 max_size: 2,
1066 scale_policy: Some(CapacityGroupScalePolicy::Fixed {
1067 machines: ComputeChoiceRange {
1068 min: 1,
1069 max: 5,
1070 default: 2,
1071 },
1072 }),
1073 nested_virtualization: Some(true),
1074 })
1075 .build();
1076 let stack = Stack {
1077 id: "vm-runtime".to_string(),
1078 resources: [
1079 (
1080 "vm-runtime-loader".to_string(),
1081 ResourceEntry {
1082 config: Resource::new(daemon),
1083 lifecycle: ResourceLifecycle::Live,
1084 dependencies: Vec::new(),
1085 remote_access: false,
1086 enabled_when: None,
1087 },
1088 ),
1089 (
1090 "vm-runtime".to_string(),
1091 ResourceEntry {
1092 config: Resource::new(cluster),
1093 lifecycle: ResourceLifecycle::Frozen,
1094 dependencies: Vec::new(),
1095 remote_access: false,
1096 enabled_when: None,
1097 },
1098 ),
1099 ]
1100 .into_iter()
1101 .collect(),
1102 permissions: crate::permissions::PermissionsConfig::default(),
1103 supported_platforms: None,
1104 inputs: vec![],
1105 dynamic_container_repositories: Vec::new(),
1106 dynamic_container_image_resources: Vec::new(),
1107 };
1108
1109 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1110 let pool = plan.pools.first().expect("general pool should exist");
1111 assert_eq!(pool.recommended.machine(), Some("m8i.2xlarge"));
1112 assert_eq!(pool.recommended.min_size(), 2);
1113 assert_eq!(pool.recommended.max_size(), 2);
1114 assert_eq!(
1115 pool.scale,
1116 CapacityGroupScalePolicy::Fixed {
1117 machines: ComputeChoiceRange {
1118 min: 1,
1119 max: 5,
1120 default: 2,
1121 },
1122 }
1123 );
1124 assert!(!pool
1125 .machines
1126 .iter()
1127 .any(|option| option.machine == "m7g.2xlarge"));
1128
1129 let invalid_settings = ComputeSettings {
1130 pools: [(
1131 "general".to_string(),
1132 ComputePoolSelection::Fixed {
1133 machines: 2,
1134 machine: Some("m7g.2xlarge".to_string()),
1135 failure_domains: None,
1136 },
1137 )]
1138 .into_iter()
1139 .collect(),
1140 };
1141 let invalid_plan = plan_compute(&stack, Platform::Aws, Some(&invalid_settings))
1142 .expect("plan should build");
1143 assert!(!invalid_plan.pools[0].errors.is_empty());
1144 }
1145
1146 #[test]
1147 fn local_plan_has_no_provider_machine_choices() {
1148 let stack = stack_with_container();
1149
1150 let plan = plan_compute(&stack, Platform::Local, None).expect("plan should build");
1151
1152 let pool = plan.pools.first().expect("general pool should exist");
1153 assert_eq!(pool.selected.machine(), None);
1154 assert!(pool.machines.is_empty());
1155 assert!(pool.errors.is_empty());
1156 }
1157}